From 73d6af785f08b17aa102f6e0923403c845385e67 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 29 Jan 2025 01:45:47 +0530 Subject: [PATCH] Update 0.23.0 - OSS release --- .dockerignore | 5 +- .pre-commit-config.yaml | 145 ++ .vscode/extensions.json | 28 + .vscode/settings.json | 43 + CHANGELOG-Windows.rst | 18 + CHANGELOG.rst | 208 ++ LICENSE | 26 +- README.md | 55 +- chained_optimizations/README.md | 18 - .../scripts/3_onnx_export.sh | 10 - .../cache_diffusion/cache_diffusion/module.py | 55 - docker/Dockerfile | 10 +- docker/build.sh | 21 +- examples/README.md | 18 + benchmark.md => examples/benchmark.md | 0 .../chained_optimizations}/.gitignore | 0 examples/chained_optimizations/README.md | 56 + .../bert_prune_distill_quantize.py | 41 +- .../chained_optimizations}/scripts/1_prune.sh | 15 + .../scripts/2_int8_quantize.sh | 15 + .../scripts/3_onnx_export.sh | 25 + .../diffusers}/cache_diffusion/README.md | 0 .../assets/SDXL_Cache_Diffusion_Img.png | Bin .../cache_diffusion/assets/sdxl_cache.png | Bin .../diffusers}/cache_diffusion/benchmarks.py | 56 +- .../cache_diffusion/cachify.py | 26 +- .../cache_diffusion/cache_diffusion/module.py | 49 + .../cache_diffusion/cache_diffusion/utils.py | 26 +- .../diffusers}/cache_diffusion/example.ipynb | 0 .../cache_diffusion/pipeline/config.py | 27 +- .../cache_diffusion/pipeline/deploy.py | 26 +- .../cache_diffusion/pipeline/models/sd3.py | 34 +- .../cache_diffusion/pipeline/models/sdxl.py | 42 +- .../cache_diffusion/pipeline/utils.py | 26 +- .../cache_diffusion/requirements.txt | 3 +- .../diffusers}/quantization/.gitignore | 0 .../diffusers/quantization/QAT_GUIDANCE.md | 108 + .../diffusers}/quantization/README.md | 55 +- .../assets/xl_base-fp16-sakura.png | Bin .../quantization/assets/xl_base-fp16.png | Bin .../assets/xl_base-fp8-sakura.png | Bin .../quantization/assets/xl_base-int8.png | Bin .../quantization/build_sdxl_8bit_engine.sh | 15 + .../quantization/calib/calib_prompts.txt | 0 .../quantization/calib/plugin_calib.py | 26 +- .../diffusers}/quantization/config.py | 270 +- .../diffusers}/quantization/diffusion_trt.py | 58 +- .../quantization/onnx_utils/export.py | 122 +- .../onnx_utils/fp8_onnx_graphsurgeon.py | 26 +- .../diffusers}/quantization/quantize.py | 79 +- .../diffusers}/quantization/requirements.txt | 5 +- .../diffusers}/quantization/utils.py | 28 +- .../llm_distill}/README.md | 0 {llm_distill => examples/llm_distill}/main.py | 46 +- examples/llm_distill/requirements.txt | 2 + {llm_eval => examples/llm_eval}/.gitignore | 0 {llm_eval => examples/llm_eval}/README.md | 18 +- examples/llm_eval/__init__.py | 14 + .../llm_eval}/gen_model_answer.py | 37 +- {llm_eval => examples/llm_eval}/humaneval.py | 31 +- {llm_eval => examples/llm_eval}/lm_eval_hf.py | 32 +- .../llm_eval}/lm_eval_tensorrt_llm.py | 62 +- {llm_eval => examples/llm_eval}/mmlu.py | 30 +- {llm_eval => examples/llm_eval}/modeling.py | 34 +- .../llm_eval}/quantization_utils.py | 83 +- .../llm_eval}/requirements.txt | 3 +- .../llm_eval}/run_fastchat.sh | 15 + {llm_ptq => examples/llm_ptq}/.gitignore | 0 {llm_ptq => examples/llm_ptq}/README.md | 99 +- .../llm_ptq/benchmarks/benchmark_suite.py | 21 + .../disable_bench_chunked_context.py | 36 + .../config/megatron_quantization.yaml | 0 .../llm_ptq}/example_utils.py | 82 +- {llm_ptq => examples/llm_ptq}/hf_ptq.py | 186 +- .../llm_ptq}/modelopt_to_tensorrt_llm.py | 224 +- {llm_ptq => examples/llm_ptq}/nemo_ptq.py | 26 +- .../llm_ptq}/requirements.txt | 3 +- .../llm_ptq}/run_tensorrt_llm.py | 36 +- .../llm_ptq}/scripts/huggingface_example.sh | 78 +- .../llm_ptq}/scripts/nemo_example.sh | 19 +- .../llm_ptq}/scripts/parser.sh | 22 +- {llm_qat => examples/llm_qat}/README.md | 0 {llm_qat => examples/llm_qat}/launch.sh | 15 + {llm_qat => examples/llm_qat}/main.py | 26 +- .../llm_qat}/requirements.txt | 2 - {llm_qat => examples/llm_qat}/utils.py | 29 +- .../llm_sparsity}/README.md | 0 .../llm_sparsity}/data_prep.py | 26 +- .../llm_sparsity}/eval.py | 34 +- .../llm_sparsity}/export_trtllm_ckpt.py | 48 +- .../llm_sparsity}/finetune.py | 47 +- .../llm_sparsity}/hf_pts.py | 50 +- .../llm_sparsity}/launch_finetune.sh | 15 + .../llm_sparsity}/requirements.txt | 1 - examples/llm_sparsity/utils.py | 31 + {model_hub => examples/model_hub}/README.md | 0 examples/model_hub/download_hf_ckpt.py | 31 + .../model_hub}/llama_fp8_deploy_trtllm.sh | 15 + .../model_hub}/llama_fp8_deploy_vllm.sh | 15 + {onnx_ptq => examples/onnx_ptq}/README.md | 10 +- .../onnx_ptq}/download_example_onnx.py | 26 +- {onnx_ptq => examples/onnx_ptq}/evaluate.py | 26 +- {onnx_ptq => examples/onnx_ptq}/evaluation.py | 30 +- {onnx_ptq => examples/onnx_ptq}/image_prep.py | 26 +- .../onnx_ptq}/requirements.txt | 1 - {pruning => examples/pruning}/README.md | 0 .../pruning}/cifar_resnet.ipynb | 0 .../speculative_decoding}/README.md | 14 +- .../speculative_decoding}/eagle_utils.py | 36 +- .../speculative_decoding}/example.ipynb | 65 +- .../speculative_decoding}/launch.sh | 15 + .../speculative_decoding}/main.py | 34 +- .../speculative_decoding}/medusa_utils.py | 35 +- examples/speculative_decoding/prepare_data.sh | 48 + .../speculative_decoding}/requirements.txt | 1 - .../speculative_decoding}/vllm_generate.py | 28 +- examples/vlm_eval/README.md | 36 + examples/vlm_eval/convert_gqa_for_eval.py | 56 + examples/vlm_eval/gqa.sh | 142 ++ examples/vlm_eval/model_gqa_loader.py | 263 ++ .../vlm_eval}/requirements-vila.txt | 2 +- examples/vlm_eval/vlm_eval_utils.py | 64 + {vlm_ptq => examples/vlm_ptq}/.gitignore | 0 {vlm_ptq => examples/vlm_ptq}/README.md | 4 +- .../vlm_ptq}/requirements-mllama.txt | 0 examples/vlm_ptq/requirements-vila.txt | 2 + .../vlm_ptq}/scripts/huggingface_example.sh | 95 +- examples/vlm_ptq/utils.py | 137 + examples/vlm_ptq/vlm_run.py | 93 + examples/vlm_ptq/vlm_visual_engine.py | 26 + {windows => examples/windows}/Benchmark.md | 0 {windows => examples/windows}/README.md | 5 +- .../windows}/accuracy_benchmark/README.md | 0 .../accuracy_benchmark/mmlu_benchmark.py | 30 +- .../windows}/accuracy_benchmark/modeling.py | 34 +- .../accuracy_benchmark/quantization_utils.py | 33 +- .../accuracy_benchmark/requirements.txt | 1 + .../accuracy_benchmark/trtllm_utils.py | 29 +- .../windows}/onnx_ptq/README.md | 0 .../windows}/onnx_ptq/quantize.py | 42 +- .../windows}/onnx_ptq/requirements.txt | 2 +- llm_distill/requirements.txt | 3 - llm_eval/__init__.py | 20 - llm_ptq/benchmarks/benchmark_suite.py | 27 - .../disable_bench_chunked_context.py | 42 - llm_sparsity/utils.py | 37 - model_hub/download_hf_ckpt.py | 37 - modelopt/__init__.py | 42 + modelopt/deploy/__init__.py | 16 + modelopt/deploy/llm/__init__.py | 33 + modelopt/deploy/llm/generate.py | 255 ++ modelopt/deploy/llm/nemo_utils.py | 180 ++ modelopt/onnx/__init__.py | 35 + modelopt/onnx/op_types.py | 287 +++ modelopt/onnx/quantization/__init__.py | 20 + modelopt/onnx/quantization/__main__.py | 233 ++ modelopt/onnx/quantization/calib_utils.py | 151 ++ modelopt/onnx/quantization/extensions.py | 34 + modelopt/onnx/quantization/fp8.py | 307 +++ modelopt/onnx/quantization/graph_utils.py | 1098 ++++++++ modelopt/onnx/quantization/gs_patching.py | 112 + modelopt/onnx/quantization/int4.py | 1326 ++++++++++ modelopt/onnx/quantization/int8.py | 260 ++ modelopt/onnx/quantization/operators.py | 96 + modelopt/onnx/quantization/ort_patching.py | 610 +++++ modelopt/onnx/quantization/ort_utils.py | 137 + modelopt/onnx/quantization/partitioning.py | 411 +++ modelopt/onnx/quantization/qdq_utils.py | 822 ++++++ modelopt/onnx/quantization/quant_utils.py | 174 ++ modelopt/onnx/quantization/quantize.py | 369 +++ .../src/modelopt_round_and_pack_ext.cpp | 127 + modelopt/onnx/quantization/trt_utils.py | 116 + modelopt/onnx/utils.py | 629 +++++ modelopt/torch/__init__.py | 31 + modelopt/torch/_deploy/__init__.py | 23 + modelopt/torch/_deploy/_runtime/__init__.py | 30 + modelopt/torch/_deploy/_runtime/common.py | 73 + modelopt/torch/_deploy/_runtime/ort_client.py | 190 ++ modelopt/torch/_deploy/_runtime/registry.py | 83 + .../torch/_deploy/_runtime/runtime_client.py | 158 ++ .../_deploy/_runtime/tensorrt/constants.py | 107 + .../_runtime/tensorrt/engine_builder.py | 332 +++ .../_runtime/tensorrt/hw_param_config.py | 51 + .../_runtime/tensorrt/layerwise_profiling.py | 175 ++ .../_runtime/tensorrt/parse_trtexec_log.py | 154 ++ .../_runtime/tensorrt/tensorrt_utils.py | 198 ++ modelopt/torch/_deploy/_runtime/trt_client.py | 217 ++ modelopt/torch/_deploy/compilation.py | 143 ++ modelopt/torch/_deploy/device_model.py | 157 ++ modelopt/torch/_deploy/profiling.py | 195 ++ modelopt/torch/_deploy/utils/__init__.py | 18 + .../torch/_deploy/utils/onnx_optimizer.py | 118 + modelopt/torch/_deploy/utils/onnx_utils.py | 58 + modelopt/torch/_deploy/utils/torch_onnx.py | 561 ++++ modelopt/torch/distill/__init__.py | 28 + modelopt/torch/distill/config.py | 130 + modelopt/torch/distill/distillation.py | 84 + modelopt/torch/distill/distillation_model.py | 317 +++ modelopt/torch/distill/loss_balancers.py | 139 + modelopt/torch/distill/losses.py | 150 ++ modelopt/torch/distill/mode.py | 224 ++ modelopt/torch/distill/plugins/__init__.py | 21 + modelopt/torch/distill/plugins/megatron.py | 477 ++++ modelopt/torch/distill/registry.py | 23 + modelopt/torch/export/__init__.py | 21 + modelopt/torch/export/distribute.py | 303 +++ modelopt/torch/export/hf_config_map.py | 81 + modelopt/torch/export/layer_utils.py | 1570 ++++++++++++ modelopt/torch/export/mcore_config_map.py | 24 + modelopt/torch/export/mcore_hf_export_map.py | 163 ++ modelopt/torch/export/model_config.py | 597 +++++ modelopt/torch/export/model_config_export.py | 549 ++++ modelopt/torch/export/model_config_utils.py | 401 +++ modelopt/torch/export/postprocess.py | 817 ++++++ modelopt/torch/export/quant_utils.py | 687 +++++ modelopt/torch/export/tensorrt_llm_type.py | 42 + modelopt/torch/export/tensorrt_llm_utils.py | 566 +++++ modelopt/torch/export/transformer_engine.py | 94 + modelopt/torch/export/unified_export_hf.py | 363 +++ .../torch/export/unified_export_megatron.py | 692 +++++ modelopt/torch/nas/__init__.py | 27 + modelopt/torch/nas/_algorithms.py | 430 ++++ modelopt/torch/nas/_patch.py | 262 ++ modelopt/torch/nas/algorithms.py | 452 ++++ modelopt/torch/nas/autonas.py | 637 +++++ modelopt/torch/nas/config.py | 94 + modelopt/torch/nas/conversion.py | 119 + modelopt/torch/nas/hparams/__init__.py | 19 + modelopt/torch/nas/hparams/concat.py | 281 ++ modelopt/torch/nas/hparams/container.py | 48 + modelopt/torch/nas/mode.py | 135 + modelopt/torch/nas/modules/__init__.py | 21 + modelopt/torch/nas/modules/container.py | 98 + modelopt/torch/nas/modules/conv.py | 253 ++ modelopt/torch/nas/modules/linear.py | 79 + modelopt/torch/nas/modules/norm.py | 198 ++ modelopt/torch/nas/modules/utils.py | 63 + modelopt/torch/nas/plugins/__init__.py | 30 + modelopt/torch/nas/plugins/torch.py | 71 + .../torch/nas/plugins/transformer_engine.py | 42 + modelopt/torch/nas/plugins/transformers.py | 207 ++ modelopt/torch/nas/registry.py | 23 + modelopt/torch/nas/search_space.py | 245 ++ modelopt/torch/nas/traced_hp.py | 149 ++ modelopt/torch/nas/utils.py | 406 +++ modelopt/torch/opt/__init__.py | 43 + modelopt/torch/opt/_hooks.py | 66 + modelopt/torch/opt/config.py | 396 +++ modelopt/torch/opt/conversion.py | 588 +++++ modelopt/torch/opt/dynamic.py | 1295 ++++++++++ modelopt/torch/opt/hparam.py | 261 ++ modelopt/torch/opt/mode.py | 313 +++ modelopt/torch/opt/plugins/__init__.py | 32 + modelopt/torch/opt/plugins/diffusers.py | 24 + modelopt/torch/opt/plugins/huggingface.py | 192 ++ .../opt/plugins/mcore_dist_checkpointing.py | 412 +++ modelopt/torch/opt/plugins/peft.py | 100 + modelopt/torch/opt/plugins/transformers.py | 24 + modelopt/torch/opt/searcher.py | 364 +++ modelopt/torch/opt/utils.py | 77 + modelopt/torch/prune/__init__.py | 26 + modelopt/torch/prune/config.py | 86 + modelopt/torch/prune/fastnas.py | 263 ++ modelopt/torch/prune/gradnas.py | 314 +++ modelopt/torch/prune/mcore_gpt_minitron.py | 155 ++ modelopt/torch/prune/mode.py | 146 ++ modelopt/torch/prune/plugins/__init__.py | 24 + modelopt/torch/prune/plugins/megatron.py | 32 + modelopt/torch/prune/plugins/transformers.py | 40 + modelopt/torch/prune/pruning.py | 211 ++ modelopt/torch/quantization/__init__.py | 25 + modelopt/torch/quantization/calib/__init__.py | 24 + .../torch/quantization/calib/calibrator.py | 63 + .../torch/quantization/calib/histogram.py | 436 ++++ modelopt/torch/quantization/calib/max.py | 103 + modelopt/torch/quantization/config.py | 812 ++++++ modelopt/torch/quantization/conversion.py | 337 +++ modelopt/torch/quantization/export_onnx.py | 544 ++++ modelopt/torch/quantization/extensions.py | 89 + modelopt/torch/quantization/mode.py | 122 + modelopt/torch/quantization/model_calib.py | 912 +++++++ modelopt/torch/quantization/model_quant.py | 448 ++++ modelopt/torch/quantization/nn/__init__.py | 27 + modelopt/torch/quantization/nn/functional.py | 72 + .../torch/quantization/nn/modules/__init__.py | 16 + .../torch/quantization/nn/modules/clip.py | 56 + .../nn/modules/quant_activations.py | 22 + .../nn/modules/quant_batchnorm.py | 24 + .../quantization/nn/modules/quant_conv.py | 123 + .../nn/modules/quant_instancenorm.py | 45 + .../quantization/nn/modules/quant_linear.py | 50 + .../quantization/nn/modules/quant_module.py | 186 ++ .../quantization/nn/modules/quant_pooling.py | 117 + .../quantization/nn/modules/quant_rnn.py | 532 ++++ .../nn/modules/tensor_quantizer.py | 984 +++++++ modelopt/torch/quantization/optim.py | 38 + .../torch/quantization/plugins/__init__.py | 51 + modelopt/torch/quantization/plugins/apex.py | 54 + .../torch/quantization/plugins/attention.py | 229 ++ modelopt/torch/quantization/plugins/custom.py | 110 + .../torch/quantization/plugins/diffusers.py | 231 ++ .../torch/quantization/plugins/fairscale.py | 44 + .../torch/quantization/plugins/huggingface.py | 251 ++ .../torch/quantization/plugins/megatron.py | 298 +++ modelopt/torch/quantization/plugins/nemo.py | 67 + modelopt/torch/quantization/plugins/peft.py | 56 + .../torch/quantization/qtensor/__init__.py | 22 + .../quantization/qtensor/base_qtensor.py | 162 ++ .../torch/quantization/qtensor/int4_tensor.py | 116 + .../torch/quantization/qtensor/nf4_tensor.py | 185 ++ modelopt/torch/quantization/quant_modules.py | 50 + .../torch/quantization/src/tensor_quant.cpp | 77 + .../torch/quantization/src/tensor_quant.h | 69 + .../quantization/src/tensor_quant_fp8.cpp | 48 + .../quantization/src/tensor_quant_gpu.cu | 361 +++ .../quantization/src/tensor_quant_gpu_fp8.cu | 122 + .../torch/quantization/src/tensor_quant_mx.cu | 401 +++ .../torch/quantization/src/tensor_quant_mx.h | 247 ++ modelopt/torch/quantization/tensor_quant.py | 653 +++++ modelopt/torch/quantization/utils.py | 184 ++ modelopt/torch/sparsity/__init__.py | 19 + modelopt/torch/sparsity/config.py | 51 + modelopt/torch/sparsity/magnitude.py | 148 ++ modelopt/torch/sparsity/mode.py | 204 ++ modelopt/torch/sparsity/module.py | 89 + modelopt/torch/sparsity/plugins/__init__.py | 27 + modelopt/torch/sparsity/plugins/megatron.py | 78 + modelopt/torch/sparsity/searcher.py | 84 + modelopt/torch/sparsity/sparsegpt.py | 276 ++ modelopt/torch/sparsity/sparsification.py | 123 + modelopt/torch/speculative/__init__.py | 20 + modelopt/torch/speculative/config.py | 63 + modelopt/torch/speculative/eagle/__init__.py | 19 + .../torch/speculative/eagle/conversion.py | 59 + .../torch/speculative/eagle/eagle_model.py | 32 + modelopt/torch/speculative/eagle/utils.py | 388 +++ modelopt/torch/speculative/medusa/__init__.py | 19 + .../torch/speculative/medusa/conversion.py | 59 + .../torch/speculative/medusa/medusa_model.py | 32 + modelopt/torch/speculative/mode.py | 115 + .../torch/speculative/plugins/__init__.py | 30 + .../torch/speculative/plugins/megatron.py | 722 ++++++ .../torch/speculative/plugins/transformers.py | 620 +++++ .../torch/speculative/redrafter/__init__.py | 19 + .../torch/speculative/redrafter/conversion.py | 62 + .../speculative/redrafter/redrafter_model.py | 32 + .../torch/speculative/speculative_decoding.py | 55 + modelopt/torch/speculative/utils.py | 53 + modelopt/torch/trace/__init__.py | 25 + modelopt/torch/trace/analyzer.py | 1380 ++++++++++ modelopt/torch/trace/modules/__init__.py | 19 + modelopt/torch/trace/modules/concat.py | 387 +++ modelopt/torch/trace/modules/nn.py | 157 ++ modelopt/torch/trace/plugins/__init__.py | 34 + modelopt/torch/trace/plugins/megatron.py | 29 + modelopt/torch/trace/plugins/transformers.py | 58 + modelopt/torch/trace/symbols.py | 544 ++++ modelopt/torch/trace/tracer.py | 330 +++ modelopt/torch/utils/__init__.py | 28 + modelopt/torch/utils/_pytree.py | 133 + modelopt/torch/utils/cpp_extension.py | 87 + .../torch/utils}/dataset_utils.py | 112 +- modelopt/torch/utils/distributed.py | 323 +++ modelopt/torch/utils/graph.py | 125 + modelopt/torch/utils/image_processor.py | 115 + modelopt/torch/utils/import_utils.py | 34 + modelopt/torch/utils/list.py | 59 + modelopt/torch/utils/logging.py | 127 + modelopt/torch/utils/network.py | 609 +++++ modelopt/torch/utils/perf.py | 90 + modelopt/torch/utils/random.py | 160 ++ modelopt/torch/utils/tensor.py | 55 + modelopt/torch/utils/vlm_dataset_utils.py | 111 + pyproject.toml | 149 ++ setup.py | 151 ++ speculative_decoding/prepare_data.sh | 33 - tests/_test_utils/import_helper.py | 57 + .../onnx_quantization/lib_test_models.py | 178 ++ .../_test_utils/torch_deploy/device_model.py | 88 + tests/_test_utils/torch_deploy/runtime.py | 71 + tests/_test_utils/torch_dist/dist_utils.py | 80 + tests/_test_utils/torch_dist/fsdp_test.py | 115 + .../torch_dist/plugins/megatron_common.py | 290 +++ tests/_test_utils/torch_misc.py | 36 + .../torch_model/benchmark_models.py | 171 ++ .../_test_utils/torch_model/deploy_models.py | 814 ++++++ .../_test_utils/torch_model/example_models.py | 2264 +++++++++++++++++ tests/_test_utils/torch_model/utils.py | 27 + .../_test_utils/torch_quantization/models.py | 122 + .../torch_quantization/onnx_export.py | 82 + .../torch_quantization/quant_utils.py | 34 + .../torch_quantization/quantize_common.py | 140 + .../torch_quantization/tensor_quant_common.py | 216 ++ .../tensor_quantizer_common.py | 508 ++++ tests/conftest.py | 38 + tests/deploy/unit/test_deploy_dummy.py | 18 + tests/onnx/gpu/test_onnx_torch_int4_awq.py | 115 + tests/onnx/gpu/test_plugin.py | 121 + tests/onnx/unit/test_onnx_utils.py | 33 + tests/onnx/unit/test_partitioning.py | 87 + tests/onnx/unit/test_qdq_rules_int8.py | 125 + tests/onnx/unit/test_qdq_utils_fp8.py | 51 + tests/onnx/unit/test_quant_utils.py | 187 ++ tests/onnx/unit/test_quantize_fp8.py | 56 + tests/onnx/unit/test_quantize_int4.py | 146 ++ tests/onnx/unit/test_quantize_int8.py | 56 + tests/torch/conftest.py | 25 + tests/torch/gpu/conftest.py | 32 + .../gpu/deploy/_runtime/test_trt_client.py | 81 + .../gpu/deploy/test_gpu_deploy_benchmark.py | 25 + .../export/test_unified_export_megatron.py | 181 ++ .../plugins/test_megatron_dynamic_modules.py | 295 +++ tests/torch/gpu/nas/test_distributed_model.py | 84 + ...test_search_space_with_benchmark_models.py | 90 + .../gpu/opt/plugins/test_megatron_chaining.py | 49 + tests/torch/gpu/opt/test_fsdp_save_restore.py | 81 + .../test_mcore_gpt_minitron_pruning.py | 163 ++ tests/torch/gpu/quantization/conftest.py | 21 + .../gpu/quantization/plugins/test_apex.py | 139 + .../gpu/quantization/plugins/test_megatron.py | 245 ++ tests/torch/gpu/quantization/test_fsdp.py | 92 + .../quantization/test_nvfp4_onnx_export.py | 97 + .../gpu/quantization/test_onnx_export_cuda.py | 46 + .../gpu/quantization/test_qtensor_cuda.py | 288 +++ .../gpu/quantization/test_quant_rnn_cuda.py | 53 + .../gpu/quantization/test_quantize_cuda.py | 75 + .../test_quantize_mxformats_cuda.py | 317 +++ .../quantization/test_real_quantize_cuda.py | 94 + .../quantization/test_tensor_quant_cuda.py | 274 ++ .../test_tensor_quantizer_cuda.py | 89 + .../gpu/quantization/test_torch_export.py | 106 + .../plugins/test_megatron_sparsity.py | 49 + tests/torch/gpu/sparsity/test_sparse_fsdp.py | 55 + .../test_speculative_megatron_modules.py | 144 ++ .../integration/nas/test_full_algorithms.py | 262 ++ .../quantization/test_onnx_export_cpu.py | 40 + tests/torch/unit/conftest.py | 54 + .../_runtime/tensorrt/test_engine_builder.py | 143 ++ .../tensorrt/test_layerwise_profiling.py | 189 ++ .../_runtime/tensorrt/test_tensorrt_utils.py | 66 + tests/torch/unit/deploy/conftest.py | 30 + .../unit/deploy/test_cpu_deploy_benchmark.py | 25 + .../torch/unit/deploy/test_runtime_config.py | 47 + .../deploy/utils/test_torch_onnx_utils.py | 261 ++ tests/torch/unit/distill/test_all.py | 287 +++ .../unit/export/test_get_quantization.py | 66 + .../unit/export/test_torch_export_config.py | 771 ++++++ .../unit/export/test_torch_export_dummy.py | 18 + tests/torch/unit/nas/modules/test_conv.py | 238 ++ tests/torch/unit/nas/modules/test_linear.py | 109 + .../torch/unit/nas/modules/test_mod_utils.py | 35 + tests/torch/unit/nas/modules/test_norm.py | 267 ++ .../torch/unit/nas/modules/test_sequential.py | 107 + .../nas/plugins/test_dynamic_hf_attention.py | 101 + .../nas/plugins/test_hf_nas_save_restore.py | 73 + .../unit/nas/plugins/test_torch_hooks.py | 76 + tests/torch/unit/nas/test_chaining.py | 231 ++ .../unit/nas/test_evaluate_constraints.py | 106 + .../unit/nas/test_latency_interpolator.py | 123 + tests/torch/unit/nas/test_nas.py | 519 ++++ tests/torch/unit/nas/test_nas_utils.py | 40 + tests/torch/unit/nas/test_registry.py | 126 + tests/torch/unit/nas/test_search_space.py | 343 +++ .../test_search_space_with_example_models.py | 85 + .../nas/test_search_space_with_torchvision.py | 89 + .../unit/opt/plugins/_transformers_helper.py | 33 + .../unit/opt/plugins/test_peft_save_load.py | 60 + .../plugins/test_transformers_save_load.py | 82 + tests/torch/unit/opt/test_config.py | 118 + tests/torch/unit/opt/test_dynamic.py | 123 + tests/torch/unit/prune/test_algorithms.py | 43 + .../unit/prune/test_fastnas_conversion.py | 29 + tests/torch/unit/prune/test_gradnas.py | 116 + .../quantization/plugins/test_accelerate.py | 53 + .../plugins/test_attention_quant.py | 74 + .../quantization/plugins/test_huggingface.py | 229 ++ .../unit/quantization/plugins/test_peft.py | 73 + .../torch/unit/quantization/test_autoquant.py | 268 ++ tests/torch/unit/quantization/test_calib.py | 366 +++ .../unit/quantization/test_calibrator.py | 433 ++++ tests/torch/unit/quantization/test_clip.py | 52 + tests/torch/unit/quantization/test_dist.py | 47 + .../unit/quantization/test_module_registry.py | 53 + tests/torch/unit/quantization/test_print.py | 51 + .../quantization/test_quant_activations.py | 54 + .../unit/quantization/test_quant_batchnorm.py | 111 + .../unit/quantization/test_quant_conv.py | 318 +++ .../quantization/test_quant_instancenorm.py | 108 + .../unit/quantization/test_quant_linear.py | 220 ++ .../unit/quantization/test_quant_pooling.py | 188 ++ .../torch/unit/quantization/test_quant_rnn.py | 261 ++ .../unit/quantization/test_quantize_cpu.py | 217 ++ .../quantization/test_quantize_replace.py | 48 + .../quantization/test_tensor_quant_cpu.py | 124 + .../quantization/test_tensor_quantizer_cpu.py | 34 + tests/torch/unit/sparsity/test_sparsify.py | 113 + .../plugins/test_hf_speculative.py | 54 + .../test_transformers_attention_symbols.py | 41 + tests/torch/unit/trace/test_analyzer.py | 60 + tests/torch/unit/trace/test_concat.py | 83 + tests/torch/unit/trace/test_model.py | 89 + tests/torch/unit/trace/test_nn.py | 142 ++ tests/torch/unit/trace/test_symbol.py | 197 ++ tests/torch/unit/trace/test_tracer.py | 61 + tests/torch/unit/utils/test_network.py | 223 ++ tests/torch/unit/utils/test_pytree.py | 74 + tests/torch/unit/utils/test_tensor.py | 33 + vlm_ptq/requirements-llava.txt | 1 - vlm_ptq/vlm_run.py | 173 -- vlm_ptq/vlm_visual_engine.py | 32 - 510 files changed, 71646 insertions(+), 2104 deletions(-) create mode 100644 .pre-commit-config.yaml create mode 100644 .vscode/extensions.json create mode 100644 .vscode/settings.json create mode 100644 CHANGELOG-Windows.rst create mode 100755 CHANGELOG.rst delete mode 100644 chained_optimizations/README.md delete mode 100755 chained_optimizations/scripts/3_onnx_export.sh delete mode 100644 diffusers/cache_diffusion/cache_diffusion/module.py create mode 100644 examples/README.md rename benchmark.md => examples/benchmark.md (100%) rename {chained_optimizations => examples/chained_optimizations}/.gitignore (100%) create mode 100644 examples/chained_optimizations/README.md rename {chained_optimizations => examples/chained_optimizations}/bert_prune_distill_quantize.py (97%) rename {chained_optimizations => examples/chained_optimizations}/scripts/1_prune.sh (60%) rename {chained_optimizations => examples/chained_optimizations}/scripts/2_int8_quantize.sh (58%) create mode 100755 examples/chained_optimizations/scripts/3_onnx_export.sh rename {diffusers => examples/diffusers}/cache_diffusion/README.md (100%) rename {diffusers => examples/diffusers}/cache_diffusion/assets/SDXL_Cache_Diffusion_Img.png (100%) rename {diffusers => examples/diffusers}/cache_diffusion/assets/sdxl_cache.png (100%) rename {diffusers => examples/diffusers}/cache_diffusion/benchmarks.py (62%) rename {diffusers => examples/diffusers}/cache_diffusion/cache_diffusion/cachify.py (79%) create mode 100644 examples/diffusers/cache_diffusion/cache_diffusion/module.py rename {diffusers => examples/diffusers}/cache_diffusion/cache_diffusion/utils.py (51%) rename {diffusers => examples/diffusers}/cache_diffusion/example.ipynb (100%) rename {diffusers => examples/diffusers}/cache_diffusion/pipeline/config.py (83%) rename {diffusers => examples/diffusers}/cache_diffusion/pipeline/deploy.py (87%) rename {diffusers => examples/diffusers}/cache_diffusion/pipeline/models/sd3.py (81%) rename {diffusers => examples/diffusers}/cache_diffusion/pipeline/models/sdxl.py (90%) rename {diffusers => examples/diffusers}/cache_diffusion/pipeline/utils.py (79%) rename {diffusers => examples/diffusers}/cache_diffusion/requirements.txt (58%) rename {diffusers => examples/diffusers}/quantization/.gitignore (100%) create mode 100644 examples/diffusers/quantization/QAT_GUIDANCE.md rename {diffusers => examples/diffusers}/quantization/README.md (83%) rename {diffusers => examples/diffusers}/quantization/assets/xl_base-fp16-sakura.png (100%) rename {diffusers => examples/diffusers}/quantization/assets/xl_base-fp16.png (100%) rename {diffusers => examples/diffusers}/quantization/assets/xl_base-fp8-sakura.png (100%) rename {diffusers => examples/diffusers}/quantization/assets/xl_base-int8.png (100%) rename {diffusers => examples/diffusers}/quantization/build_sdxl_8bit_engine.sh (66%) rename {diffusers => examples/diffusers}/quantization/calib/calib_prompts.txt (100%) rename {diffusers => examples/diffusers}/quantization/calib/plugin_calib.py (72%) rename {diffusers => examples/diffusers}/quantization/config.py (51%) rename {diffusers => examples/diffusers}/quantization/diffusion_trt.py (81%) rename {diffusers => examples/diffusers}/quantization/onnx_utils/export.py (71%) rename {diffusers => examples/diffusers}/quantization/onnx_utils/fp8_onnx_graphsurgeon.py (87%) rename {diffusers => examples/diffusers}/quantization/quantize.py (76%) rename {diffusers => examples/diffusers}/quantization/requirements.txt (67%) rename {diffusers => examples/diffusers}/quantization/utils.py (83%) rename {llm_distill => examples/llm_distill}/README.md (100%) rename {llm_distill => examples/llm_distill}/main.py (81%) create mode 100644 examples/llm_distill/requirements.txt rename {llm_eval => examples/llm_eval}/.gitignore (100%) rename {llm_eval => examples/llm_eval}/README.md (92%) create mode 100644 examples/llm_eval/__init__.py rename {llm_eval => examples/llm_eval}/gen_model_answer.py (93%) rename {llm_eval => examples/llm_eval}/humaneval.py (84%) rename {llm_eval => examples/llm_eval}/lm_eval_hf.py (78%) rename {llm_eval => examples/llm_eval}/lm_eval_tensorrt_llm.py (70%) rename {llm_eval => examples/llm_eval}/mmlu.py (90%) rename {llm_eval => examples/llm_eval}/modeling.py (93%) rename {llm_eval => examples/llm_eval}/quantization_utils.py (69%) rename {llm_eval => examples/llm_eval}/requirements.txt (61%) rename {llm_eval => examples/llm_eval}/run_fastchat.sh (84%) rename {llm_ptq => examples/llm_ptq}/.gitignore (100%) rename {llm_ptq => examples/llm_ptq}/README.md (79%) create mode 100644 examples/llm_ptq/benchmarks/benchmark_suite.py create mode 100644 examples/llm_ptq/benchmarks/disable_bench_chunked_context.py rename {llm_ptq => examples/llm_ptq}/config/megatron_quantization.yaml (100%) rename {llm_ptq => examples/llm_ptq}/example_utils.py (75%) rename {llm_ptq => examples/llm_ptq}/hf_ptq.py (69%) rename {llm_ptq => examples/llm_ptq}/modelopt_to_tensorrt_llm.py (52%) rename {llm_ptq => examples/llm_ptq}/nemo_ptq.py (92%) rename {llm_ptq => examples/llm_ptq}/requirements.txt (75%) rename {llm_ptq => examples/llm_ptq}/run_tensorrt_llm.py (63%) rename {llm_ptq => examples/llm_ptq}/scripts/huggingface_example.sh (80%) rename {llm_ptq => examples/llm_ptq}/scripts/nemo_example.sh (90%) rename {llm_ptq => examples/llm_ptq}/scripts/parser.sh (83%) rename {llm_qat => examples/llm_qat}/README.md (100%) rename {llm_qat => examples/llm_qat}/launch.sh (83%) rename {llm_qat => examples/llm_qat}/main.py (87%) rename {llm_qat => examples/llm_qat}/requirements.txt (77%) rename {llm_qat => examples/llm_qat}/utils.py (84%) rename {llm_sparsity => examples/llm_sparsity}/README.md (100%) rename {llm_sparsity => examples/llm_sparsity}/data_prep.py (60%) rename {llm_sparsity => examples/llm_sparsity}/eval.py (87%) rename {llm_sparsity => examples/llm_sparsity}/export_trtllm_ckpt.py (77%) rename {llm_sparsity => examples/llm_sparsity}/finetune.py (91%) rename {llm_sparsity => examples/llm_sparsity}/hf_pts.py (79%) rename {llm_sparsity => examples/llm_sparsity}/launch_finetune.sh (76%) rename {llm_sparsity => examples/llm_sparsity}/requirements.txt (77%) create mode 100644 examples/llm_sparsity/utils.py rename {model_hub => examples/model_hub}/README.md (100%) create mode 100644 examples/model_hub/download_hf_ckpt.py rename {model_hub => examples/model_hub}/llama_fp8_deploy_trtllm.sh (70%) rename {model_hub => examples/model_hub}/llama_fp8_deploy_vllm.sh (67%) rename {onnx_ptq => examples/onnx_ptq}/README.md (94%) rename {onnx_ptq => examples/onnx_ptq}/download_example_onnx.py (67%) rename {onnx_ptq => examples/onnx_ptq}/evaluate.py (79%) rename {onnx_ptq => examples/onnx_ptq}/evaluation.py (86%) rename {onnx_ptq => examples/onnx_ptq}/image_prep.py (56%) rename {onnx_ptq => examples/onnx_ptq}/requirements.txt (77%) rename {pruning => examples/pruning}/README.md (100%) rename {pruning => examples/pruning}/cifar_resnet.ipynb (100%) rename {speculative_decoding => examples/speculative_decoding}/README.md (93%) rename {speculative_decoding => examples/speculative_decoding}/eagle_utils.py (84%) rename {speculative_decoding => examples/speculative_decoding}/example.ipynb (82%) rename {speculative_decoding => examples/speculative_decoding}/launch.sh (84%) rename {speculative_decoding => examples/speculative_decoding}/main.py (84%) rename {speculative_decoding => examples/speculative_decoding}/medusa_utils.py (84%) create mode 100644 examples/speculative_decoding/prepare_data.sh rename {speculative_decoding => examples/speculative_decoding}/requirements.txt (80%) rename {speculative_decoding => examples/speculative_decoding}/vllm_generate.py (82%) create mode 100644 examples/vlm_eval/README.md create mode 100644 examples/vlm_eval/convert_gqa_for_eval.py create mode 100644 examples/vlm_eval/gqa.sh create mode 100644 examples/vlm_eval/model_gqa_loader.py rename {vlm_ptq => examples/vlm_eval}/requirements-vila.txt (73%) create mode 100644 examples/vlm_eval/vlm_eval_utils.py rename {vlm_ptq => examples/vlm_ptq}/.gitignore (100%) rename {vlm_ptq => examples/vlm_ptq}/README.md (89%) rename {vlm_ptq => examples/vlm_ptq}/requirements-mllama.txt (100%) create mode 100644 examples/vlm_ptq/requirements-vila.txt rename {vlm_ptq => examples/vlm_ptq}/scripts/huggingface_example.sh (71%) create mode 100644 examples/vlm_ptq/utils.py create mode 100644 examples/vlm_ptq/vlm_run.py create mode 100644 examples/vlm_ptq/vlm_visual_engine.py rename {windows => examples/windows}/Benchmark.md (100%) rename {windows => examples/windows}/README.md (95%) rename {windows => examples/windows}/accuracy_benchmark/README.md (100%) rename {windows => examples/windows}/accuracy_benchmark/mmlu_benchmark.py (94%) rename {windows => examples/windows}/accuracy_benchmark/modeling.py (90%) rename {windows => examples/windows}/accuracy_benchmark/quantization_utils.py (79%) rename {windows => examples/windows}/accuracy_benchmark/requirements.txt (86%) rename {windows => examples/windows}/accuracy_benchmark/trtllm_utils.py (92%) rename {windows => examples/windows}/onnx_ptq/README.md (100%) rename {windows => examples/windows}/onnx_ptq/quantize.py (92%) rename {windows => examples/windows}/onnx_ptq/requirements.txt (66%) delete mode 100644 llm_distill/requirements.txt delete mode 100644 llm_eval/__init__.py delete mode 100644 llm_ptq/benchmarks/benchmark_suite.py delete mode 100644 llm_ptq/benchmarks/disable_bench_chunked_context.py delete mode 100644 llm_sparsity/utils.py delete mode 100644 model_hub/download_hf_ckpt.py create mode 100644 modelopt/__init__.py create mode 100644 modelopt/deploy/__init__.py create mode 100644 modelopt/deploy/llm/__init__.py create mode 100644 modelopt/deploy/llm/generate.py create mode 100644 modelopt/deploy/llm/nemo_utils.py create mode 100644 modelopt/onnx/__init__.py create mode 100644 modelopt/onnx/op_types.py create mode 100644 modelopt/onnx/quantization/__init__.py create mode 100644 modelopt/onnx/quantization/__main__.py create mode 100644 modelopt/onnx/quantization/calib_utils.py create mode 100644 modelopt/onnx/quantization/extensions.py create mode 100755 modelopt/onnx/quantization/fp8.py create mode 100755 modelopt/onnx/quantization/graph_utils.py create mode 100644 modelopt/onnx/quantization/gs_patching.py create mode 100644 modelopt/onnx/quantization/int4.py create mode 100755 modelopt/onnx/quantization/int8.py create mode 100644 modelopt/onnx/quantization/operators.py create mode 100755 modelopt/onnx/quantization/ort_patching.py create mode 100755 modelopt/onnx/quantization/ort_utils.py create mode 100644 modelopt/onnx/quantization/partitioning.py create mode 100644 modelopt/onnx/quantization/qdq_utils.py create mode 100644 modelopt/onnx/quantization/quant_utils.py create mode 100755 modelopt/onnx/quantization/quantize.py create mode 100644 modelopt/onnx/quantization/src/modelopt_round_and_pack_ext.cpp create mode 100644 modelopt/onnx/quantization/trt_utils.py create mode 100644 modelopt/onnx/utils.py create mode 100644 modelopt/torch/__init__.py create mode 100644 modelopt/torch/_deploy/__init__.py create mode 100644 modelopt/torch/_deploy/_runtime/__init__.py create mode 100644 modelopt/torch/_deploy/_runtime/common.py create mode 100644 modelopt/torch/_deploy/_runtime/ort_client.py create mode 100644 modelopt/torch/_deploy/_runtime/registry.py create mode 100644 modelopt/torch/_deploy/_runtime/runtime_client.py create mode 100644 modelopt/torch/_deploy/_runtime/tensorrt/constants.py create mode 100644 modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py create mode 100644 modelopt/torch/_deploy/_runtime/tensorrt/hw_param_config.py create mode 100644 modelopt/torch/_deploy/_runtime/tensorrt/layerwise_profiling.py create mode 100644 modelopt/torch/_deploy/_runtime/tensorrt/parse_trtexec_log.py create mode 100644 modelopt/torch/_deploy/_runtime/tensorrt/tensorrt_utils.py create mode 100644 modelopt/torch/_deploy/_runtime/trt_client.py create mode 100644 modelopt/torch/_deploy/compilation.py create mode 100644 modelopt/torch/_deploy/device_model.py create mode 100644 modelopt/torch/_deploy/profiling.py create mode 100644 modelopt/torch/_deploy/utils/__init__.py create mode 100644 modelopt/torch/_deploy/utils/onnx_optimizer.py create mode 100644 modelopt/torch/_deploy/utils/onnx_utils.py create mode 100644 modelopt/torch/_deploy/utils/torch_onnx.py create mode 100644 modelopt/torch/distill/__init__.py create mode 100644 modelopt/torch/distill/config.py create mode 100644 modelopt/torch/distill/distillation.py create mode 100644 modelopt/torch/distill/distillation_model.py create mode 100644 modelopt/torch/distill/loss_balancers.py create mode 100644 modelopt/torch/distill/losses.py create mode 100644 modelopt/torch/distill/mode.py create mode 100644 modelopt/torch/distill/plugins/__init__.py create mode 100644 modelopt/torch/distill/plugins/megatron.py create mode 100644 modelopt/torch/distill/registry.py create mode 100644 modelopt/torch/export/__init__.py create mode 100644 modelopt/torch/export/distribute.py create mode 100644 modelopt/torch/export/hf_config_map.py create mode 100755 modelopt/torch/export/layer_utils.py create mode 100644 modelopt/torch/export/mcore_config_map.py create mode 100644 modelopt/torch/export/mcore_hf_export_map.py create mode 100755 modelopt/torch/export/model_config.py create mode 100644 modelopt/torch/export/model_config_export.py create mode 100755 modelopt/torch/export/model_config_utils.py create mode 100644 modelopt/torch/export/postprocess.py create mode 100644 modelopt/torch/export/quant_utils.py create mode 100644 modelopt/torch/export/tensorrt_llm_type.py create mode 100755 modelopt/torch/export/tensorrt_llm_utils.py create mode 100644 modelopt/torch/export/transformer_engine.py create mode 100644 modelopt/torch/export/unified_export_hf.py create mode 100644 modelopt/torch/export/unified_export_megatron.py create mode 100644 modelopt/torch/nas/__init__.py create mode 100644 modelopt/torch/nas/_algorithms.py create mode 100644 modelopt/torch/nas/_patch.py create mode 100644 modelopt/torch/nas/algorithms.py create mode 100644 modelopt/torch/nas/autonas.py create mode 100644 modelopt/torch/nas/config.py create mode 100644 modelopt/torch/nas/conversion.py create mode 100644 modelopt/torch/nas/hparams/__init__.py create mode 100644 modelopt/torch/nas/hparams/concat.py create mode 100644 modelopt/torch/nas/hparams/container.py create mode 100644 modelopt/torch/nas/mode.py create mode 100644 modelopt/torch/nas/modules/__init__.py create mode 100644 modelopt/torch/nas/modules/container.py create mode 100644 modelopt/torch/nas/modules/conv.py create mode 100644 modelopt/torch/nas/modules/linear.py create mode 100644 modelopt/torch/nas/modules/norm.py create mode 100644 modelopt/torch/nas/modules/utils.py create mode 100644 modelopt/torch/nas/plugins/__init__.py create mode 100644 modelopt/torch/nas/plugins/torch.py create mode 100644 modelopt/torch/nas/plugins/transformer_engine.py create mode 100644 modelopt/torch/nas/plugins/transformers.py create mode 100644 modelopt/torch/nas/registry.py create mode 100644 modelopt/torch/nas/search_space.py create mode 100644 modelopt/torch/nas/traced_hp.py create mode 100644 modelopt/torch/nas/utils.py create mode 100644 modelopt/torch/opt/__init__.py create mode 100644 modelopt/torch/opt/_hooks.py create mode 100644 modelopt/torch/opt/config.py create mode 100644 modelopt/torch/opt/conversion.py create mode 100644 modelopt/torch/opt/dynamic.py create mode 100644 modelopt/torch/opt/hparam.py create mode 100644 modelopt/torch/opt/mode.py create mode 100644 modelopt/torch/opt/plugins/__init__.py create mode 100644 modelopt/torch/opt/plugins/diffusers.py create mode 100644 modelopt/torch/opt/plugins/huggingface.py create mode 100644 modelopt/torch/opt/plugins/mcore_dist_checkpointing.py create mode 100644 modelopt/torch/opt/plugins/peft.py create mode 100644 modelopt/torch/opt/plugins/transformers.py create mode 100644 modelopt/torch/opt/searcher.py create mode 100644 modelopt/torch/opt/utils.py create mode 100644 modelopt/torch/prune/__init__.py create mode 100644 modelopt/torch/prune/config.py create mode 100644 modelopt/torch/prune/fastnas.py create mode 100644 modelopt/torch/prune/gradnas.py create mode 100644 modelopt/torch/prune/mcore_gpt_minitron.py create mode 100644 modelopt/torch/prune/mode.py create mode 100644 modelopt/torch/prune/plugins/__init__.py create mode 100644 modelopt/torch/prune/plugins/megatron.py create mode 100644 modelopt/torch/prune/plugins/transformers.py create mode 100644 modelopt/torch/prune/pruning.py create mode 100644 modelopt/torch/quantization/__init__.py create mode 100644 modelopt/torch/quantization/calib/__init__.py create mode 100644 modelopt/torch/quantization/calib/calibrator.py create mode 100644 modelopt/torch/quantization/calib/histogram.py create mode 100644 modelopt/torch/quantization/calib/max.py create mode 100644 modelopt/torch/quantization/config.py create mode 100644 modelopt/torch/quantization/conversion.py create mode 100644 modelopt/torch/quantization/export_onnx.py create mode 100644 modelopt/torch/quantization/extensions.py create mode 100644 modelopt/torch/quantization/mode.py create mode 100644 modelopt/torch/quantization/model_calib.py create mode 100644 modelopt/torch/quantization/model_quant.py create mode 100644 modelopt/torch/quantization/nn/__init__.py create mode 100644 modelopt/torch/quantization/nn/functional.py create mode 100644 modelopt/torch/quantization/nn/modules/__init__.py create mode 100644 modelopt/torch/quantization/nn/modules/clip.py create mode 100644 modelopt/torch/quantization/nn/modules/quant_activations.py create mode 100644 modelopt/torch/quantization/nn/modules/quant_batchnorm.py create mode 100644 modelopt/torch/quantization/nn/modules/quant_conv.py create mode 100644 modelopt/torch/quantization/nn/modules/quant_instancenorm.py create mode 100644 modelopt/torch/quantization/nn/modules/quant_linear.py create mode 100644 modelopt/torch/quantization/nn/modules/quant_module.py create mode 100644 modelopt/torch/quantization/nn/modules/quant_pooling.py create mode 100644 modelopt/torch/quantization/nn/modules/quant_rnn.py create mode 100644 modelopt/torch/quantization/nn/modules/tensor_quantizer.py create mode 100644 modelopt/torch/quantization/optim.py create mode 100644 modelopt/torch/quantization/plugins/__init__.py create mode 100644 modelopt/torch/quantization/plugins/apex.py create mode 100644 modelopt/torch/quantization/plugins/attention.py create mode 100644 modelopt/torch/quantization/plugins/custom.py create mode 100644 modelopt/torch/quantization/plugins/diffusers.py create mode 100644 modelopt/torch/quantization/plugins/fairscale.py create mode 100644 modelopt/torch/quantization/plugins/huggingface.py create mode 100644 modelopt/torch/quantization/plugins/megatron.py create mode 100644 modelopt/torch/quantization/plugins/nemo.py create mode 100644 modelopt/torch/quantization/plugins/peft.py create mode 100644 modelopt/torch/quantization/qtensor/__init__.py create mode 100644 modelopt/torch/quantization/qtensor/base_qtensor.py create mode 100644 modelopt/torch/quantization/qtensor/int4_tensor.py create mode 100644 modelopt/torch/quantization/qtensor/nf4_tensor.py create mode 100644 modelopt/torch/quantization/quant_modules.py create mode 100644 modelopt/torch/quantization/src/tensor_quant.cpp create mode 100644 modelopt/torch/quantization/src/tensor_quant.h create mode 100644 modelopt/torch/quantization/src/tensor_quant_fp8.cpp create mode 100644 modelopt/torch/quantization/src/tensor_quant_gpu.cu create mode 100644 modelopt/torch/quantization/src/tensor_quant_gpu_fp8.cu create mode 100644 modelopt/torch/quantization/src/tensor_quant_mx.cu create mode 100644 modelopt/torch/quantization/src/tensor_quant_mx.h create mode 100644 modelopt/torch/quantization/tensor_quant.py create mode 100644 modelopt/torch/quantization/utils.py create mode 100644 modelopt/torch/sparsity/__init__.py create mode 100644 modelopt/torch/sparsity/config.py create mode 100644 modelopt/torch/sparsity/magnitude.py create mode 100644 modelopt/torch/sparsity/mode.py create mode 100644 modelopt/torch/sparsity/module.py create mode 100644 modelopt/torch/sparsity/plugins/__init__.py create mode 100644 modelopt/torch/sparsity/plugins/megatron.py create mode 100644 modelopt/torch/sparsity/searcher.py create mode 100644 modelopt/torch/sparsity/sparsegpt.py create mode 100644 modelopt/torch/sparsity/sparsification.py create mode 100644 modelopt/torch/speculative/__init__.py create mode 100644 modelopt/torch/speculative/config.py create mode 100644 modelopt/torch/speculative/eagle/__init__.py create mode 100644 modelopt/torch/speculative/eagle/conversion.py create mode 100644 modelopt/torch/speculative/eagle/eagle_model.py create mode 100644 modelopt/torch/speculative/eagle/utils.py create mode 100644 modelopt/torch/speculative/medusa/__init__.py create mode 100644 modelopt/torch/speculative/medusa/conversion.py create mode 100644 modelopt/torch/speculative/medusa/medusa_model.py create mode 100644 modelopt/torch/speculative/mode.py create mode 100644 modelopt/torch/speculative/plugins/__init__.py create mode 100644 modelopt/torch/speculative/plugins/megatron.py create mode 100644 modelopt/torch/speculative/plugins/transformers.py create mode 100644 modelopt/torch/speculative/redrafter/__init__.py create mode 100644 modelopt/torch/speculative/redrafter/conversion.py create mode 100644 modelopt/torch/speculative/redrafter/redrafter_model.py create mode 100644 modelopt/torch/speculative/speculative_decoding.py create mode 100644 modelopt/torch/speculative/utils.py create mode 100644 modelopt/torch/trace/__init__.py create mode 100644 modelopt/torch/trace/analyzer.py create mode 100644 modelopt/torch/trace/modules/__init__.py create mode 100644 modelopt/torch/trace/modules/concat.py create mode 100644 modelopt/torch/trace/modules/nn.py create mode 100644 modelopt/torch/trace/plugins/__init__.py create mode 100644 modelopt/torch/trace/plugins/megatron.py create mode 100644 modelopt/torch/trace/plugins/transformers.py create mode 100644 modelopt/torch/trace/symbols.py create mode 100644 modelopt/torch/trace/tracer.py create mode 100644 modelopt/torch/utils/__init__.py create mode 100644 modelopt/torch/utils/_pytree.py create mode 100644 modelopt/torch/utils/cpp_extension.py rename {common => modelopt/torch/utils}/dataset_utils.py (73%) create mode 100644 modelopt/torch/utils/distributed.py create mode 100644 modelopt/torch/utils/graph.py create mode 100644 modelopt/torch/utils/image_processor.py create mode 100644 modelopt/torch/utils/import_utils.py create mode 100644 modelopt/torch/utils/list.py create mode 100644 modelopt/torch/utils/logging.py create mode 100644 modelopt/torch/utils/network.py create mode 100644 modelopt/torch/utils/perf.py create mode 100644 modelopt/torch/utils/random.py create mode 100644 modelopt/torch/utils/tensor.py create mode 100644 modelopt/torch/utils/vlm_dataset_utils.py create mode 100644 pyproject.toml create mode 100644 setup.py delete mode 100644 speculative_decoding/prepare_data.sh create mode 100644 tests/_test_utils/import_helper.py create mode 100644 tests/_test_utils/onnx_quantization/lib_test_models.py create mode 100644 tests/_test_utils/torch_deploy/device_model.py create mode 100644 tests/_test_utils/torch_deploy/runtime.py create mode 100644 tests/_test_utils/torch_dist/dist_utils.py create mode 100644 tests/_test_utils/torch_dist/fsdp_test.py create mode 100644 tests/_test_utils/torch_dist/plugins/megatron_common.py create mode 100644 tests/_test_utils/torch_misc.py create mode 100644 tests/_test_utils/torch_model/benchmark_models.py create mode 100644 tests/_test_utils/torch_model/deploy_models.py create mode 100644 tests/_test_utils/torch_model/example_models.py create mode 100644 tests/_test_utils/torch_model/utils.py create mode 100644 tests/_test_utils/torch_quantization/models.py create mode 100644 tests/_test_utils/torch_quantization/onnx_export.py create mode 100644 tests/_test_utils/torch_quantization/quant_utils.py create mode 100644 tests/_test_utils/torch_quantization/quantize_common.py create mode 100644 tests/_test_utils/torch_quantization/tensor_quant_common.py create mode 100644 tests/_test_utils/torch_quantization/tensor_quantizer_common.py create mode 100644 tests/conftest.py create mode 100644 tests/deploy/unit/test_deploy_dummy.py create mode 100644 tests/onnx/gpu/test_onnx_torch_int4_awq.py create mode 100755 tests/onnx/gpu/test_plugin.py create mode 100644 tests/onnx/unit/test_onnx_utils.py create mode 100644 tests/onnx/unit/test_partitioning.py create mode 100644 tests/onnx/unit/test_qdq_rules_int8.py create mode 100644 tests/onnx/unit/test_qdq_utils_fp8.py create mode 100644 tests/onnx/unit/test_quant_utils.py create mode 100644 tests/onnx/unit/test_quantize_fp8.py create mode 100644 tests/onnx/unit/test_quantize_int4.py create mode 100644 tests/onnx/unit/test_quantize_int8.py create mode 100644 tests/torch/conftest.py create mode 100644 tests/torch/gpu/conftest.py create mode 100644 tests/torch/gpu/deploy/_runtime/test_trt_client.py create mode 100644 tests/torch/gpu/deploy/test_gpu_deploy_benchmark.py create mode 100644 tests/torch/gpu/export/test_unified_export_megatron.py create mode 100644 tests/torch/gpu/nas/plugins/test_megatron_dynamic_modules.py create mode 100644 tests/torch/gpu/nas/test_distributed_model.py create mode 100644 tests/torch/gpu/nas/test_search_space_with_benchmark_models.py create mode 100644 tests/torch/gpu/opt/plugins/test_megatron_chaining.py create mode 100644 tests/torch/gpu/opt/test_fsdp_save_restore.py create mode 100644 tests/torch/gpu/prune/plugins/test_mcore_gpt_minitron_pruning.py create mode 100644 tests/torch/gpu/quantization/conftest.py create mode 100644 tests/torch/gpu/quantization/plugins/test_apex.py create mode 100644 tests/torch/gpu/quantization/plugins/test_megatron.py create mode 100644 tests/torch/gpu/quantization/test_fsdp.py create mode 100644 tests/torch/gpu/quantization/test_nvfp4_onnx_export.py create mode 100644 tests/torch/gpu/quantization/test_onnx_export_cuda.py create mode 100644 tests/torch/gpu/quantization/test_qtensor_cuda.py create mode 100644 tests/torch/gpu/quantization/test_quant_rnn_cuda.py create mode 100644 tests/torch/gpu/quantization/test_quantize_cuda.py create mode 100644 tests/torch/gpu/quantization/test_quantize_mxformats_cuda.py create mode 100644 tests/torch/gpu/quantization/test_real_quantize_cuda.py create mode 100644 tests/torch/gpu/quantization/test_tensor_quant_cuda.py create mode 100644 tests/torch/gpu/quantization/test_tensor_quantizer_cuda.py create mode 100644 tests/torch/gpu/quantization/test_torch_export.py create mode 100644 tests/torch/gpu/sparsity/plugins/test_megatron_sparsity.py create mode 100644 tests/torch/gpu/sparsity/test_sparse_fsdp.py create mode 100644 tests/torch/gpu/speculative/plugins/test_speculative_megatron_modules.py create mode 100644 tests/torch/integration/nas/test_full_algorithms.py create mode 100644 tests/torch/integration/quantization/test_onnx_export_cpu.py create mode 100644 tests/torch/unit/conftest.py create mode 100755 tests/torch/unit/deploy/_runtime/tensorrt/test_engine_builder.py create mode 100755 tests/torch/unit/deploy/_runtime/tensorrt/test_layerwise_profiling.py create mode 100755 tests/torch/unit/deploy/_runtime/tensorrt/test_tensorrt_utils.py create mode 100644 tests/torch/unit/deploy/conftest.py create mode 100644 tests/torch/unit/deploy/test_cpu_deploy_benchmark.py create mode 100644 tests/torch/unit/deploy/test_runtime_config.py create mode 100644 tests/torch/unit/deploy/utils/test_torch_onnx_utils.py create mode 100644 tests/torch/unit/distill/test_all.py create mode 100644 tests/torch/unit/export/test_get_quantization.py create mode 100644 tests/torch/unit/export/test_torch_export_config.py create mode 100644 tests/torch/unit/export/test_torch_export_dummy.py create mode 100644 tests/torch/unit/nas/modules/test_conv.py create mode 100644 tests/torch/unit/nas/modules/test_linear.py create mode 100644 tests/torch/unit/nas/modules/test_mod_utils.py create mode 100644 tests/torch/unit/nas/modules/test_norm.py create mode 100644 tests/torch/unit/nas/modules/test_sequential.py create mode 100644 tests/torch/unit/nas/plugins/test_dynamic_hf_attention.py create mode 100644 tests/torch/unit/nas/plugins/test_hf_nas_save_restore.py create mode 100644 tests/torch/unit/nas/plugins/test_torch_hooks.py create mode 100644 tests/torch/unit/nas/test_chaining.py create mode 100644 tests/torch/unit/nas/test_evaluate_constraints.py create mode 100644 tests/torch/unit/nas/test_latency_interpolator.py create mode 100644 tests/torch/unit/nas/test_nas.py create mode 100644 tests/torch/unit/nas/test_nas_utils.py create mode 100644 tests/torch/unit/nas/test_registry.py create mode 100644 tests/torch/unit/nas/test_search_space.py create mode 100644 tests/torch/unit/nas/test_search_space_with_example_models.py create mode 100644 tests/torch/unit/nas/test_search_space_with_torchvision.py create mode 100644 tests/torch/unit/opt/plugins/_transformers_helper.py create mode 100644 tests/torch/unit/opt/plugins/test_peft_save_load.py create mode 100644 tests/torch/unit/opt/plugins/test_transformers_save_load.py create mode 100644 tests/torch/unit/opt/test_config.py create mode 100644 tests/torch/unit/opt/test_dynamic.py create mode 100644 tests/torch/unit/prune/test_algorithms.py create mode 100644 tests/torch/unit/prune/test_fastnas_conversion.py create mode 100644 tests/torch/unit/prune/test_gradnas.py create mode 100644 tests/torch/unit/quantization/plugins/test_accelerate.py create mode 100644 tests/torch/unit/quantization/plugins/test_attention_quant.py create mode 100644 tests/torch/unit/quantization/plugins/test_huggingface.py create mode 100644 tests/torch/unit/quantization/plugins/test_peft.py create mode 100644 tests/torch/unit/quantization/test_autoquant.py create mode 100644 tests/torch/unit/quantization/test_calib.py create mode 100644 tests/torch/unit/quantization/test_calibrator.py create mode 100644 tests/torch/unit/quantization/test_clip.py create mode 100644 tests/torch/unit/quantization/test_dist.py create mode 100644 tests/torch/unit/quantization/test_module_registry.py create mode 100644 tests/torch/unit/quantization/test_print.py create mode 100644 tests/torch/unit/quantization/test_quant_activations.py create mode 100644 tests/torch/unit/quantization/test_quant_batchnorm.py create mode 100644 tests/torch/unit/quantization/test_quant_conv.py create mode 100644 tests/torch/unit/quantization/test_quant_instancenorm.py create mode 100644 tests/torch/unit/quantization/test_quant_linear.py create mode 100644 tests/torch/unit/quantization/test_quant_pooling.py create mode 100644 tests/torch/unit/quantization/test_quant_rnn.py create mode 100644 tests/torch/unit/quantization/test_quantize_cpu.py create mode 100644 tests/torch/unit/quantization/test_quantize_replace.py create mode 100644 tests/torch/unit/quantization/test_tensor_quant_cpu.py create mode 100644 tests/torch/unit/quantization/test_tensor_quantizer_cpu.py create mode 100644 tests/torch/unit/sparsity/test_sparsify.py create mode 100644 tests/torch/unit/speculative/plugins/test_hf_speculative.py create mode 100644 tests/torch/unit/trace/plugins/test_transformers_attention_symbols.py create mode 100644 tests/torch/unit/trace/test_analyzer.py create mode 100644 tests/torch/unit/trace/test_concat.py create mode 100644 tests/torch/unit/trace/test_model.py create mode 100644 tests/torch/unit/trace/test_nn.py create mode 100644 tests/torch/unit/trace/test_symbol.py create mode 100644 tests/torch/unit/trace/test_tracer.py create mode 100644 tests/torch/unit/utils/test_network.py create mode 100644 tests/torch/unit/utils/test_pytree.py create mode 100644 tests/torch/unit/utils/test_tensor.py delete mode 100644 vlm_ptq/requirements-llava.txt delete mode 100644 vlm_ptq/vlm_run.py delete mode 100644 vlm_ptq/vlm_visual_engine.py diff --git a/.dockerignore b/.dockerignore index 6a70d6d1ea..a5d89047c1 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,6 +1,7 @@ docker -**/.git -llm_ptq/saved_models* +examples/**/.git +examples/llm_ptq/saved_models* +**/experimental ##### Copied from .gitignore ##### # Byte-compiled / optimized / DLL files diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000000..0d0c65987d --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,145 @@ +# NOTE: Make sure to update version in dev requirements (setup.py) as well! +exclude: > + (?x)^( + experimental/.*| + )$ + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.6.0 + hooks: + - id: trailing-whitespace + - id: mixed-line-ending + args: [--fix=lf] + - id: end-of-file-fixer + - id: check-merge-conflict + - id: requirements-txt-fixer + - id: debug-statements + - id: check-json + exclude: ^.vscode/.*.json # vscode files can take comments + - id: check-yaml + args: [--allow-multiple-documents] + - id: check-toml + - id: check-added-large-files + args: [--maxkb=500, --enforce-all] + exclude: > + (?x)^( + examples/diffusers/quantization/assets/.*.png| + examples/diffusers/cache_diffusion/assets/.*.png| + )$ + + - repo: https://github.com/executablebooks/mdformat + rev: 0.7.17 + hooks: + - id: mdformat + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.6.4 + hooks: + - id: ruff + args: [--fix, --exit-non-zero-on-fix] + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.11.2 + hooks: + - id: mypy + + - repo: https://github.com/pre-commit/mirrors-clang-format + rev: v16.0.4 + hooks: + - id: clang-format + types_or: [c++, c, c#, cuda, java, javascript, objective-c, proto] # no json! + args: ["--style={ColumnLimit: 100}"] + + - repo: https://github.com/pre-commit/pygrep-hooks + rev: v1.10.0 + hooks: + - id: rst-backticks + - id: rst-directive-colons + - id: rst-inline-touching-normal + + - repo: https://github.com/jumanjihouse/pre-commit-hook-yamlfmt + rev: 0.2.3 + hooks: + - id: yamlfmt + args: [--mapping=2, --sequence=4, --offset=2, --implicit_start, --implicit_end, --preserve-quotes] + exclude: ^.github/workflows/ + + # Instructions to change license file if ever needed: + # https://github.com/Lucas-C/pre-commit-hooks#removing-old-license-and-replacing-it-with-a-new-one + - repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.5.5 + hooks: + # Default hook for Apache 2.0 in core python files + - id: insert-license + alias: insert-license-py + args: + - --license-filepath + - ./LICENSE + - --comment-style + - "#" + - --allow-past-years + types: [python] + # NOTE: Exclude files that have copyright or license headers from another company or individual + # since we want to keep those above the license header added by this hook. + # Instead, we should manually add the license header to those files after the original header. + exclude: > + (?x)^( + modelopt/onnx/quantization/operators.py| + modelopt/onnx/quantization/ort_patching.py| + modelopt/torch/export/transformer_engine.py| + modelopt/torch/quantization/export_onnx.py| + modelopt/torch/quantization/plugins/attention.py| + modelopt/torch/speculative/plugins/transformers.py| + modelopt/torch/speculative/eagle/utils.py| + modelopt/torch/_deploy/utils/onnx_utils.py| + examples/chained_optimizations/bert_prune_distill_quantize.py| + examples/diffusers/quantization/onnx_utils/export.py| + examples/diffusers/cache_diffusion/pipeline/models/sdxl.py| + examples/llm_eval/gen_model_answer.py| + examples/llm_eval/humaneval.py| + examples/llm_eval/lm_eval_hf.py| + examples/llm_eval/mmlu.py| + examples/llm_eval/modeling.py| + examples/llm_sparsity/finetune.py| + examples/llm_qat/main.py| + examples/speculative_decoding/main.py| + examples/speculative_decoding/medusa_utils.py| + examples/speculative_decoding/vllm_generate.py| + )$ + + # Default hook for Apache 2.0 in core c/c++/cuda files + - id: insert-license + alias: insert-license-c + args: + - --license-filepath + - ./LICENSE + - --comment-style + - "/*| *| */" + - --allow-past-years + types_or: [c++, cuda, c] + + # Default hook for Apache 2.0 in shell files + - id: insert-license + alias: insert-license-sh + args: + - --license-filepath + - ./LICENSE + - --comment-style + - "#" + - --allow-past-years + types_or: [shell] + + - repo: https://github.com/keith/pre-commit-buildifier + rev: 6.4.0 + hooks: + - id: buildifier + - id: buildifier-lint + + - repo: https://github.com/PyCQA/bandit + rev: 1.7.9 + hooks: + - id: bandit + args: ["-c", "pyproject.toml", "-q"] + additional_dependencies: ["bandit[toml]"] diff --git a/.vscode/extensions.json b/.vscode/extensions.json new file mode 100644 index 0000000000..f0c728ce98 --- /dev/null +++ b/.vscode/extensions.json @@ -0,0 +1,28 @@ +{ + // See https://go.microsoft.com/fwlink/?LinkId=827846 to learn about workspace recommendations. + // Extension identifier format: ${publisher}.${name}. Example: vscode.csharp + // List of extensions which should be recommended for users of this workspace. + "recommendations": [ + "ms-vscode.cpptools", + "ms-azuretools.vscode-docker", + "tamasfe.even-better-toml", + "GitHub.copilot", + "GitLab.gitlab-workflow", + "eamodio.gitlens", + "VisualStudioExptTeam.vscodeintellicode", + "ms-toolsai.jupyter", + "ms-python.vscode-pylance", + "ms-python.python", + "ms-vscode-remote.remote-ssh", + "ms-vscode.remote-explorer", + "charliermarsh.ruff", + "redhat.vscode-yaml", + ], + // List of extensions recommended by VS Code that should not be recommended for users of this workspace. + "unwantedRecommendations": [ + "ms-python.black-formatter", + "ms-python.mypy-type-checker", + "ms-python.pylint", + "ms-python.flake8", + ] +} diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000000..f68aaf626d --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,43 @@ +// VSCode workspace settings for modelopt +{ + "editor.rulers": [ + 100, + 120 + ], // 100 for black auto-formatter, 120 for hard limit in ruff + "[python]": { + "editor.defaultFormatter": "charliermarsh.ruff", + "editor.formatOnSave": true, + "editor.codeActionsOnSave": { + "source.fixAll": "explicit" + }, + }, + "files.exclude": { + "build": true, + }, + "files.watcherExclude": { + ".ipynb_checkpoints": true, + ".mypy_cache": true, + ".pytest_cache": true, + ".ruff_cache": true, + ".tox": true, + "**/__pycache__/**": true, + "**/*.pyc": true, + "**/runs": true, + "build": true + }, + "[yaml]": { + "editor.defaultFormatter": "redhat.vscode-yaml", + }, + "yaml.format.enable": true, + "yaml.format.printWidth": 150, + "yaml.format.bracketSpacing": false, + "yaml.customTags": [ + "!reference sequence" + ], + "python.testing.pytestEnabled": true, + "python.testing.pytestArgs": [ + "./tests", + "--no-cov", + ], + "evenBetterToml.schema.enabled": false, // disable toml/json schema since we have custom fields +} diff --git a/CHANGELOG-Windows.rst b/CHANGELOG-Windows.rst new file mode 100644 index 0000000000..53b4e388ce --- /dev/null +++ b/CHANGELOG-Windows.rst @@ -0,0 +1,18 @@ +=================================== +Model Optimizer Changelog (Windows) +=================================== + +0.19 (2024-11-18) +^^^^^^^^^^^^^^^^^ + +**New Features** + +- This is the first official release of TensorRT Model Optimizer for Windows +- **ONNX INT4 Quantization:** :meth:`modelopt.onnx.quantization.quantize_int4 ` now supports ONNX INT4 quantization for DirectML and TensorRT* deployment. See :ref:`Support_Matrix` for details about supported features and models. +- **LLM Quantization with Olive:** Enabled LLM quantization through Olive, streamlining model optimization workflows. Refer `example `_ +- **DirectML Deployment Guide:** Added DML deployment guide. Refer :ref:`DirectML_Deployment`. +- **MMLU Benchmark for Accuracy Evaluations:** Introduced `MMLU benchmarking `_ for accuracy evaluation of ONNX models on DirectML (DML). +- **Published quantized ONNX models collection:** Published quantized ONNX models at HuggingFace `NVIDIA collections `_. + + +\* *This version includes experimental features such as TensorRT deployment of ONNX INT4 models, PyTorch quantization and sparsity. These are currently unverified on Windows.* diff --git a/CHANGELOG.rst b/CHANGELOG.rst new file mode 100755 index 0000000000..b68f7bd13a --- /dev/null +++ b/CHANGELOG.rst @@ -0,0 +1,208 @@ +Model Optimizer Changelog (Linux) +================================= + +0.23 (2025-01-29) +^^^^^^^^^^^^^^^^^ + +**Backward Breaking Changes** + +- Nvidia TensorRT Model Optimizer has changed its LICENSE from NVIDIA Proprietary (library wheel) and MIT (examples) to Apache 2.0 in this first full OSS release. +- Deprecate Python 3.8, Torch 2.0, and Cuda 11.x support. +- ONNX Runtime dependency upgraded to 1.20 which no longer supports Python 3.9. +- In the Huggingface examples, the ``trust_remote_code`` is by default set to false and require users to explicitly turning it on with ``--trust_remote_code`` flag. + +**New Features** + +- Added OCP Microscaling Formats (MX) for fake quantization support, including FP8 (E5M2, E4M3), FP6 (E3M2, E2M3), FP4, INT8. +- Added NVFP4 quantization support for NVIDIA Blackwell GPUs along with updated examples. +- Allows export lm_head quantized TensorRT-LLM checkpoint. Quantize lm_head could benefit smaller sized models at a potential cost of additional accuracy loss. +- TensorRT-LLM now supports Moe FP8 and w4a8_awq inference on SM89 (Ada) GPUs. +- New models support in the ``llm_ptq`` example: Llama 3.3, Phi 4. +- Added Minitron pruning support for NeMo 2.0 GPT models. +- Exclude modules in TensorRT-LLM export configs are now wildcards +- The unified llama3.1 FP8 huggingface checkpoints can be deployed on `SGLang `_. + +0.21 (2024-12-03) +^^^^^^^^^^^^^^^^^ + +**Backward Breaking Changes** + +- Support TensorRT-LLM to 0.15. Examples (e.g. benchmark task in llm_ptq) may not be fully compatible with TensorRT-LLM 0.14. +- Remove the deprecated arg ``export_npz`` from the :meth:`mt.export.export_tensorrt_llm_checkpoint ` API +- Deprecate :meth:`mt.export.export_to_vllm ` API for :meth:`mt.export.export_hf_checkpoint ` +- Rename decoder type ``gptnext`` to ``gpt`` in ``llm_ptq`` to align with TensorRT-LLM model definition. + +**New Features** + +- Added new tutorial notebooks for Minitron pruning and distillation in NVIDIA NeMo framework. +- New models support in the ``llm_ptq`` example: Minitron, Phi3.5 MOE. +- New models support in the ``vlm_ptq`` example: Llama3.2(Mllama) +- :meth:`mt.export.export_tensorrt_llm_checkpoint ` and :meth:`mt.export.export_hf_checkpoint ` no longer requires the ``dtype`` arg. +- Added an example to deploy and run quantized fp8 llama3.1 8B instruct model from HuggingFace modelopt model hub on both TensorRT and vLLM. + +**Bug Fixes** + +- Improve Minitron pruning quality by avoiding possible bf16 overflow in importance calculation and minor change in ``hidden_size`` importance ranking. + +**Misc** + +- Added deprecation warnings for Python 3.8, torch 2.0, and CUDA 11.x. Support will be dropped in the next release. + +0.19 (2024-10-23) +^^^^^^^^^^^^^^^^^ + +**Backward Breaking Changes** + +- Deprecated the summarize task in the ``llm_ptq`` example. +- Deprecated the ``type`` flag in the `huggingface_example.sh `_ +- Deprecated Python plugin support in ONNX. +- Support TensorRT-LLM 0.13. Examples not compatible with TensorRT-LLM 0.12. +- :meth:`mtq.auto_quantize ` API has been updated. The API now + accepts ``forward_step`` and ``forward_backward_step`` as arguments instead of ``loss_func`` and ``collect_func``. + Please see the API documentation for more details. + +**New Features** + +- ModelOpt is compatible for SBSA aarch64 (e.g. GH200) now! + Except ONNX PTQ with plugins is not supported. +- Add ``effective_bits`` as a constraint for :meth:`mtq.auto_qauntize `. +- ``lm_evaluation_harness`` is fully integrated to modelopt backed by TensorRT-LLM. + ``lm_evaluation_harness`` benchmarks are now available in the examples for LLM accuracy evaluation. +- A new ``--perf`` flag is introduced in the ``modelopt_to_tensorrt_llm.py`` example to build engines with max perf. +- Users can choose the execution provider to run the calibration in ONNX quantization. +- Added automatic detection of custom ops in ONNX models using TensorRT plugins. + This requires the ``tensorrt`` python package to be installed. +- Replaced ``jax`` with ``cupy`` for faster INT4 ONNX quantization. +- :meth:`mtq.auto_quantize ` now supports search based automatic + quantization for NeMo & MCore models (in addition to HuggingFace models). +- Add ``num_layers`` and ``hidden_size`` pruning support for NeMo / Megatron-core models. + + +0.17 (2024-09-11) +^^^^^^^^^^^^^^^^^ + +**Backward Breaking Changes** + +- Deprecated ``torch<2.0`` support. +- :meth:`modelopt.torch.utils.dataset_utils.get_dataset_dataloader` now returns a key value pair instead of the tensor. + +**New Features** + +- New APIs and examples: :mod:`modelopt.torch.prune` for pruning Conv, Linear, and Attention heads for + NVIDIA Megatron-core GPT-style models (e.g. Llama 3), PyTorch Computer Vision models, and HuggingFace Bert/GPT-J models. +- New API: :mod:`modelopt.torch.distill` for knowledge distillation, along with guides and example. +- New Example: `HF BERT Prune, Distill & Quantize `_ + showcasing how to chain pruning, distillation, and quantization to achieve the best performance on a given model. +- Added INT8/FP8 DQ-only support for ONNX model. +- New API: :mod:`modelopt.torch.speculative` for end-to-end support of Medusa models. +- Added Medusa QAT and End-to-end examples. +- Modelopt now supports automatic save/restore of ``modelopt_state`` with the ``.save_pretrained`` and ``.from_pretrained`` APIs + from Huggingface libraries, such as ``transformers`` and ``diffusers``. This feature can be enabled by calling + :meth:`mto.enable_huggingface_checkpointing() `. +- ONNX FP8 quantization support with amax calibration. +- TensorRT-LLM dependency upgraded to 0.12.0. Huggingface tokenizer files are now also stored in the engine dir. +- The unified model export API :meth:`modelopt.torch.export.export_hf_checkpoint ` + supports exporting ``fp8`` and ``int4_awq`` quantized checkpoints with packed weights for + Hugging Face models with namings aligned with its original checkpoints. The exported ``fp8`` checkpoints can be deployed with both TensorRT-LLM and VLLM. +- Add int8 and fp8 quantization support for the FLUX.1-dev model. +- Add a Python-friendly TensorRT inference pipeline for diffusion models. + +**Misc** + +- Added deprecation warning for :meth:`set_data_parallel_group ` + and :meth:`set_tensor_parallel_group `. These APIs are + no longer needed for supporting distributed data and tensor parallelism in quantization. They will be removed in + a future release. + + +0.15 (2024-07-25) +^^^^^^^^^^^^^^^^^ + +**Backward Breaking Changes** + +- Deprecated :class:`QuantDescriptor `. + Use :class:`QuantizerAttributeConfig ` to + configure :class:`TensorQuantizer `. + :meth:`set_from_attribute_config ` + can be used to set the quantizer attributes from the config class or attribute dictionary. This change applies only + to backend APIs. The change is backward compatible if you are using + only the :meth:`mtq.quantize ` API. + +**New Features** + +- Added quantization support for torch ``RNN, LSTM, GRU`` modules. Only available for ``torch>=2.0``. +- ``modelopt.torch.quantization`` now supports module class based quantizer attribute setting for + :meth:`mtq.quantize ` API. +- Added new LLM PTQ example for DBRX model. +- Added new LLM (Gemma 2) PTQ and TensorRT-LLM checkpoint export support. +- Added new LLM QAT example for NVIDIA NeMo framework. +- TensorRT-LLM dependency upgraded to 0.11.0. +- (Experimental): :meth:`mtq.auto_quantize ` API which quantizes a model + by searching for the best per-layer quantization formats. +- (Experimental): Added new LLM QLoRA example with NF4 and INT4_AWQ quantization. +- (Experimental): ``modelopt.torch.export`` now supports exporting quantized checkpoints with packed weights for + Hugging Face models with namings aligned with its original checkpoints. +- (Experimental) Added support for quantization of ONNX models with TensorRT plugin. + +**Misc** + +- Added deprecation warning for ``torch<2.0``. Support will be dropped in next release. + + +0.13 (2024-06-14) +^^^^^^^^^^^^^^^^^ + +**Backward Breaking Changes** + +- `PTQ examples `_ have been + upgraded to use TensorRT-LLM 0.10. + +**New Features** + +- Adding TensorRT-LLM checkpoint export support for Medusa decoding (official ``MedusaModel`` and Megatron Core ``GPTModel``). +- Enable support for mixtral, recurrentgemma, starcoder, qwen in `PTQ examples `_. +- Adding TensorRT-LLM checkpoint export and engine building support for sparse models. +- Import scales from TensorRT calibration cache and use them for quantization. +- (Experimental) Enable low GPU memory FP8 calibration for the Hugging Face models when the original model size does not fit into the GPU memory. +- (Experimental) Support exporting FP8 calibrated model to VLLM deployment. +- (Experimental) Python 3.12 support added. + + +0.11 (2024-05-07) +^^^^^^^^^^^^^^^^^ + +**Backward Breaking Changes** + +- [!!!] The package was renamed from ``ammo`` to ``modelopt``. The new full product + name is *Nvidia TensorRT Model Optimizer*. PLEASE CHANGE ALL YOUR REFERENCES FROM ``ammo`` to + ``modelopt`` including any paths and links! +- Default installation ``pip install nvidia-modelopt`` will now only install minimal core + dependencies. Following optional dependencies are available depending on the features that are + being used: ``[deploy], [onnx], [torch], [hf]``. To install all dependencies, use + ``pip install "nvidia-modelopt[all]"``. +- Deprecated ``inference_gpus`` arg in ``modelopt.torch.export.model_config_export.torch_to_tensorrt_llm_checkpoint``. User should use ``inference_tensor_parallel`` instead. +- Experimental ``modelopt.torch.deploy`` module is now available as ``modelopt.torch._deploy``. + +**New Features** + +- ``modelopt.torch.sparsity`` now supports sparsity-aware training (SAT). Both SAT and post-training + sparsification supports chaining with other modes, e.g. SAT + QAT. +- ``modelopt.torch.quantization`` natively support distributed data and tensor parallelism while estimating quantization parameters. + The data and tensor parallel groups needs to be registered with ``modelopt.torch.utils.distributed.set_data_parallel_group`` and ``modelopt.torch.utils.distributed.set_tensor_parallel_group`` APIs. + By default, the data parallel group is set as the default distributed group and the tensor parallel group is disabled. +- ``modelopt.torch.opt`` now supports chaining multiple optimization techniques that each require + modifications to the same model, e.g., you can now sparsify and quantize a model at the same time. +- ``modelopt.onnx.quantization`` supports FLOAT8 quantization format with Distribution calibration algorithm. +- Native support of ``modelopt.torch.opt`` with FSDP (Fully Sharded Data Parallel) for ``torch>=2.1``. This includes + sparsity, quantization, and any other model modification & optimization. +- Added FP8 ONNX quantization support in ``modelopt.onnx.quantization``. +- Added Windows (``win_amd64``) support for ModelOpt released wheels. Currently supported for ``modelopt.onnx`` submodule only. + +**Bug Fixes** + +- Fixed the compatibility issue of ``modelopt.torch.sparsity`` with FSDP. +- Fixed an issue in dynamic dim handling in ``modelopt.onnx.quantization`` with random calibration data. +- Fixed graph node naming issue after opset conversion operation. +- Fixed an issue in negative dim handling like dynamic dim in ``modelopt.onnx.quantization`` with random calibration data. +- Fixed allowing to accept ``.pb`` file for input file. +- Fixed copy extra data to tmp folder issue for ONNX PTQ. diff --git a/LICENSE b/LICENSE index cf4ede00fb..7b0d9f06b0 100644 --- a/LICENSE +++ b/LICENSE @@ -1,20 +1,14 @@ SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -SPDX-License-Identifier: MIT +SPDX-License-Identifier: Apache-2.0 -Permission is hereby granted, free of charge, to any person obtaining a -copy of this software and associated documentation files (the "Software"), -to deal in the Software without restriction, including without limitation -the rights to use, copy, modify, merge, publish, distribute, sublicense, -and/or sell copies of the Software, and to permit persons to whom the -Software is furnished to do so, subject to the following conditions: +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 -The above copyright notice and this permission notice shall be included in -all copies or substantial portions of the Software. +http://www.apache.org/licenses/LICENSE-2.0 -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -DEALINGS IN THE SOFTWARE. +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/README.md b/README.md index 0aef01f78d..049fe32fb0 100644 --- a/README.md +++ b/README.md @@ -6,13 +6,13 @@ [![Documentation](https://img.shields.io/badge/Documentation-latest-brightgreen.svg?style=flat)](https://nvidia.github.io/TensorRT-Model-Optimizer) [![version](https://img.shields.io/pypi/v/nvidia-modelopt?label=Release)](https://pypi.org/project/nvidia-modelopt/) -[![license](https://img.shields.io/badge/License-MIT-blue)](./LICENSE) +[![license](https://img.shields.io/badge/License-Apache%202.0-blue)](./LICENSE) [Examples](#examples) | [Documentation](https://nvidia.github.io/TensorRT-Model-Optimizer) | [Benchmark Results](#benchmark) | -[Roadmap](https://github.com/NVIDIA/TensorRT-Model-Optimizer/issues/108) | -[ModelOpt-Windows](./windows/README.md) +[Roadmap](#roadmap) | +[ModelOpt-Windows](./examples/windows/README.md) @@ -22,8 +22,8 @@ - \[2024/9/10\] [Post-Training Quantization of LLMs with NVIDIA NeMo and TensorRT Model Optimizer](https://developer.nvidia.com/blog/post-training-quantization-of-llms-with-nvidia-nemo-and-nvidia-tensorrt-model-optimizer/) - \[2024/8/28\] [Boosting Llama 3.1 405B Performance up to 44% with TensorRT Model Optimizer on NVIDIA H200 GPUs](https://developer.nvidia.com/blog/boosting-llama-3-1-405b-performance-by-up-to-44-with-nvidia-tensorrt-model-optimizer-on-nvidia-h200-gpus/) - \[2024/8/28\] [Up to 1.9X Higher Llama 3.1 Performance with Medusa](https://developer.nvidia.com/blog/low-latency-inference-chapter-1-up-to-1-9x-higher-llama-3-1-performance-with-medusa-on-nvidia-hgx-h200-with-nvlink-switch/) -- \[2024/08/15\] New features in recent releases: [Cache Diffusion](https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/diffusers/cache_diffusion), [QLoRA workflow with NVIDIA NeMo](https://docs.nvidia.com/nemo-framework/user-guide/latest/sft_peft/qlora.html), and more. Check out [our blog](https://developer.nvidia.com/blog/nvidia-tensorrt-model-optimizer-v0-15-boosts-inference-performance-and-expands-model-support/) for details. -- \[2024/06/03\] Model Optimizer now has an experimental feature to deploy to vLLM as part of our effort to support popular deployment frameworks. Check out the workflow [here](./llm_ptq/README.md#deploy-fp8-quantized-model-using-vllm) +- \[2024/08/15\] New features in recent releases: [Cache Diffusion](https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/diffusers/cache_diffusion), [QLoRA workflow with NVIDIA NeMo](https://docs.nvidia.com/nemo-framework/user-guide/latest/sft_peft/qlora.html), and more. Check out [our blog](https://developer.nvidia.com/blog/nvidia-tensorrt-model-optimizer-v0-15-boosts-inference-performance-and-expands-model-support/) for details. +- \[2024/06/03\] Model Optimizer now has an experimental feature to deploy to vLLM as part of our effort to support popular deployment frameworks. Check out the workflow [here](./examples/llm_ptq/README.md#deploy-fp8-quantized-model-using-vllm) - \[2024/05/08\] [Announcement: Model Optimizer Now Formally Available to Further Accelerate GenAI Inference Performance](https://developer.nvidia.com/blog/accelerate-generative-ai-inference-performance-with-nvidia-tensorrt-model-optimizer-now-publicly-available/) - \[2024/03/27\] [Model Optimizer supercharges TensorRT-LLM to set MLPerf LLM inference records](https://developer.nvidia.com/blog/nvidia-h200-tensor-core-gpus-and-nvidia-tensorrt-llm-set-mlperf-llm-inference-records/) - \[2024/03/18\] [GTC Session: Optimize Generative AI Inference with Quantization in TensorRT-LLM and TensorRT](https://www.nvidia.com/en-us/on-demand/session/gtc24-s63213/) @@ -56,11 +56,9 @@ Seamlessly integrated within the NVIDIA AI software ecosystem, the quantized che ModelOpt is integrated with [NVIDIA NeMo](https://github.com/NVIDIA/NeMo) and [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) for training-in-the-loop optimization techniques. For enterprise users, the 8-bit quantization with Stable Diffusion is also available on [NVIDIA NIM](https://developer.nvidia.com/blog/nvidia-nim-offers-optimized-inference-microservices-for-deploying-ai-models-at-scale/). -Model Optimizer for both Linux and Windows are available for free for all developers on [NVIDIA PyPI](https://pypi.org/project/nvidia-modelopt/). This repository is for sharing examples and GPU-optimized recipes as well as collecting feedback from the community. - ## Installation / Docker -Easiest way to get started with using Model Optimizer and additional dependencies (e.g. TensorRT-LLM deployment) is to start from our docker image. +To use Model Optimizer with full dependencies (e.g. TensorRT-LLM deployment), we recommend using the provided docker image. After installing the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html), please run the following commands to build the Model Optimizer docker container which has all the necessary @@ -82,6 +80,20 @@ docker run --gpus all -it --shm-size 20g --rm docker.io/library/modelopt_example python -c "import modelopt; print(modelopt.__version__)" ``` +Alternatively, you can install it from [NVIDIA PyPI](https://pypi.org/project/nvidia-modelopt/) without TRT-LLM etc. + +```bash +pip install "nvidia-modelopt[all]" -U --extra-index-url https://pypi.nvidia.com +``` + +To install from source for local development, you can install it as follows: + +```bash +pip install -e ".[all]" --extra-index-url https://pypi.nvidia.com +``` + +When installing from source, please make sure to re-run the install command everytime you pull new changes in the repository so dependencies are also updated. + See the [installation guide](https://nvidia.github.io/TensorRT-Model-Optimizer/getting_started/2_installation.html) for more details on alternate pre-built docker images or installation in a local environment. NOTE: Unless specified otherwise, all example READMEs assume they are using the above ModelOpt docker image for running the examples. The example specific dependencies are required to be install separately from their respective `requirements.txt` files if not using the ModelOpt's docker image. @@ -107,33 +119,18 @@ Sparsity is a technique to further reduce the memory footprint of deep learning ## Examples -- [PTQ for LLMs](./llm_ptq/README.md) covers how to use Post-training quantization (PTQ) and export to [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) for deployment for popular pre-trained models from frameworks like - - [Hugging Face](https://huggingface.co/docs/hub/en/models-the-hub) - - [NVIDIA NeMo](https://github.com/NVIDIA/NeMo) - - [NVIDIA Megatron-LM](https://github.com/NVIDIA/Megatron-LM) - - [Medusa](https://github.com/FasterDecoding/Medusa) -- [PTQ for Diffusers](./diffusers/quantization/README.md) walks through how to quantize a diffusion model with FP8 or INT8, export to ONNX, and deploy with [TensorRT](https://github.com/NVIDIA/TensorRT/tree/release/10.0/demo/Diffusion). The Diffusers example in this repo is complementary to the [demoDiffusion example in TensorRT repo](https://github.com/NVIDIA/TensorRT/tree/release/10.0/demo/Diffusion#introduction) and includes FP8 plugins as well as the latest updates on INT8 quantization. -- [QAT for LLMs](./llm_qat/README.md) demonstrates the recipe and workflow for Quantization-aware Training (QAT), which can further preserve model accuracy at low precisions (e.g., INT4, or 4-bit in [NVIDIA Blackwell platform](https://www.nvidia.com/en-us/data-center/technologies/blackwell-architecture/)). -- [Sparsity for LLMs](./llm_sparsity/README.md) shows how to perform Post-training Sparsification and Sparsity-aware fine-tuning on a pre-trained Hugging Face model. -- [Pruning](./pruning/README.md) demonstrates how to optimally prune Linear and Conv layers, and Transformer attention heads, MLP, and depth using the Model Optimizer for following frameworks: - - [NVIDIA NeMo](https://github.com/NVIDIA/NeMo) / [NVIDIA Megatron-LM](https://github.com/NVIDIA/Megatron-LM) GPT-style models (e.g. Llama 3, Mistral NeMo, etc.) - - Hugging Face language models like BERT and GPT-J - - Computer Vision models like [NVIDIA Tao](https://developer.nvidia.com/tao-toolkit) framework detection models. -- [ONNX PTQ](./onnx_ptq/README.md) shows how to quantize the ONNX models in INT4 or INT8 quantization mode. The examples also include the deployment of quantized ONNX models using TensorRT. -- [Distillation for LLMs](./llm_distill/README.md) demonstrates how to use Knowledge Distillation, which can increasing the accuracy and/or convergence speed for finetuning / QAT. -- [Chained Optimizations](./chained_optimizations/README.md) shows how to chain multiple optimizations together (e.g. Pruning + Distillation + Quantization). -- [Model Hub](./model_hub/) provides an example to deploy and run quantized Llama 3.1 8B instruct model from Nvidia's Hugging Face model hub on both TensorRT-LLM and vLLM. +Please see examples [here](./examples/README.md). ## Model Support Matrix -- For LLM quantization, please refer to this [support matrix](./llm_ptq/README.md#model-support-list). -- For VLM quantization, please refer to this [support matrix](./vlm_ptq/README.md#model-support-list). +- For LLM quantization, please refer to this [support matrix](./examples/llm_ptq/README.md#model-support-list). +- For VLM quantization, please refer to this [support matrix](./examples/vlm_ptq/README.md#model-support-list). - For Diffusion, Model Optimizer supports [FLUX](https://huggingface.co/black-forest-labs/FLUX.1-dev), [Stable Diffusion 3](https://huggingface.co/stabilityai/stable-diffusion-3-medium), [Stable Diffusion XL](https://huggingface.co/papers/2307.01952), [SDXL-Turbo](https://huggingface.co/stabilityai/sdxl-turbo), and [Stable Diffusion 2.1](https://huggingface.co/stabilityai/stable-diffusion-2-1). -- For speculative decoding, please refer to this [support matrix](./speculative_decoding/README.md#model-support-list). +- For speculative decoding, please refer to this [support matrix](./examples/speculative_decoding/README.md#model-support-list). ## Benchmark -Please find the benchmarks at [here](./benchmark.md). +Please find the benchmarks at [here](./examples/benchmark.md). ## Quantized Checkpoints @@ -149,4 +146,4 @@ Please see Model Optimizer Changelog [here](https://nvidia.github.io/TensorRT-Mo ## Contributing -At the moment, we are not accepting external contributions. However, this will soon change after we open source our library in early 2025 with a focus on extensibility. We welcome any feedback and feature requests. Please open an issue if you have any suggestions or questions. +At the moment, we are not accepting external contributions. However, this will soon change with a focus on extensibility. We welcome any feedback and feature requests. Please open an issue if you have any suggestions or questions. diff --git a/chained_optimizations/README.md b/chained_optimizations/README.md deleted file mode 100644 index 8f82b5f384..0000000000 --- a/chained_optimizations/README.md +++ /dev/null @@ -1,18 +0,0 @@ -# Chaining Multiple Optimizations Techniques - -This example demonstrates how to chain multiple optimization techniques like Pruning, Distillation, and Quantization together to -achieve the best performance on a given model. - -## HuggingFace BERT Pruning + Distillation + Quantization - -This example shows how to compress a [Hugging Face Bert large model for Question Answering](https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad) -using the combination of `modelopt.torch.prune`, `modelopt.torch.distill` and `modelopt.torch.quantize`. More specifically, we will: - -1. Prune the Bert large model to 50% FLOPs with GradNAS algorithm and fine-tune with distillation -1. Quantize the fine-tuned model to INT8 precision with Post-Training Quantization (PTQ) and Quantize Aware Training (QAT) with distillation -1. Export the quantized model to ONNX format for deployment with TensorRT - -The main python file is [bert_prune_distill_quantize.py](./bert_prune_distill_quantize.py) and scripts for running it -for all 3 steps are available in the [scripts](./scripts/) directory. -More details on this example (including highlighted code snippets) can be found in the Model Optimizer documentation -[here](https://nvidia.github.io/TensorRT-Model-Optimizer/examples/2_bert_prune_distill_quantize.html) diff --git a/chained_optimizations/scripts/3_onnx_export.sh b/chained_optimizations/scripts/3_onnx_export.sh deleted file mode 100755 index 9ef5dd2931..0000000000 --- a/chained_optimizations/scripts/3_onnx_export.sh +++ /dev/null @@ -1,10 +0,0 @@ -#!/bin/bash -set -ex - -FLOPS_PERCENT=50 -BASE_DIR=results/bert_large_pruned_${FLOPS_PERCENT}_percent - -# Export to ONNX on a single GPU -python3 bert_prune_distill_quantize.py \ - --model_name_or_path ${BASE_DIR}/int8_qat/ \ - --onnx_export_path ${BASE_DIR}/pruned_model_int8.onnx \ diff --git a/diffusers/cache_diffusion/cache_diffusion/module.py b/diffusers/cache_diffusion/cache_diffusion/module.py deleted file mode 100644 index a1ed434c30..0000000000 --- a/diffusers/cache_diffusion/cache_diffusion/module.py +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT -# -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. - -from torch import nn - - -class CachedModule(nn.Module): - def __init__(self, block, select_cache_step_func) -> None: - super().__init__() - self.block = block - self.select_cache_step_func = select_cache_step_func - self.cur_step = 0 - self.cached_results = None - self.enabled = True - - def __getattr__(self, name): - try: - return super().__getattr__(name) - except AttributeError: - return getattr(self.block, name) - - def if_cache(self): - return self.select_cache_step_func(self.cur_step) and self.enabled - - def enable_cache(self): - self.enabled = True - - def disable_cache(self): - self.enabled = False - self.cur_step = 0 - - def forward(self, *args, **kwargs): - if not self.if_cache(): - self.cached_results = self.block(*args, **kwargs) - if self.enabled: - self.cur_step += 1 - return self.cached_results diff --git a/docker/Dockerfile b/docker/Dockerfile index becad2638b..e94d351d33 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -2,9 +2,9 @@ FROM nvidia/cuda:12.6.2-devel-ubuntu22.04 WORKDIR /workspace -RUN apt-get update && apt-get -y install python3.10 python3-pip python-is-python3 openmpi-bin libopenmpi-dev wget git git-lfs unzip jq +RUN apt-get update && apt-get -y install python3.10 python3-pip python-is-python3 openmpi-bin libopenmpi-dev wget git git-lfs unzip jq cmake -ARG PIP_EXTRA_INDEX_URL="https://pypi.nvidia.com https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/" +ARG PIP_EXTRA_INDEX_URL="https://pypi.nvidia.com" ENV PIP_EXTRA_INDEX_URL=$PIP_EXTRA_INDEX_URL ENV PIP_NO_CACHE_DIR=off @@ -44,9 +44,9 @@ RUN python -c "import modelopt.torch.quantization.extensions as ext; ext.precomp # Find and install requirements.txt files for all examples excluding windows COPY . TensorRT-Model-Optimizer -RUN find TensorRT-Model-Optimizer -name "requirements.txt" | grep -v "windows" | while read req_file; do \ - echo "Installing from $req_file"; \ - pip install -r "$req_file"; \ +RUN find TensorRT-Model-Optimizer/examples -name "requirements.txt" | grep -v "windows" | while read req_file; do \ + echo "Installing from $req_file"; \ + pip install -r "$req_file"; \ done # Allow users to run without root diff --git a/docker/build.sh b/docker/build.sh index 9d35b294a3..da8f00bca3 100755 --- a/docker/build.sh +++ b/docker/build.sh @@ -1,9 +1,18 @@ +# 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. + set -e -# Resolve symlinks if there is any -find . -type l -exec sh -c 'cp -L "$0" "$0.tmp" && mv "$0.tmp" "$0"' {} \; - docker build --progress=plain . -f docker/Dockerfile -t modelopt_examples:latest "$@" - -# restore symlinks if there is any -git status | grep 'typechange:' | awk '{print $2}' | xargs git checkout -- diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000000..83e11ba4fc --- /dev/null +++ b/examples/README.md @@ -0,0 +1,18 @@ +# NVIDIA TensorRT Model Optimizer Examples + +- [PTQ for LLMs](./llm_ptq/README.md) covers how to use Post-training quantization (PTQ) and export to [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) for deployment for popular pre-trained models from frameworks like + - [Hugging Face](https://huggingface.co/docs/hub/en/models-the-hub) + - [NVIDIA NeMo](https://github.com/NVIDIA/NeMo) + - [NVIDIA Megatron-LM](https://github.com/NVIDIA/Megatron-LM) + - [Medusa](https://github.com/FasterDecoding/Medusa) +- [PTQ for Diffusers](./diffusers/quantization/README.md) walks through how to quantize a diffusion model with FP8 or INT8, export to ONNX, and deploy with [TensorRT](https://github.com/NVIDIA/TensorRT/tree/release/10.0/demo/Diffusion). The Diffusers example in this repo is complementary to the [demoDiffusion example in TensorRT repo](https://github.com/NVIDIA/TensorRT/tree/release/10.0/demo/Diffusion#introduction) and includes FP8 plugins as well as the latest updates on INT8 quantization. +- [QAT for LLMs](./llm_qat/README.md) demonstrates the recipe and workflow for Quantization-aware Training (QAT), which can further preserve model accuracy at low precisions (e.g., INT4, or 4-bit in [NVIDIA Blackwell platform](https://www.nvidia.com/en-us/data-center/technologies/blackwell-architecture/)). +- [Sparsity for LLMs](./llm_sparsity/README.md) shows how to perform Post-training Sparsification and Sparsity-aware fine-tuning on a pre-trained Hugging Face model. +- [Pruning](./pruning/README.md) demonstrates how to optimally prune Linear and Conv layers, and Transformer attention heads, MLP, and depth using the Model Optimizer for following frameworks: + - [NVIDIA NeMo](https://github.com/NVIDIA/NeMo) / [NVIDIA Megatron-LM](https://github.com/NVIDIA/Megatron-LM) GPT-style models (e.g. Llama 3, Mistral NeMo, etc.) + - Hugging Face language models like BERT and GPT-J + - Computer Vision models like [NVIDIA Tao](https://developer.nvidia.com/tao-toolkit) framework detection models. +- [ONNX PTQ](./onnx_ptq/README.md) shows how to quantize the ONNX models in INT4 or INT8 quantization mode. The examples also include the deployment of quantized ONNX models using TensorRT. +- [Distillation for LLMs](./llm_distill/README.md) demonstrates how to use Knowledge Distillation, which can increasing the accuracy and/or convergence speed for finetuning / QAT. +- [Chained Optimizations](./chained_optimizations/README.md) shows how to chain multiple optimizations together (e.g. Pruning + Distillation + Quantization). +- [Model Hub](./model_hub/) provides an example to deploy and run quantized Llama 3.1 8B instruct model from Nvidia's Hugging Face model hub on both TensorRT-LLM and vLLM. diff --git a/benchmark.md b/examples/benchmark.md similarity index 100% rename from benchmark.md rename to examples/benchmark.md diff --git a/chained_optimizations/.gitignore b/examples/chained_optimizations/.gitignore similarity index 100% rename from chained_optimizations/.gitignore rename to examples/chained_optimizations/.gitignore diff --git a/examples/chained_optimizations/README.md b/examples/chained_optimizations/README.md new file mode 100644 index 0000000000..969d14d44d --- /dev/null +++ b/examples/chained_optimizations/README.md @@ -0,0 +1,56 @@ +# Chaining Multiple Optimizations Techniques + +This directory demonstrates how to chain multiple optimization techniques like Pruning, Distillation, and Quantization together to +achieve the best performance on a given model. + +## HuggingFace BERT Pruning + Distillation + Quantization + +This example shows how to compress a [Hugging Face Bert large model for Question Answering](https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad) +using the combination of `modelopt.torch.prune`, `modelopt.torch.distill` and `modelopt.torch.quantize`. More specifically, we will: + +1. Prune the Bert large model to 50% FLOPs with GradNAS algorithm and fine-tune with distillation +1. Quantize the fine-tuned model to INT8 precision with Post-Training Quantization (PTQ) and Quantize Aware Training (QAT) with distillation +1. Export the quantized model to ONNX format for deployment with TensorRT + +The main python file is [bert_prune_distill_quantize.py](./bert_prune_distill_quantize.py) and scripts for running it +for all 3 steps are available in the [scripts](./scripts/) directory. + +NOTE: This example has been tested on 8 x 24GB A5000 GPUs with PyTorch 2.4 and CUDA 12.4. It takes +about 2 hours to complete all the stages of the optimization. Most of the time is spent +on fine-tuning and QAT. + +### Pre-requisites + +Install Model Optimizer with optional torch and huggingface dependencies: + +```bash +pip install "nvidia-modelopt[torch,hf]" --extra-index-url https://pypi.nvidia.com +``` + +### Running the example + +To run the example, execute the following scripts in order: + +1. First we prune the Bert large model to 50% FLOPs with GradNAS algorithm. Then, we fine-tune the pruned + model with distillation from unpruned teacher model to recover 99+% of the initial F1 score (93.15). + We recommend using multiple GPUs for fine-tuning. Note that we use more epochs + for fine-tuning, which is different from the 2 epochs used originally in fine-tuning Bert without distillation since + distillation requires more epochs to converge but achieves much better results. + + ```bash + bash scripts/1_prune.sh + ``` + +1. Quantize the fine-tuned model to INT8 precision and run calibration (PTQ). + Note that PTQ will result in a slight drop in F1 score but we will be able to recover the F1 score with QAT. + We run QAT with distillation as well from unpruned teacher model. + + ```bash + bash scripts/2_int8_quantize.sh + ``` + +1. Export the quantized model to ONNX format for deployment with TensorRT. + + ```bash + bash scripts/3_onnx_export.sh + ``` diff --git a/chained_optimizations/bert_prune_distill_quantize.py b/examples/chained_optimizations/bert_prune_distill_quantize.py similarity index 97% rename from chained_optimizations/bert_prune_distill_quantize.py rename to examples/chained_optimizations/bert_prune_distill_quantize.py index 1912f1049a..5a134c6d00 100644 --- a/chained_optimizations/bert_prune_distill_quantize.py +++ b/examples/chained_optimizations/bert_prune_distill_quantize.py @@ -16,25 +16,20 @@ # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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. + """ Example showcasing how to do end-to-end optimization of a BERT model on SQuAD using Model Optimizer. This includes GradNAS pruning, INT8 quantization, fine-tuning / QAT with distillation, and ONNX export. @@ -47,7 +42,7 @@ import logging import math import os import random -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Optional import datasets import evaluate @@ -86,7 +81,7 @@ logger = get_logger(__name__) SEED = 123 -def parse_args(input_args: Optional[List[str]] = None): +def parse_args(input_args: Optional[list[str]] = None): parser = argparse.ArgumentParser( description="Finetune a transformers model on a Question Answering task" ) @@ -585,7 +580,7 @@ def evaluate_model( def postprocess_qa_predictions( examples, features, - predictions: Tuple[np.ndarray, np.ndarray], + predictions: tuple[np.ndarray, np.ndarray], version_2_with_negative: bool = False, n_best_size: int = 20, max_answer_length: int = 30, @@ -885,9 +880,9 @@ def train_and_evaluate_model( args, model: nn.Module, accelerator: Accelerator, - examples: Dict, - dataset: Dict, - dataloader: Dict[str, DataLoader], + examples: dict, + dataset: dict, + dataloader: dict[str, DataLoader], answer_column_name, ): # Optimizer @@ -1095,7 +1090,7 @@ def train_and_evaluate_model( json.dump(eval_metric, f, indent=4) -def main(input_args: Optional[List[str]] = None) -> None: +def main(input_args: Optional[list[str]] = None) -> None: args = parse_args(input_args) # Initialize the accelerator diff --git a/chained_optimizations/scripts/1_prune.sh b/examples/chained_optimizations/scripts/1_prune.sh similarity index 60% rename from chained_optimizations/scripts/1_prune.sh rename to examples/chained_optimizations/scripts/1_prune.sh index bf6a7e664f..e7e6f79740 100755 --- a/chained_optimizations/scripts/1_prune.sh +++ b/examples/chained_optimizations/scripts/1_prune.sh @@ -1,4 +1,19 @@ #!/bin/bash +# 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. + set -ex MODEL_NAME_OR_PATH=bert-large-uncased-whole-word-masking-finetuned-squad diff --git a/chained_optimizations/scripts/2_int8_quantize.sh b/examples/chained_optimizations/scripts/2_int8_quantize.sh similarity index 58% rename from chained_optimizations/scripts/2_int8_quantize.sh rename to examples/chained_optimizations/scripts/2_int8_quantize.sh index 8a629dafae..19c9909d81 100755 --- a/chained_optimizations/scripts/2_int8_quantize.sh +++ b/examples/chained_optimizations/scripts/2_int8_quantize.sh @@ -1,4 +1,19 @@ #!/bin/bash +# 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. + set -ex FLOPS_PERCENT=50 diff --git a/examples/chained_optimizations/scripts/3_onnx_export.sh b/examples/chained_optimizations/scripts/3_onnx_export.sh new file mode 100755 index 0000000000..6e30a9a207 --- /dev/null +++ b/examples/chained_optimizations/scripts/3_onnx_export.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# 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. + +set -ex + +FLOPS_PERCENT=50 +BASE_DIR=results/bert_large_pruned_${FLOPS_PERCENT}_percent + +# Export to ONNX on a single GPU +python3 bert_prune_distill_quantize.py \ + --model_name_or_path ${BASE_DIR}/int8_qat/ \ + --onnx_export_path ${BASE_DIR}/pruned_model_int8.onnx \ diff --git a/diffusers/cache_diffusion/README.md b/examples/diffusers/cache_diffusion/README.md similarity index 100% rename from diffusers/cache_diffusion/README.md rename to examples/diffusers/cache_diffusion/README.md diff --git a/diffusers/cache_diffusion/assets/SDXL_Cache_Diffusion_Img.png b/examples/diffusers/cache_diffusion/assets/SDXL_Cache_Diffusion_Img.png similarity index 100% rename from diffusers/cache_diffusion/assets/SDXL_Cache_Diffusion_Img.png rename to examples/diffusers/cache_diffusion/assets/SDXL_Cache_Diffusion_Img.png diff --git a/diffusers/cache_diffusion/assets/sdxl_cache.png b/examples/diffusers/cache_diffusion/assets/sdxl_cache.png similarity index 100% rename from diffusers/cache_diffusion/assets/sdxl_cache.png rename to examples/diffusers/cache_diffusion/assets/sdxl_cache.png diff --git a/diffusers/cache_diffusion/benchmarks.py b/examples/diffusers/cache_diffusion/benchmarks.py similarity index 62% rename from diffusers/cache_diffusion/benchmarks.py rename to examples/diffusers/cache_diffusion/benchmarks.py index 5dfd2a89ba..74f9714693 100644 --- a/diffusers/cache_diffusion/benchmarks.py +++ b/examples/diffusers/cache_diffusion/benchmarks.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 time @@ -78,13 +72,14 @@ def main(args): total_time = 0 cachify.disable(pipe) for _ in range(args.num_iter): - start_time = time.time() - _ = pipe( - prompt=[prompt] * args.batch_size, - num_inference_steps=args.num_inference_steps, - generator=generator, - ) - end_time = time.time() + with torch.autocast("cuda"): + start_time = time.time() + _ = pipe( + prompt=[prompt] * args.batch_size, + num_inference_steps=args.num_inference_steps, + generator=generator, + ) + end_time = time.time() total_time += end_time - start_time total_time = total_time / args.num_iter latency = total_time / args.batch_size @@ -94,13 +89,14 @@ def main(args): total_time = 0 cachify.enable(pipe) for _ in range(args.num_iter): - start_time = time.time() - _ = pipe( - prompt=[prompt] * args.batch_size, - num_inference_steps=args.num_inference_steps, - generator=generator, - ) - end_time = time.time() + with torch.autocast("cuda"): + start_time = time.time() + _ = pipe( + prompt=[prompt] * args.batch_size, + num_inference_steps=args.num_inference_steps, + generator=generator, + ) + end_time = time.time() cachify.reset_status(pipe) total_time += end_time - start_time total_time = total_time / args.num_iter diff --git a/diffusers/cache_diffusion/cache_diffusion/cachify.py b/examples/diffusers/cache_diffusion/cache_diffusion/cachify.py similarity index 79% rename from diffusers/cache_diffusion/cache_diffusion/cachify.py rename to examples/diffusers/cache_diffusion/cache_diffusion/cachify.py index df5b3efbb2..53be91891f 100644 --- a/diffusers/cache_diffusion/cache_diffusion/cachify.py +++ b/examples/diffusers/cache_diffusion/cache_diffusion/cachify.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 fnmatch from contextlib import contextmanager diff --git a/examples/diffusers/cache_diffusion/cache_diffusion/module.py b/examples/diffusers/cache_diffusion/cache_diffusion/module.py new file mode 100644 index 0000000000..a0e17813e1 --- /dev/null +++ b/examples/diffusers/cache_diffusion/cache_diffusion/module.py @@ -0,0 +1,49 @@ +# 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. + +from torch import nn + + +class CachedModule(nn.Module): + def __init__(self, block, select_cache_step_func) -> None: + super().__init__() + self.block = block + self.select_cache_step_func = select_cache_step_func + self.cur_step = 0 + self.cached_results = None + self.enabled = True + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.block, name) + + def if_cache(self): + return self.select_cache_step_func(self.cur_step) and self.enabled + + def enable_cache(self): + self.enabled = True + + def disable_cache(self): + self.enabled = False + self.cur_step = 0 + + def forward(self, *args, **kwargs): + if not self.if_cache(): + self.cached_results = self.block(*args, **kwargs) + if self.enabled: + self.cur_step += 1 + return self.cached_results diff --git a/diffusers/cache_diffusion/cache_diffusion/utils.py b/examples/diffusers/cache_diffusion/cache_diffusion/utils.py similarity index 51% rename from diffusers/cache_diffusion/cache_diffusion/utils.py rename to examples/diffusers/cache_diffusion/cache_diffusion/utils.py index d8b2c27726..e49c7e5fb6 100644 --- a/diffusers/cache_diffusion/cache_diffusion/utils.py +++ b/examples/diffusers/cache_diffusion/cache_diffusion/utils.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 re diff --git a/diffusers/cache_diffusion/example.ipynb b/examples/diffusers/cache_diffusion/example.ipynb similarity index 100% rename from diffusers/cache_diffusion/example.ipynb rename to examples/diffusers/cache_diffusion/example.ipynb diff --git a/diffusers/cache_diffusion/pipeline/config.py b/examples/diffusers/cache_diffusion/pipeline/config.py similarity index 83% rename from diffusers/cache_diffusion/pipeline/config.py rename to examples/diffusers/cache_diffusion/pipeline/config.py index 3be8b35261..c6ef07e173 100644 --- a/diffusers/cache_diffusion/pipeline/config.py +++ b/examples/diffusers/cache_diffusion/pipeline/config.py @@ -1,23 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel diff --git a/diffusers/cache_diffusion/pipeline/deploy.py b/examples/diffusers/cache_diffusion/pipeline/deploy.py similarity index 87% rename from diffusers/cache_diffusion/pipeline/deploy.py rename to examples/diffusers/cache_diffusion/pipeline/deploy.py index 1fa465b1a6..74112089b0 100644 --- a/diffusers/cache_diffusion/pipeline/deploy.py +++ b/examples/diffusers/cache_diffusion/pipeline/deploy.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 types from pathlib import Path diff --git a/diffusers/cache_diffusion/pipeline/models/sd3.py b/examples/diffusers/cache_diffusion/pipeline/models/sd3.py similarity index 81% rename from diffusers/cache_diffusion/pipeline/models/sd3.py rename to examples/diffusers/cache_diffusion/pipeline/models/sd3.py index 33fb5f940b..63bfcce557 100755 --- a/diffusers/cache_diffusion/pipeline/models/sd3.py +++ b/examples/diffusers/cache_diffusion/pipeline/models/sd3.py @@ -1,25 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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, List, Optional, Union +from typing import Any, Optional, Union import torch from diffusers.models.modeling_outputs import Transformer2DModelOutput @@ -37,8 +31,8 @@ def sd3_forward( encoder_hidden_states: torch.FloatTensor = None, pooled_projections: torch.FloatTensor = None, timestep: torch.LongTensor = None, - block_controlnet_hidden_states: List = None, - joint_attention_kwargs: Optional[Dict[str, Any]] = None, + block_controlnet_hidden_states: list = None, + joint_attention_kwargs: Optional[dict[str, Any]] = None, return_dict: bool = True, ) -> Union[torch.FloatTensor, Transformer2DModelOutput]: """ @@ -95,7 +89,7 @@ def sd3_forward( return custom_forward - ckpt_kwargs: Dict[str, Any] = ( + ckpt_kwargs: dict[str, Any] = ( {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} ) encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint( diff --git a/diffusers/cache_diffusion/pipeline/models/sdxl.py b/examples/diffusers/cache_diffusion/pipeline/models/sdxl.py similarity index 90% rename from diffusers/cache_diffusion/pipeline/models/sdxl.py rename to examples/diffusers/cache_diffusion/pipeline/models/sdxl.py index 92aaece6ab..f8e0c91832 100755 --- a/diffusers/cache_diffusion/pipeline/models/sdxl.py +++ b/examples/diffusers/cache_diffusion/pipeline/models/sdxl.py @@ -16,21 +16,23 @@ # 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. -# -# Not a contribution -# Changes made by NVIDIA CORPORATION & AFFILIATES or otherwise documented as -# NVIDIA-proprietary are not a contribution and subject to the following terms and conditions: -# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: LicenseRef-NvidiaProprietary -# -# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual -# property and proprietary rights in and to this material, related -# documentation and any modifications thereto. Any use, reproduction, -# disclosure or distribution of this material and related documentation -# without an express license agreement from NVIDIA CORPORATION or -# its affiliates is strictly prohibited. -from typing import Any, Dict, Optional, Tuple, Union +# 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. + +from typing import Any, Optional, Union import torch from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput @@ -44,7 +46,7 @@ def cachecrossattnupblock2d_forward( res_hidden_states_2: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, + cross_attention_kwargs: Optional[dict[str, Any]] = None, upsample_size: Optional[int] = None, attention_mask: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, @@ -108,14 +110,14 @@ def cacheunet_forward( class_labels: Optional[torch.Tensor] = None, timestep_cond: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + cross_attention_kwargs: Optional[dict[str, Any]] = None, + added_cond_kwargs: Optional[dict[str, torch.Tensor]] = None, + down_block_additional_residuals: Optional[tuple[torch.Tensor]] = None, mid_block_additional_residual: Optional[torch.Tensor] = None, - down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + down_intrablock_additional_residuals: Optional[tuple[torch.Tensor]] = None, encoder_attention_mask: Optional[torch.Tensor] = None, return_dict: bool = True, -) -> Union[UNet2DConditionOutput, Tuple]: +) -> Union[UNet2DConditionOutput, tuple]: # 1. time t_emb = self.get_time_embed(sample=sample, timestep=timestep) emb = self.time_embedding(t_emb, timestep_cond) diff --git a/diffusers/cache_diffusion/pipeline/utils.py b/examples/diffusers/cache_diffusion/pipeline/utils.py similarity index 79% rename from diffusers/cache_diffusion/pipeline/utils.py rename to examples/diffusers/cache_diffusion/pipeline/utils.py index 340f2ea58f..a50de0d732 100644 --- a/diffusers/cache_diffusion/pipeline/utils.py +++ b/examples/diffusers/cache_diffusion/pipeline/utils.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 diff --git a/diffusers/cache_diffusion/requirements.txt b/examples/diffusers/cache_diffusion/requirements.txt similarity index 58% rename from diffusers/cache_diffusion/requirements.txt rename to examples/diffusers/cache_diffusion/requirements.txt index c4c8f0ce0d..f8d7c33c8e 100644 --- a/diffusers/cache_diffusion/requirements.txt +++ b/examples/diffusers/cache_diffusion/requirements.txt @@ -1,6 +1,5 @@ cuda-python -datasets>=2.14.5 -opencv-python==4.8.1.78 +opencv-python>=4.8.1.78 peft>=0.10.0 polygraphy==0.49.9 sentencepiece diff --git a/diffusers/quantization/.gitignore b/examples/diffusers/quantization/.gitignore similarity index 100% rename from diffusers/quantization/.gitignore rename to examples/diffusers/quantization/.gitignore diff --git a/examples/diffusers/quantization/QAT_GUIDANCE.md b/examples/diffusers/quantization/QAT_GUIDANCE.md new file mode 100644 index 0000000000..3cf5cf5f65 --- /dev/null +++ b/examples/diffusers/quantization/QAT_GUIDANCE.md @@ -0,0 +1,108 @@ +# Quantization-Aware Training (QAT) + +Quantization-Aware Training (QAT) is a powerful technique for optimizing your models, particularly when post-training quantization (PTQ) methods fail to meet the requirements for your tasks. By simulating the effects of quantization during training, QAT allows the model to learn to minimize the quantization error, ultimately delivering better accuracy. + +If PTQ methods on the market fall short of your needs, QAT should be your next step. + +## Implementing QAT with ModelOPT + +**Overview** + +ModelOPT makes QAT straightforward and flexible. While the example below uses Hugging Face Accelerate for simplicity. You can integrate QAT into your workflow using your preferred training setup. + +**Example:** + +Here’s an example workflow using Hugging Face Accelerate to manage distributed training: + +```python +# Restore the model in its quantized state using ModelOPT's API +mto.restore(transformer, args.restore_quantized_ckpt) + +# Move the model to the appropriate device and set the desired weight precision +transformer.to(accelerator.device, dtype=weight_dtype) +transformer.requires_grad_(True) + +transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer, optimizer, train_dataloader, lr_scheduler +) + +``` + +Once the model is loaded in its quantized state through ModelOPT, you can proceed with regular training. The QAT process will automatically take place during the forward pass. + +**How QAT Works in ModelOPT** + +With ModelOPT, the forward pass simulates quantization, allowing the model to adjust its weights to minimize training loss and reduce quantization errors. This enables the model to better handle the constraints of quantized hardware without significant loss of performance. + +**Combining Other Techniques in ModelOPT** + +ModelOPT enables seamless integration of additional techniques, such as distillation, into your workflow without requiring significant changes to your training code. + +**Example: Combining Distillation with QAT** + +Distillation is a powerful approach where a high-precision model (the teacher) guides the training of a quantized model (the student). ModelOPT simplifies the process of combining distillation with QAT by handling most of the complexity for you. + +For more details about distillation, please refer to this [link](../../../docs/source/getting_started/6_distillation.rst). + +Here’s an example of how to implement this: + +```diff +# Restore the model in its quantized state using ModelOPT's API +mto.restore(transformer, args.restore_quantized_ckpt) + +''' +After mtd.convert, the model structure becomes: + +model: + transformer_0 + transformer_1 + teacher_model: + transformer_0 + transformer_1 + +And the forward pass is automatically monkey-patched to: + +def forward(input): + student_output = model(input) + _ = teacher_model(input) + return student_output +''' + ++ # Configuration for knowledge distillation (KD) ++ kd_config = { ++ "teacher_model": teacher_model, ++ "criterion": distill_config["criterion"], ++ "loss_balancer": distill_config["loss_balancer"], ++ "expose_minimal_state_dict": False, ++ } ++ transformer = mtd.convert(transformer, mode=[("kd_loss", kd_config)]) + +# Move the model to the appropriate device and set the desired weight precision +transformer.to(accelerator.device, dtype=weight_dtype) +transformer.requires_grad_(True) + +# Making sure to freeze the weights from model._teacher_model +transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + transformer, optimizer, train_dataloader, lr_scheduler +) + +# Compute the distillation loss using ModelOPT's compute_kd_loss ++ ... ++ loss = transformer.compute_kd_loss(...) ++ ... + +``` + +## Flexibility to Customize QAT + +ModelOPT supports fine-tuning weights during QAT out of the box only. However, if you have specific QAT algorithms or techniques, ModelOPT’s API is designed to accommodate custom implementations. + +**Get Involved** + +If you’re interested in extending ModelOPT’s QAT capabilities, we encourage you to contribute by: + +- Raising a feature request +- Opening an issue +- Submitting a merge request with your implementation + +By leveraging ModelOPT for QAT, you can achieve superior performance for complex image and video generation tasks, tailored to your specific requirements. diff --git a/diffusers/quantization/README.md b/examples/diffusers/quantization/README.md similarity index 83% rename from diffusers/quantization/README.md rename to examples/diffusers/quantization/README.md index a0633c6c96..1998eb8118 100644 --- a/diffusers/quantization/README.md +++ b/examples/diffusers/quantization/README.md @@ -16,7 +16,7 @@ If you prefer to customize parameters in calibration or run other models, please ## Calibration with Model Optimizer -We support calibration for both INT8 and FP8 precision and for both weights and activations. +We support calibration for INT8, FP8 and FP4 precision and for both weights and activations. Note: Model calibration requires relatively more GPU computing powers and it does not need to be on the same GPUs as the deployment target GPUs. Using the command line below will execute both calibration and ONNX export. @@ -30,21 +30,25 @@ python quantize.py \ --calib-size 32 --collect-method min-mean \ --percentile 1.0 --alpha 0.8 \ --quant-level 3.0 --n-steps 20 \ - --model-dtype {Half/BFloat16} \ - --quantized-torch-ckpt-save-path ./{MODEL}_int8.pt --onnx-dir {ONNX_DIR} + --model-dtype {Half/BFloat16} --trt-high-precision-dtype {Half|BFloat16} \ + --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --onnx-dir {ONNX_DIR} ``` -### FLUX-Dev|SDXL|SD2.1|SDXL-Turbo FP8 +### FLUX-Dev|SDXL|SD2.1|SDXL-Turbo FP8/FP4 + +*In our example code, FP4 is only supported for Flux. However, you can modify our script to enable FP4 format support for your own model.* + +*The FP4 format relies on the onnx-weekly build rather than the stable ONNX version.* ```sh python quantize.py \ - --model {flux-dev|sdxl-1.0|sdxl-turbo|sd2.1|sd2.1-base} --model-dtype {Half|BFloat16} \ - --format fp8 --batch-size 2 --calib-size {128|256} --quant-level {3.0|4.0} \ - --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL}_fp8.pt --collect-method default \ + --model {flux-dev|sdxl-1.0|sdxl-turbo|sd2.1|sd2.1-base} --model-dtype {Half|BFloat16} --trt-high-precision-dtype {Half|BFloat16} \ + --format {fp8|fp4} --batch-size 2 --calib-size {128|256} --quant-level {3.0|4.0} \ + --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \ --onnx-dir {ONNX_DIR} ``` -We recommend using a device with a minimum of 48GB of combined CPU and GPU memory for exporting ONNX models. Quant-level 4.0 requires additional memory. Exporting the **FLUX-Dev** ONNX model requires a minimum of 60GB of memory. +We recommend using a device with a minimum of 48GB of combined CPU and GPU memory for exporting ONNX models. Quant-level 4.0 requires additional memory. #### Important Parameters @@ -60,6 +64,10 @@ We recommend using a device with a minimum of 48GB of combined CPU and GPU memor **Then, we can load the generated checkpoint and export the INT8/FP8 quantized model in the next step. For FP8, we only support the TRT deployment on Ada/Hopper GPUs.** +## QAT + +For guidance on using our tool for QAT, please refer to this [link](./QAT_GUIDANCE.md) + ## Build the TRT engine for the Quantized ONNX Backbone We assume you already have TensorRT environment setup. INT8 requires **TensorRT version >= 9.2.0**. If you prefer to use the FP8 TensorRT, ensure you have **TensorRT version 10.2.0 or higher**. You can download the latest version of TensorRT at [here](https://developer.nvidia.com/tensorrt/download). @@ -69,25 +77,25 @@ Then generate the INT8/FP8 Backbone Engine ```bash # For SDXL -trtexec --builderOptimizationLevel=4 --stronglyTyped --onnx=./backbone.onnx \ +trtexec --builderOptimizationLevel=4 --stronglyTyped --onnx=./model.onnx \ --minShapes=sample:2x4x128x128,timestep:1,encoder_hidden_states:2x77x2048,text_embeds:2x1280,time_ids:2x6 \ --optShapes=sample:16x4x128x128,timestep:1,encoder_hidden_states:16x77x2048,text_embeds:16x1280,time_ids:16x6 \ --maxShapes=sample:16x4x128x128,timestep:1,encoder_hidden_states:16x77x2048,text_embeds:16x1280,time_ids:16x6 \ - --saveEngine=backbone.plan + --saveEngine=model.plan # For SD3-Medium -trtexec --builderOptimizationLevel=4 --stronglyTyped --onnx=./backbone.onnx \ +trtexec --builderOptimizationLevel=4 --stronglyTyped --onnx=./model.onnx \ --minShapes=hidden_states:2x16x128x128,timestep:2,encoder_hidden_states:2x333x4096,pooled_projections:2x2048 \ --optShapes=hidden_states:16x16x128x128,timestep:16,encoder_hidden_states:16x333x4096,pooled_projections:16x2048 \ --maxShapes=hidden_states:16x16x128x128,timestep:16,encoder_hidden_states:16x333x4096,pooled_projections:16x2048 \ - --saveEngine=backbone.plan + --saveEngine=model.plan # For FLUX-Dev FP8 -trtexec --onnx=./backbone.onnx --fp8 --bf16 --stronglyTyped \ - --minShapes=hidden_states:1x4096x64,img_ids:1x4096x3,encoder_hidden_states:1x512x4096,txt_ids:1x512x3,timestep:1,pooled_projections:1x768,guidance:1 \ - --optShapes=hidden_states:1x4096x64,img_ids:1x4096x3,encoder_hidden_states:1x512x4096,txt_ids:1x512x3,timestep:1,pooled_projections:1x768,guidance:1 \ - --maxShapes=hidden_states:1x4096x64,img_ids:1x4096x3,encoder_hidden_states:1x512x4096,txt_ids:1x512x3,timestep:1,pooled_projections:1x768,guidance:1 \ - --saveEngine=backbone.plan +trtexec --onnx=./model.onnx --fp8 --bf16 --stronglyTyped \ + --minShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 \ + --optShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 \ + --maxShapes=hidden_states:1x4096x64,img_ids:4096x3,encoder_hidden_states:1x512x4096,txt_ids:512x3,timestep:1,pooled_projections:1x768,guidance:1 \ + --saveEngine=model.plan ``` **Please note that `maxShapes` represents the maximum shape of the given tensor. If you want to use a larger batch size or any other dimensions, feel free to adjust the value accordingly.** @@ -129,7 +137,7 @@ Generate a quantized torch checkpoint using the command shown below: ```bash python quantize.py \ - --model {sdxl-1.0|sdxl-turbo|sd2.1|sd2.1-base|sd3-medium|flux-dev|flux-schnell} \ + --model {sdxl-1.0|sdxl-turbo|sd2.1|sd2.1-base|sd3-medium|flux-dev} \ --format fp8 \ --batch-size {1|2} \ --calib-size 128 \ @@ -143,15 +151,18 @@ Generate images for the quantized checkpoint with the following command: ```bash python diffusion_trt.py \ - --model {sdxl-1.0|sdxl-turbo|sd2.1|sd2.1-base|sd3-medium|flux-dev|flux-schnell} \ + --model {sdxl-1.0|sdxl-turbo|sd2.1|sd2.1-base|sd3-medium|flux-dev} \ --prompt "A cat holding a sign that says hello world" \ [--restore-from ./{MODEL}_fp8.pt] \ [--onnx-load-path {ONNX_DIR}] \ - [--trt_engine-path {ENGINE_DIR}] + [--trt-engine-load-path {ENGINE_DIR}] \ + [--dq_only] \ + [--torch] ``` This script will save the output image as `./{MODEL}.png` and report the latency of the TensorRT backbone. -To generate the image with FP16|BF16 precision, you can run the command shown above without the `--restore-from` argument.

+To generate the image with FP16|BF16 precision, you can run the command shown above without the `--restore-from` argument.
+While loading a TensorRT engine using the --trt-engine-load-path argument, it is recommended to load only engines generated using this pipeline.

## Demo Images @@ -221,7 +232,7 @@ mto.restore(pipe.unet, your_quantized_ckpt) # Export the onnx model ``` -By following these steps, your PEFT LoRA model should be efficiently quantized using ModelOPT, ready for deployment while maximizing performance. +By following these steps, your PEFT LoRA model should be efficiently quantized using ModelOpt, ready for deployment while maximizing performance. ### Notes About Randomness diff --git a/diffusers/quantization/assets/xl_base-fp16-sakura.png b/examples/diffusers/quantization/assets/xl_base-fp16-sakura.png similarity index 100% rename from diffusers/quantization/assets/xl_base-fp16-sakura.png rename to examples/diffusers/quantization/assets/xl_base-fp16-sakura.png diff --git a/diffusers/quantization/assets/xl_base-fp16.png b/examples/diffusers/quantization/assets/xl_base-fp16.png similarity index 100% rename from diffusers/quantization/assets/xl_base-fp16.png rename to examples/diffusers/quantization/assets/xl_base-fp16.png diff --git a/diffusers/quantization/assets/xl_base-fp8-sakura.png b/examples/diffusers/quantization/assets/xl_base-fp8-sakura.png similarity index 100% rename from diffusers/quantization/assets/xl_base-fp8-sakura.png rename to examples/diffusers/quantization/assets/xl_base-fp8-sakura.png diff --git a/diffusers/quantization/assets/xl_base-int8.png b/examples/diffusers/quantization/assets/xl_base-int8.png similarity index 100% rename from diffusers/quantization/assets/xl_base-int8.png rename to examples/diffusers/quantization/assets/xl_base-int8.png diff --git a/diffusers/quantization/build_sdxl_8bit_engine.sh b/examples/diffusers/quantization/build_sdxl_8bit_engine.sh similarity index 66% rename from diffusers/quantization/build_sdxl_8bit_engine.sh rename to examples/diffusers/quantization/build_sdxl_8bit_engine.sh index 7df198ac4a..e29759a165 100755 --- a/diffusers/quantization/build_sdxl_8bit_engine.sh +++ b/examples/diffusers/quantization/build_sdxl_8bit_engine.sh @@ -1,4 +1,19 @@ #!/bin/bash +# 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. + set -e set -x diff --git a/diffusers/quantization/calib/calib_prompts.txt b/examples/diffusers/quantization/calib/calib_prompts.txt similarity index 100% rename from diffusers/quantization/calib/calib_prompts.txt rename to examples/diffusers/quantization/calib/calib_prompts.txt diff --git a/diffusers/quantization/calib/plugin_calib.py b/examples/diffusers/quantization/calib/plugin_calib.py similarity index 72% rename from diffusers/quantization/calib/plugin_calib.py rename to examples/diffusers/quantization/calib/plugin_calib.py index 90af241998..c8c5850c8d 100644 --- a/diffusers/quantization/calib/plugin_calib.py +++ b/examples/diffusers/quantization/calib/plugin_calib.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 diff --git a/diffusers/quantization/config.py b/examples/diffusers/quantization/config.py similarity index 51% rename from diffusers/quantization/config.py rename to examples/diffusers/quantization/config.py index c4bf1b4a61..d0b3c63035 100644 --- a/diffusers/quantization/config.py +++ b/examples/diffusers/quantization/config.py @@ -1,186 +1,35 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 calib.plugin_calib import PercentileCalibrator from utils import filter_func -FP8_BF16_DEFAULT_CONFIG = { - "quant_cfg": { - "*weight_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "*input_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "*output_quantizer": {"enable": False}, - "*q_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "*k_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "*v_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "*softmax_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "default": {"enable": False}, - }, - "algorithm": "max", -} +from modelopt.core.torch.quantization.config import NVFP4_FP8_MHA_CONFIG # noqa: F401 -FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG = { +FP8_DEFAULT_CONFIG = { "quant_cfg": { - "*weight_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "*input_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, + "*weight_quantizer": {"num_bits": (4, 3), "axis": None}, + "*input_quantizer": {"num_bits": (4, 3), "axis": None}, "*output_quantizer": {"enable": False}, - "*q_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "*k_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "*v_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, + "*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None}, + "*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None}, + "*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None}, "*softmax_quantizer": { "num_bits": (4, 3), "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "transformer_blocks*bmm2_output_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "BFloat16", - }, - "default": {"enable": False}, - }, - "algorithm": "max", -} - -FP8_FP16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG = { - "quant_cfg": { - "*weight_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Half", - }, - "*input_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Half", - }, - "*output_quantizer": {"enable": False}, - "*q_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Half", - }, - "*k_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Half", - }, - "*v_bmm_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Half", - }, - "*softmax_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Half", - }, - "transformer_blocks*bmm2_output_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Half", - }, - "default": {"enable": False}, - }, - "algorithm": "max", -} - -FP8_FP16_DEFAULT_CONFIG = { - "quant_cfg": { - "*weight_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, - "*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, - "*output_quantizer": {"enable": False}, - "*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, - "*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, - "*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Half"}, - "*softmax_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Half", - }, - "default": {"enable": False}, - }, - "algorithm": "max", -} - -FP8_FP32_DEFAULT_CONFIG = { - "quant_cfg": { - "*weight_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Float", - }, - "*input_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"}, - "*output_quantizer": {"enable": False}, - "*q_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"}, - "*k_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"}, - "*v_bmm_quantizer": {"num_bits": (4, 3), "axis": None, "trt_high_precision_dtype": "Float"}, - "*softmax_quantizer": { - "num_bits": (4, 3), - "axis": None, - "trt_high_precision_dtype": "Float", }, "default": {"enable": False}, }, @@ -195,7 +44,6 @@ def get_int8_config( percentile=1.0, num_inference_steps=20, collect_method="global_min", - trt_precision_flag="Half", ): quant_config = { "quant_cfg": { @@ -221,23 +69,19 @@ def get_int8_config( quant_config["quant_cfg"][w_name] = { "num_bits": 8, "axis": 0, - "trt_high_precision_dtype": trt_precision_flag, } quant_config["quant_cfg"][i_name] = { "num_bits": 8, "axis": -1, - "trt_high_precision_dtype": trt_precision_flag, } elif isinstance(module, torch.nn.Conv2d): quant_config["quant_cfg"][w_name] = { "num_bits": 8, "axis": 0, - "trt_high_precision_dtype": trt_precision_flag, } quant_config["quant_cfg"][i_name] = { "num_bits": 8, "axis": None, - "trt_high_precision_dtype": trt_precision_flag, "calibrator": ( PercentileCalibrator, (), @@ -253,9 +97,53 @@ def get_int8_config( return quant_config +def get_fp4_config(model, fp4_linear_only=False): + """fp4 for linear, optionally fp8 for conv""" + + quant_config = { + "quant_cfg": {}, + "algorithm": "max", + } + for name, module in model.named_modules(): + w_name = f"{name}*weight_quantizer" + i_name = f"{name}*input_quantizer" + + if ( + w_name in quant_config["quant_cfg"].keys() # type: ignore + or i_name in quant_config["quant_cfg"].keys() # type: ignore + ): + continue + if isinstance(module, torch.nn.Linear): + quant_config["quant_cfg"][w_name] = { # type: ignore + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + } + quant_config["quant_cfg"][i_name] = { # type: ignore + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + } + elif isinstance(module, torch.nn.Conv2d): + if fp4_linear_only: + quant_config["quant_cfg"][w_name] = {"enable": False} # type: ignore + quant_config["quant_cfg"][i_name] = {"enable": False} # type: ignore + else: + # fp8 for conv + quant_config["quant_cfg"][w_name] = {"num_bits": (4, 3), "axis": None} # type: ignore + quant_config["quant_cfg"][i_name] = {"num_bits": (4, 3), "axis": None} # type: ignore + return quant_config + + +def set_quant_config_attr(quant_config, trt_high_precision_dtype): + for _, p in quant_config["quant_cfg"].items(): + if "num_bits" in p.keys() and "trt_high_precision_dtype" not in p.keys(): + p["trt_high_precision_dtype"] = trt_high_precision_dtype + + def update_dynamic_axes(model, dynamic_axes): if model in ["flux-dev", "flux-schnell"]: - dynamic_axes["out.0"] = dynamic_axes.pop("output") + dynamic_axes["out.0"] = dynamic_axes.pop("latent") elif model in ["sdxl-1.0", "sdxl-turbo"]: dynamic_axes["added_cond_kwargs.text_embeds"] = dynamic_axes.pop("text_embeds") dynamic_axes["added_cond_kwargs.time_ids"] = dynamic_axes.pop("time_ids") @@ -266,6 +154,21 @@ def update_dynamic_axes(model, dynamic_axes): dynamic_axes["out.0"] = dynamic_axes.pop("sample") +def remove_nesting(compilation_args): + compilation_args["dynamic_shapes"]["minShapes"]["text_embeds"] = compilation_args[ + "dynamic_shapes" + ]["minShapes"].pop("added_cond_kwargs.text_embeds") + compilation_args["dynamic_shapes"]["minShapes"]["time_ids"] = compilation_args[ + "dynamic_shapes" + ]["minShapes"].pop("added_cond_kwargs.time_ids") + compilation_args["dynamic_shapes"]["optShapes"]["text_embeds"] = compilation_args[ + "dynamic_shapes" + ]["optShapes"].pop("added_cond_kwargs.text_embeds") + compilation_args["dynamic_shapes"]["optShapes"]["time_ids"] = compilation_args[ + "dynamic_shapes" + ]["optShapes"].pop("added_cond_kwargs.time_ids") + + SDXL_DYNAMIC_SHAPES = { "sample": {"min": [2, 4, 128, 128], "opt": [16, 4, 128, 128]}, "timestep": {"min": [1], "opt": [1]}, @@ -299,8 +202,8 @@ FLUX_DEV_DYNAMIC_SHAPES = { "guidance": {"min": [1], "opt": [1]}, "pooled_projections": {"min": [1, 768], "opt": [1, 768]}, "encoder_hidden_states": {"min": [1, 512, 4096], "opt": [1, 512, 4096]}, - "txt_ids": {"min": [1, 512, 3], "opt": [1, 512, 3]}, - "img_ids": {"min": [1, 4096, 3], "opt": [1, 4096, 3]}, + "txt_ids": {"min": [512, 3], "opt": [512, 3]}, + "img_ids": {"min": [4096, 3], "opt": [4096, 3]}, } FLUX_SCHNELL_DYNAMIC_SHAPES = FLUX_DEV_DYNAMIC_SHAPES.copy() @@ -322,21 +225,6 @@ def create_dynamic_shapes(dynamic_shapes): } -def remove_nesting(compilation_args): - compilation_args["dynamic_shapes"]["minShapes"]["text_embeds"] = compilation_args[ - "dynamic_shapes" - ]["minShapes"].pop("added_cond_kwargs.text_embeds") - compilation_args["dynamic_shapes"]["minShapes"]["time_ids"] = compilation_args[ - "dynamic_shapes" - ]["minShapes"].pop("added_cond_kwargs.time_ids") - compilation_args["dynamic_shapes"]["optShapes"]["text_embeds"] = compilation_args[ - "dynamic_shapes" - ]["optShapes"].pop("added_cond_kwargs.text_embeds") - compilation_args["dynamic_shapes"]["optShapes"]["time_ids"] = compilation_args[ - "dynamic_shapes" - ]["optShapes"].pop("added_cond_kwargs.time_ids") - - DYNAMIC_SHAPES = { "sdxl-1.0": create_dynamic_shapes(SDXL_DYNAMIC_SHAPES), "sdxl-turbo": create_dynamic_shapes(SDXL_DYNAMIC_SHAPES), diff --git a/diffusers/quantization/diffusion_trt.py b/examples/diffusers/quantization/diffusion_trt.py similarity index 81% rename from diffusers/quantization/diffusion_trt.py rename to examples/diffusers/quantization/diffusion_trt.py index 8b07e58340..8591c2ad37 100644 --- a/diffusers/quantization/diffusion_trt.py +++ b/examples/diffusers/quantization/diffusion_trt.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 @@ -45,11 +39,11 @@ def generate_image(pipe, prompt, image_name): image = pipe( prompt, output_type="pil", - num_inference_steps=20, + num_inference_steps=30, generator=torch.Generator("cuda").manual_seed(seed), ).images[0] - image.save(f"{image_name}.png") - print(f"Image generated using {image_name} model saved as {image_name}.png") + image.save(image_name) + print(f"Image generated saved as {image_name}") def main(): @@ -89,8 +83,11 @@ def main(): parser.add_argument( "--torch", action="store_true", help="Generate an image using the torch pipeline" ) + parser.add_argument("--save-image-as", type=str, default=None, help="Name of the image to save") args = parser.parse_args() + image_name = args.save_image_as if args.save_image_as else f"{args.model}.png" + if args.model in ["sd2.1", "sd2.1-base"]: pipe = StableDiffusionPipeline.from_pretrained( MODEL_ID[args.model], torch_dtype=torch.float16, safety_checker=None @@ -112,11 +109,6 @@ def main(): use_safetensors=True, ) - if args.torch: - pipe.to("cuda") - generate_image(pipe, args.prompt, args.model) - return - # Save the backbone of the pipeline and move it to the GPU add_embedding = None backbone = None @@ -133,10 +125,19 @@ def main(): if args.restore_from: mto.restore(backbone, args.restore_from) + if args.torch: + if hasattr(pipe, "transformer"): + pipe.transformer = backbone + elif hasattr(pipe, "unet"): + pipe.unet = backbone + pipe.to("cuda") + generate_image(pipe, args.prompt, image_name) + return + backbone.to("cuda") # Generate dummy inputs for the backbone - dummy_inputs = generate_dummy_inputs(args.model, "cuda", True) + dummy_inputs = generate_dummy_inputs(args.model, "cuda", "BFloat16") # Define dynamic axes for ONNX export dynamic_axes = AXES_NAME[args.model] @@ -149,8 +150,10 @@ def main(): # We only need to remove the nesting for SDXL models as they contain the nested input added_cond_kwargs # which are renamed by the DeviceModel + ignore_nesting = False if args.onnx_load_path != "" and args.model in ["sdxl-1.0", "sdxl-turbo"]: remove_nesting(compilation_args) + ignore_nesting = True # Define deployment configuration deployment = { @@ -193,6 +196,7 @@ def main(): metadata, compilation_args, get_io_shapes(args.model, args.onnx_load_path), + ignore_nesting, ) if hasattr(pipe, "unet") and add_embedding: @@ -208,8 +212,8 @@ def main(): raise ValueError("Pipeline does not have a transformer or unet backbone") pipe.to("cuda") - generate_image(pipe, args.prompt, args.model) - print(f"Image generated using {args.model} model saved as {args.model}.png") + generate_image(pipe, args.prompt, image_name) + print(f"Image generated using {args.model} model saved as {image_name}") print(f"Inference latency of the backbone of the pipeline is {device_model.get_latency()} ms") diff --git a/diffusers/quantization/onnx_utils/export.py b/examples/diffusers/quantization/onnx_utils/export.py similarity index 71% rename from diffusers/quantization/onnx_utils/export.py rename to examples/diffusers/quantization/onnx_utils/export.py index fa031abfa9..3861011bbe 100644 --- a/diffusers/quantization/onnx_utils/export.py +++ b/examples/diffusers/quantization/onnx_utils/export.py @@ -15,28 +15,23 @@ # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. - +# 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 +import shutil +import tempfile from pathlib import Path import onnx @@ -45,8 +40,7 @@ import torch from onnxmltools.utils.float16_converter import convert_float_to_float16 from torch.onnx import export as onnx_export -from modelopt.torch._deploy.utils.onnx_utils import check_model_uses_external_data -from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer +from modelopt.onnx.quantization.qdq_utils import fp4qdq_to_2dq from .fp8_onnx_graphsurgeon import cast_resize_io, convert_fp16_io, convert_zp_fp8 @@ -95,14 +89,13 @@ AXES_NAME = { "sample": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}, }, "flux-dev": { - "hidden_states": {0: "batch_size", 1: "sequence"}, + "hidden_states": {0: "batch_size", 1: "latent_dim"}, "encoder_hidden_states": {0: "batch_size"}, "pooled_projections": {0: "batch_size"}, "timestep": {0: "batch_size"}, - "img_ids": {0: "batch_size", 1: "sequence"}, - "txt_ids": {0: "batch_size"}, + "img_ids": {0: "latent_dim"}, "guidance": {0: "batch_size"}, - "output": {0: "batch_size"}, + "latent": {0: "batch_size"}, }, } @@ -123,8 +116,7 @@ def flux_convert_rope_weight_type(onnx_graph): graph = gs.import_onnx(onnx_graph) for node in graph.nodes: if node.op == "Einsum": - node.inputs[1].dtype == "float32" - print(node.name) + node.inputs[1].dtype = "float32" return gs.export_onnx(graph) @@ -140,7 +132,7 @@ def generate_fp8_scales(backbone): module.weight_quantizer._amax = module.weight_quantizer._amax * (127 / 448.0) -def generate_dummy_inputs(model_id, device, model_dtype="bf16"): +def generate_dummy_inputs(model_id, device, model_dtype="BFloat16"): dummy_input = {} if model_id == "sdxl-1.0" or model_id == "sdxl-turbo": dummy_input["sample"] = torch.ones(2, 4, 128, 128).to(device).half() @@ -167,18 +159,17 @@ def generate_dummy_inputs(model_id, device, model_dtype="bf16"): dummy_input["encoder_hidden_states"] = torch.ones(2, 77, 1024).to(device).half() dummy_input["return_dict"] = False elif model_id == "flux-dev": + text_maxlen = 512 torch_dtype = torch.bfloat16 if model_dtype == "BFloat16" else torch.float16 dummy_input["hidden_states"] = torch.randn(1, 1024, 64, dtype=torch_dtype, device=device) dummy_input["encoder_hidden_states"] = torch.randn( - 1, 512, 4096, dtype=torch_dtype, device=device + 1, text_maxlen, 4096, dtype=torch_dtype, device=device ) dummy_input["pooled_projections"] = torch.randn(1, 768, dtype=torch_dtype, device=device) - dummy_input["timestep"] = torch.randn(1, dtype=torch_dtype, device=device) - dummy_input["img_ids"] = torch.randn(1, 1024, 3, dtype=torch.float32, device=device) - dummy_input["txt_ids"] = torch.randn(1, 512, 3, dtype=torch.float32, device=device) - dummy_input["guidance"] = ( - torch.randn(1, dtype=torch.float32, device=device) if model_id == "flux-dev" else None - ) + dummy_input["timestep"] = torch.tensor(data=[1.0] * 1, dtype=torch_dtype, device=device) + dummy_input["img_ids"] = torch.randn(1024, 3, dtype=torch.float32, device=device) + dummy_input["txt_ids"] = torch.randn(text_maxlen, 3, dtype=torch.float32, device=device) + dummy_input["guidance"] = torch.full((1,), 3.5, dtype=torch.float32, device=device) dummy_input["return_dict"] = False else: raise NotImplementedError(f"Unsupported model_id: {model_id}") @@ -198,36 +189,27 @@ def save_onnx(onnx_model, output): print(f"ONNX model saved to {output}") -def set_onnx_export_attr(model, dtype): - for name, module in model.named_modules(): - if isinstance(module, (TensorQuantizer, SequentialQuantizer)): - if any( - q_tag in name - for q_tag in [ - "q_bmm_quantizer", - "k_bmm_quantizer", - "v_bmm_quantizer", - "softmax_quantizer", - "bmm2_output_quantizer", - ] - ): - module._trt_high_precision_dtype = dtype - elif isinstance(module, torch.nn.Linear): - module.input_quantizer._trt_high_precision_dtype = dtype +def set_onnx_export_attr(model): + for _, module in model.named_modules(): + if isinstance(module, torch.nn.Linear): module.input_quantizer._onnx_quantizer_type = "dynamic" module.weight_quantizer._onnx_quantizer_type = "static" - else: - assert not hasattr(module, "input_quantizer") - assert not hasattr(module, "weight_quantizer") def modelopt_export_sd(backbone, onnx_dir, model_name, precision, model_dtype="BFloat16"): + model_file_name = "model.onnx" os.makedirs(f"{onnx_dir}", exist_ok=True) + tmp_subfolder = tempfile.mkdtemp(prefix="myapp_") + tmp_output = Path(f"{tmp_subfolder}/{model_file_name}") + q_output = Path(f"{onnx_dir}/{model_file_name}") + + if precision == "fp4": + set_onnx_export_attr(backbone) + dummy_inputs = generate_dummy_inputs( model_name, device=backbone.device, model_dtype=model_dtype ) - output = Path(f"{onnx_dir}/backbone.onnx") if model_name == "sdxl-1.0" or model_name == "sdxl-turbo": input_names = ["sample", "timestep", "encoder_hidden_states", "text_embeds", "time_ids"] output_names = ["latent"] @@ -244,7 +226,7 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision, model_dtype="B "txt_ids", "guidance", ] - output_names = ["output"] + output_names = ["latent"] elif model_name == "sd2.1" or model_name == "sd2.1-base": input_names = ["sample", "timestep", "encoder_hidden_states"] output_names = ["latent"] @@ -255,39 +237,20 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision, model_dtype="B do_constant_folding = True opset_version = 20 - # Copied from Huggingface's Optimum onnx_export( backbone, (dummy_inputs,), - f=output.as_posix(), + f=tmp_output.as_posix(), input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, do_constant_folding=do_constant_folding, opset_version=opset_version, ) - - onnx_model = onnx.load(str(output), load_external_data=False) - model_uses_external_data = check_model_uses_external_data(onnx_model) - - if model_uses_external_data: - onnx_model = onnx.load(str(output), load_external_data=True) - # Iterate over all files in the folder and delete them - output_folder = output.parent - for file in output_folder.iterdir(): - if file.is_file(): - file.unlink() - save_onnx(onnx_model, output) - + print(f"Saved at {tmp_output}") + onnx_model = onnx.load(str(tmp_output), load_external_data=True) if precision == "fp8": - onnx_model = onnx.load(str(output), load_external_data=True) - # Iterate over all files in the folder and delete them - output_folder = output.parent - for file in output_folder.iterdir(): - if file.is_file(): - file.unlink() - - if model_name != "flux-dev": + if not model_name.startswith("flux"): graph = gs.import_onnx(onnx_model) graph.cleanup().toposort() convert_fp16_io(graph) @@ -302,4 +265,7 @@ def modelopt_export_sd(backbone, onnx_dir, model_name, precision, model_dtype="B else: flux_convert_rope_weight_type(onnx_model) - save_onnx(onnx_model, output) + if precision == "fp4": + onnx_model = fp4qdq_to_2dq(onnx_model) + save_onnx(onnx_model, q_output) + shutil.rmtree(tmp_subfolder, ignore_errors=True) diff --git a/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py b/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py similarity index 87% rename from diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py rename to examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py index df83df0c05..3cc2824739 100644 --- a/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py +++ b/examples/diffusers/quantization/onnx_utils/fp8_onnx_graphsurgeon.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 re diff --git a/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py similarity index 76% rename from diffusers/quantization/quantize.py rename to examples/diffusers/quantization/quantize.py index 01517dadca..43fc554f09 100644 --- a/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -1,33 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 torch -from config import ( - FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG, - FP8_FP16_DEFAULT_CONFIG, - FP8_FP32_DEFAULT_CONFIG, - get_int8_config, -) +from config import FP8_DEFAULT_CONFIG, NVFP4_FP8_MHA_CONFIG, get_int8_config, set_quant_config_attr from diffusers import ( DiffusionPipeline, FluxPipeline, @@ -103,6 +92,7 @@ def main(): type=str, default=None, ) + parser.add_argument( "--n-steps", type=int, @@ -110,9 +100,17 @@ def main(): help="Number of denoising steps, for SDXL-turbo, use 1-4 steps", ) parser.add_argument("--model-dtype", type=str, default="Half", choices=["Half", "BFloat16"]) + parser.add_argument( + "--trt-high-precision-dtype", + type=str, + default="Half", + choices=["Half", "BFloat16", "Float"], + ) # Calibration and quantization parameters - parser.add_argument("--format", type=str, default="int8", choices=["int8", "fp8"]) + parser.add_argument( + "--format", type=str, default="int8", choices=["fp16", "int8", "fp8", "fp4"] + ) parser.add_argument("--percentile", type=float, default=1.0, required=False) parser.add_argument( "--collect-method", @@ -130,7 +128,7 @@ def main(): parser.add_argument("--alpha", type=float, default=1.0, help="SmoothQuant Alpha") parser.add_argument( "--quant-level", - default=3.0, + default=4.0, type=float, choices=[1.0, 2.0, 2.5, 3.0, 4.0], help="Quantization level, 1: CNN, 2: CNN+FFN, 2.5: CNN+FFN+QKV, 3: CNN+FC, 4: CNN+FC+fMHA", @@ -140,28 +138,28 @@ def main(): ) args = parser.parse_args() - if args.model not in ["flux-dev"]: assert args.model_dtype != "BFloat16", "Only Flux-dev can be BF16 precision" args.calib_size = args.calib_size // args.batch_size + torch_dtype = torch.float16 if args.model_dtype != "BFloat16" else torch.bfloat16 if args.model == "sd2.1" or args.model == "sd2.1-base": pipe = StableDiffusionPipeline.from_pretrained( - MODEL_ID[args.model], torch_dtype=torch.float16, safety_checker=None + MODEL_ID[args.model], torch_dtype=torch_dtype, safety_checker=None ) elif args.model == "sd3-medium": pipe = StableDiffusion3Pipeline.from_pretrained( - MODEL_ID[args.model], torch_dtype=torch.float16 + MODEL_ID[args.model], torch_dtype=torch_dtype ) elif args.model in ["flux-dev"]: pipe = FluxPipeline.from_pretrained( MODEL_ID[args.model], - torch_dtype=torch.bfloat16 if args.model_dtype == "BFloat16" else torch.float16, + torch_dtype=torch_dtype, ) else: pipe = DiffusionPipeline.from_pretrained( MODEL_ID[args.model], - torch_dtype=torch.float16, + torch_dtype=torch.float16, # Hardcoded into FP16 variant="fp16", use_safetensors=True, ) @@ -171,7 +169,7 @@ def main(): if args.quant_level == 4.0: assert args.format != "int8", "We only support fp8 for Level 4 Quantization" - if not args.restore_from: + if not args.restore_from and args.format != "fp16": # This is a list of prompts cali_prompts = load_calib_prompts( args.batch_size, @@ -192,20 +190,20 @@ def main(): args.percentile, args.n_steps + extra_step, collect_method=args.collect_method, - trt_precision_flag=args.model_dtype, ) elif args.format == "fp8": if args.collect_method == "default": - if args.model == "flux-dev": - quant_config = FP8_BF16_FLUX_MMDIT_BMM2_FP8_OUTPUT_CONFIG - elif args.model == "sd2.1": - quant_config = FP8_FP32_DEFAULT_CONFIG - else: - quant_config = FP8_FP16_DEFAULT_CONFIG + quant_config = FP8_DEFAULT_CONFIG else: raise NotImplementedError + elif args.format == "fp4": + assert args.model == "flux-dev", "In this example, only FP4 is supported for Flux." + quant_config = mtq.NVFP4_DEFAULT_CFG # type: ignore[attr-defined] + if args.quant_level == 4: + quant_config = NVFP4_FP8_MHA_CONFIG else: raise NotImplementedError + set_quant_config_attr(quant_config, args.trt_high_precision_dtype) def forward_loop(backbone): if args.model not in ["sd3-medium", "flux-dev"]: @@ -223,12 +221,13 @@ def main(): # All the LoRA layers should be fused check_lora(backbone) mtq.quantize(backbone, quant_config, forward_loop) - quantize_lvl(args.model, backbone, args.quant_level) - mtq.disable_quantizer(backbone, filter_func) mto.save(backbone, f"{args.quantized_torch_ckpt_save_path}") - else: + elif args.restore_from and args.format != "fp16": mto.restore(backbone, args.restore_from) + quantize_lvl(args.model, backbone, args.quant_level) + mtq.disable_quantizer(backbone, filter_func) + # if you want to export the model on CPU, move the dummy input and the model to cpu and float32 if args.onnx_dir is not None: if args.format == "fp8" and args.model != "flux-dev": diff --git a/diffusers/quantization/requirements.txt b/examples/diffusers/quantization/requirements.txt similarity index 67% rename from diffusers/quantization/requirements.txt rename to examples/diffusers/quantization/requirements.txt index b484297430..3100717ae2 100644 --- a/diffusers/quantization/requirements.txt +++ b/examples/diffusers/quantization/requirements.txt @@ -1,7 +1,8 @@ cuda-python -datasets>=2.14.5 diffusers>=0.30.1 nvtx +onnx-weekly onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com -opencv-python==4.8.1.78 +opencv-python>=4.8.1.78 peft>=0.10.0 +sentencepiece diff --git a/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py similarity index 83% rename from diffusers/quantization/utils.py rename to examples/diffusers/quantization/utils.py index 6283aa4ede..f8ff9ceb85 100644 --- a/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 import re @@ -40,7 +34,7 @@ except ModuleNotFoundError: def filter_func(name): pattern = re.compile( - r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|pos_embed|time_text_embed|context_embedder|norm_out).*" + r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|pos_embed|time_text_embed|context_embedder|norm_out|x_embedder).*" ) return pattern.match(name) is not None diff --git a/llm_distill/README.md b/examples/llm_distill/README.md similarity index 100% rename from llm_distill/README.md rename to examples/llm_distill/README.md diff --git a/llm_distill/main.py b/examples/llm_distill/main.py similarity index 81% rename from llm_distill/main.py rename to examples/llm_distill/main.py index 58be8121e3..58122dd8e9 100644 --- a/llm_distill/main.py +++ b/examples/llm_distill/main.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 logging import os @@ -29,7 +23,6 @@ import torch import torch.distributed import transformers from accelerate import PartialState -from accelerate.accelerator import extract_model_from_parallel from accelerate.logging import get_logger from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel, StateDictType from transformers import AutoTokenizer @@ -93,17 +86,17 @@ def save_model(trainer: transformers.Trainer): class KDSFTTrainer(SFTTrainer): - def compute_loss(self, model, inputs, return_outputs=False): + def compute_loss(self, model, inputs, *args, **kwargs): if not model.training: - return super().compute_loss(model, inputs, return_outputs) + _compute_loss_func = self.compute_loss_func + self.compute_loss_func = None - outputs = model(**inputs) - loss_kd = extract_model_from_parallel(model).compute_kd_loss() + loss = super().compute_loss(model, inputs, *args, **kwargs) - if self.args.past_index >= 0: - self._past = outputs[self.args.past_index] + if not model.training: + self.compute_loss_func = _compute_loss_func - return (loss_kd, outputs) if return_outputs else loss_kd + return loss def _teacher_factory(model_name_or_path): @@ -188,8 +181,11 @@ def train(): train_dataset=dset_train, eval_dataset=dset_eval, formatting_func=llama_text_format_func, - tokenizer=tokenizer, + processing_class=tokenizer, ) + if isinstance(trainer, KDSFTTrainer): + # Use our distillation aggregate loss + trainer.compute_loss_func = lambda *a, **kw: model.compute_kd_loss() # Do training if training_args.do_train: diff --git a/examples/llm_distill/requirements.txt b/examples/llm_distill/requirements.txt new file mode 100644 index 0000000000..11f4b7097e --- /dev/null +++ b/examples/llm_distill/requirements.txt @@ -0,0 +1,2 @@ +pyarrow +trl==0.13.0 diff --git a/llm_eval/.gitignore b/examples/llm_eval/.gitignore similarity index 100% rename from llm_eval/.gitignore rename to examples/llm_eval/.gitignore diff --git a/llm_eval/README.md b/examples/llm_eval/README.md similarity index 92% rename from llm_eval/README.md rename to examples/llm_eval/README.md index 444c2b786b..bf42eda1dd 100644 --- a/llm_eval/README.md +++ b/examples/llm_eval/README.md @@ -43,7 +43,7 @@ accelerate launch --multi_gpu --num_processes \ Multi-GPU evaluation without data-parallelism: ```sh -# MODELOPT_QUANT_CFG: Choose from [INT8_SMOOTHQUANT_CFG|FP8_DEFAULT_CFG|INT4_AWQ_CFG|W4A8_AWQ_BETA_CFG] +# MODELOPT_QUANT_CFG: Choose from [INT8_SMOOTHQUANT_CFG|FP8_DEFAULT_CFG|NVFP4_DEFAULT_CFG|INT4_AWQ_CFG|W4A8_AWQ_BETA_CFG] python lm_eval_hf.py --model hf \ --tasks \ --model_args pretrained=,parallelize=True \ @@ -58,11 +58,11 @@ For data-parallel evaluation, launch with `accelerate launch --multi_gpu --num_p Multi-GPU evaluation without data-parallelism: ```sh -# MODELOPT_QUANT_CFG_TO_SEARCH: Choose the formats to search separated by commas from [W4A8_AWQ_BETA_CFG,FP8_DEFAULT_CFG,NONE] +# MODELOPT_QUANT_CFG_TO_SEARCH: Choose the formats to search separated by commas from [W4A8_AWQ_BETA_CFG,FP8_DEFAULT_CFG|NVFP4_DEFAULT_CFG,NONE] # EFFECTIVE_BITS: Effective bits constraint for AutoQuantize # Examples settings for optimally quantized model with W4A8 & FP8 with effective bits to 4.8: -# MODELOPT_QUANT_CFG_TO_SEARCH=W4A8_AWQ_BETA_CFG,FP8_DEFAULT_CFG,NONE +# MODELOPT_QUANT_CFG_TO_SEARCH=W4A8_AWQ_BETA_CFG,FP8_DEFAULT_CFG|NVFP4_DEFAULT_CFG,NONE # EFFECTIVE_BITS=4.8 python lm_eval_hf.py --model hf \ @@ -80,10 +80,12 @@ For data-parallel evaluation, launch with `accelerate launch --multi_gpu --num_p - use `--model hf-seq2seq` instead. ```sh -# MODELOPT_QUANT_CFG: Choose from [INT8_SMOOTHQUANT_CFG|FP8_DEFAULT_CFG|INT4_AWQ_CFG|W4A8_AWQ_BETA_CFG] +# MODELOPT_QUANT_CFG: Choose from [INT8_SMOOTHQUANT_CFG|FP8_DEFAULT_CFG|NVFP4_DEFAULT_CFG|INT4_AWQ_CFG|W4A8_AWQ_BETA_CFG] python lm_eval_hf.py --model hf-seq2seq --model_args pretrained=t5-small --quant_cfg= --tasks --batch_size 4 ``` +If `trust_remote_code` needs to be true, please append the command with the `--trust_remote_code` flag. + ### TensorRT-LLM ```sh @@ -114,18 +116,18 @@ python mmlu.py --model_name causal --model_path ### Quantized (simulated) ```bash -# MODELOPT_QUANT_CFG: Choose from [INT8_SMOOTHQUANT_CFG|FP8_DEFAULT_CFG|INT4_AWQ_CFG|W4A8_AWQ_BETA_CFG] +# MODELOPT_QUANT_CFG: Choose from [INT8_SMOOTHQUANT_CFG|FP8_DEFAULT_CFG|NVFP4_DEFAULT_CFG|INT4_AWQ_CFG|W4A8_AWQ_BETA_CFG] python mmlu.py --model_name causal --model_path --quant_cfg MODELOPT_QUANT_CFG ``` ### AutoQuantize (simulated) ```bash -# MODELOPT_QUANT_CFG_TO_SEARCH: Choose the formats to search separated by commas from [W4A8_AWQ_BETA_CFG,FP8_DEFAULT_CFG,NONE] +# MODELOPT_QUANT_CFG_TO_SEARCH: Choose the formats to search separated by commas from [W4A8_AWQ_BETA_CFG,FP8_DEFAULT_CFG|NVFP4_DEFAULT_CFG,NONE] # EFFECTIVE_BITS: Effective bits constraint for AutoQuantize # Examples settings for optimally quantized model with W4A8 & FP8 with effective bits to 4.8: -# MODELOPT_QUANT_CFG_TO_SEARCH=W4A8_AWQ_BETA_CFG,FP8_DEFAULT_CFG,NONE +# MODELOPT_QUANT_CFG_TO_SEARCH=W4A8_AWQ_BETA_CFG,FP8_DEFAULT_CFG|NVFP4_DEFAULT_CFG,NONE # EFFECTIVE_BITS=4.8 python mmlu.py --model_name causal --model_path --quant_cfg $MODELOPT_QUANT_CFG_TO_SEARCH --auto_quantize_bits $EFFECTIVE_BITS --batch_size 4 @@ -156,7 +158,7 @@ python humaneval.py --model_name causal --model_path --n_sample 1 --quant_cfg MODELOPT_QUANT_CFG ``` diff --git a/examples/llm_eval/__init__.py b/examples/llm_eval/__init__.py new file mode 100644 index 0000000000..a08b2c2049 --- /dev/null +++ b/examples/llm_eval/__init__.py @@ -0,0 +1,14 @@ +# 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/llm_eval/gen_model_answer.py b/examples/llm_eval/gen_model_answer.py similarity index 93% rename from llm_eval/gen_model_answer.py rename to examples/llm_eval/gen_model_answer.py index c3c1e8157c..df78da0616 100644 --- a/llm_eval/gen_model_answer.py +++ b/examples/llm_eval/gen_model_answer.py @@ -15,26 +15,19 @@ # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. - +# 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. """Generate answers with local models. Usage: @@ -196,7 +189,7 @@ def get_model_answers( tokenizer = get_nemo_tokenizer(vocab_file) else: model_ckpt_path = model_path - tokenizer = get_tokenizer(model_ckpt_path) + tokenizer = get_tokenizer(model_ckpt_path, trust_remote_code=args.trust_remote_code) if engine_dir: # get model type last_part = os.path.basename(engine_dir) @@ -222,7 +215,7 @@ def get_model_answers( cpu_offloading=False, debug=False, ) - tokenizer = get_tokenizer(model_path) + tokenizer = get_tokenizer(model_path, trust_remote_code=args.trust_remote_code) if args.quant_cfg: quantize_model( model, @@ -488,6 +481,12 @@ if __name__ == "__main__": "regular quantization without AutoQuantize search will be applied." ), ) + parser.add_argument( + "--trust_remote_code", + help="Set trust_remote_code for Huggingface models and tokenizers", + default=False, + action="store_true", + ) args = parser.parse_args() print(args) diff --git a/llm_eval/humaneval.py b/examples/llm_eval/humaneval.py similarity index 84% rename from llm_eval/humaneval.py rename to examples/llm_eval/humaneval.py index 98f51386a3..70fcbebe01 100755 --- a/llm_eval/humaneval.py +++ b/examples/llm_eval/humaneval.py @@ -23,25 +23,20 @@ # SOFTWARE. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 argparse import Namespace from typing import Union @@ -186,7 +181,9 @@ def main( tokenizer = get_nemo_tokenizer(vocab_file) else: model_ckpt_path = kwargs["model_path"] - tokenizer = get_tokenizer(model_ckpt_path) + tokenizer = get_tokenizer( + model_ckpt_path, trust_remote_code=kwargs.get("trust_remote_code", False) + ) if kwargs.get("engine_dir", None): medusa_choices = kwargs.get("medusa_choices", None) model = LLM( diff --git a/llm_eval/lm_eval_hf.py b/examples/llm_eval/lm_eval_hf.py similarity index 78% rename from llm_eval/lm_eval_hf.py rename to examples/llm_eval/lm_eval_hf.py index 2033c958b3..3290b0b844 100755 --- a/llm_eval/lm_eval_hf.py +++ b/examples/llm_eval/lm_eval_hf.py @@ -23,26 +23,21 @@ # SOFTWARE. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. -from typing import Optional, Type +# 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 lm_eval import utils from lm_eval.__main__ import cli_evaluate, parse_eval_args, setup_parser @@ -52,7 +47,7 @@ from quantization_utils import quantize_model def create_from_arg_obj( - cls: Type[T], arg_dict: dict, additional_config: Optional[dict] = None + cls: type[T], arg_dict: dict, additional_config: Optional[dict] = None ) -> T: """Overrides the HFLM.create_from_arg_obj""" @@ -65,6 +60,7 @@ def create_from_arg_obj( additional_config = {k: v for k, v in additional_config.items() if v is not None} model_obj = cls(**arg_dict, **additional_config) + model_obj.tokenizer.padding_side = "left" if quant_cfg: if not calib_batch_size: diff --git a/llm_eval/lm_eval_tensorrt_llm.py b/examples/llm_eval/lm_eval_tensorrt_llm.py similarity index 70% rename from llm_eval/lm_eval_tensorrt_llm.py rename to examples/llm_eval/lm_eval_tensorrt_llm.py index 9da82eabc8..3a71a431d1 100644 --- a/llm_eval/lm_eval_tensorrt_llm.py +++ b/examples/llm_eval/lm_eval_tensorrt_llm.py @@ -1,28 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 gc import logging -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any, Iterable, Optional, Union import torch import torch.nn.functional as F @@ -35,7 +29,7 @@ from modelopt.deploy.llm.generate import LLM logger = logging.getLogger(__name__) -TokenSequence = Union[List[int], torch.LongTensor, torch.Tensor, BatchEncoding] +TokenSequence = Union[list[int], torch.LongTensor, torch.Tensor, BatchEncoding] @register_model("trt-llm") @@ -49,7 +43,9 @@ class TRTLLM(TemplateAPI): **kwargs, ): assert isinstance(tokenizer, str) - super().__init__(tokenizer=tokenizer, batch_size=int(batch_size), max_gen_toks=max_gen_toks) + super().__init__( + tokenizer=tokenizer, batch_size=int(batch_size), max_gen_toks=max_gen_toks, **kwargs + ) if self.tokenizer.pad_token_id is None: self.tokenizer.pad_token_id = self.tokenizer.eos_token_id @@ -79,20 +75,20 @@ class TRTLLM(TemplateAPI): def model_call( self, - messages: Iterable[List[int]], + messages: Iterable[list[int]], *, generate: bool = True, - gen_kwargs: Optional[Dict] = None, + gen_kwargs: Optional[dict] = None, **kwargs, ): # !!! Copy: shared dict for each request, need new object !!! gen_kwargs = copy.deepcopy(gen_kwargs) - assert isinstance(messages, Iterable), "Expect the messages to be Iterable[List[int]]" + assert isinstance(messages, Iterable), "Expect the messages to be Iterable[list[int]]" first_element = next(iter(messages)) - assert isinstance(first_element, List) and isinstance( + assert isinstance(first_element, list) and isinstance( first_element[0], int - ), "Expect the messages to be Iterable[List[int]]" + ), "Expect the messages to be Iterable[list[int]]" if not generate: return self.llm.generate_context_logits(prompts=messages) @@ -118,12 +114,12 @@ class TRTLLM(TemplateAPI): async def amodel_call( self, session, - messages: Iterable[List[int]], + messages: Iterable[list[int]], *, generate: bool = True, cache_keys: list = None, - ctxlens: Optional[List[int]] = None, - gen_kwargs: Optional[Dict] = None, + ctxlens: Optional[list[int]] = None, + gen_kwargs: Optional[dict] = None, **kwargs, ): raise NotImplementedError @@ -133,7 +129,7 @@ class TRTLLM(TemplateAPI): def _create_payload( self, - messages: Union[List[List[int]], List[dict], List[str], str], + messages: Union[list[list[int]], list[dict], list[str], str], *, generate: bool = True, gen_kwargs: Optional[dict] = None, @@ -144,17 +140,17 @@ class TRTLLM(TemplateAPI): raise NotImplementedError @staticmethod - def parse_generations(outputs: Union[Any, List[Any]], **kwargs) -> List[str]: + def parse_generations(outputs: Union[Any, list[Any]], **kwargs) -> list[str]: """Method used to parse the generations from the (batched) API response.""" return outputs @staticmethod def parse_logprobs( - outputs: Union[Any, List[Any]], - tokens: List[List[int]] = None, - ctxlens: List[int] = None, + outputs: Union[Any, list[Any]], + tokens: list[list[int]] = None, + ctxlens: list[int] = None, **kwargs, - ) -> List[Tuple[float, bool]]: + ) -> list[tuple[float, bool]]: """Method used to parse the logprobs from the (batched) API response. The provided tokens have two parts: The context tokens (length as ctxlens) and the continuation tokens. diff --git a/llm_eval/mmlu.py b/examples/llm_eval/mmlu.py similarity index 90% rename from llm_eval/mmlu.py rename to examples/llm_eval/mmlu.py index 5abae93734..38166dfdf4 100755 --- a/llm_eval/mmlu.py +++ b/examples/llm_eval/mmlu.py @@ -24,25 +24,19 @@ # SOFTWARE. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 import random @@ -259,7 +253,9 @@ def main( tokenizer = get_nemo_tokenizer(vocab_file) else: model_ckpt_path = kwargs["model_path"] - tokenizer = get_tokenizer(model_ckpt_path) + tokenizer = get_tokenizer( + model_ckpt_path, trust_remote_code=kwargs.get("trust_remote_code", False) + ) if kwargs.get("engine_dir", None): # get model type last_part = os.path.basename(kwargs["engine_dir"]) diff --git a/llm_eval/modeling.py b/examples/llm_eval/modeling.py similarity index 93% rename from llm_eval/modeling.py rename to examples/llm_eval/modeling.py index 16378e64be..5973519ddc 100644 --- a/llm_eval/modeling.py +++ b/examples/llm_eval/modeling.py @@ -23,25 +23,19 @@ # SOFTWARE. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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. # type: ignore # ruff: noqa @@ -52,7 +46,7 @@ import json import signal import time from pathlib import Path -from typing import Optional, Tuple +from typing import Optional import openai import rwkv @@ -220,7 +214,7 @@ class SeqToSeqModel(EvalModel): self.load() return len(self.tokenizer(text).input_ids) - def get_choice(self, text: str, **kwargs) -> Tuple[float, float]: + def get_choice(self, text: str, **kwargs) -> tuple[float, float]: self.load() inputs = self.tokenizer(text, return_tensors="pt").to(self.device) start_token = torch.tensor([[self.tokenizer.pad_token_id]], dtype=torch.long).to( @@ -292,7 +286,7 @@ class CausalModel(SeqToSeqModel): output_ids = outputs[:, length:].cpu().tolist() return self.tokenizer.batch_decode(output_ids, skip_special_tokens=True) - def get_choice(self, text: str, **kwargs) -> Tuple[float, float]: + def get_choice(self, text: str, **kwargs) -> tuple[float, float]: self.load() inputs = self.tokenizer(text, return_tensors="pt").to(self.device) with torch.no_grad(): @@ -364,7 +358,7 @@ class LlamaModel(SeqToSeqModel): batch_size, length = inputs.input_ids.shape return self.tokenizer.decode(outputs[0, length:], skip_special_tokens=True) - def get_choice(self, text: str, **kwargs) -> Tuple[float, float]: + def get_choice(self, text: str, **kwargs) -> tuple[float, float]: self.load() inputs = self.tokenizer(text, return_tensors="pt").to(self.device) with torch.no_grad(): diff --git a/llm_eval/quantization_utils.py b/examples/llm_eval/quantization_utils.py similarity index 69% rename from llm_eval/quantization_utils.py rename to examples/llm_eval/quantization_utils.py index 71ffe1f735..8e01293a22 100644 --- a/llm_eval/quantization_utils.py +++ b/examples/llm_eval/quantization_utils.py @@ -1,45 +1,42 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 sys -from pathlib import Path -from typing import List, Union +from typing import Union from transformers import AutoTokenizer import modelopt.torch.quantization as mtq from modelopt.torch.quantization.plugins import register_hf_attentions_on_the_fly - -sys.path.append(str(Path(__file__).resolve().parent / "../common")) -from dataset_utils import create_forward_loop, get_dataset_dataloader, get_max_batch_size +from modelopt.torch.utils.dataset_utils import ( + create_forward_loop, + get_dataset_dataloader, + get_max_batch_size, +) MAX_SEQ_LEN = 2048 MAX_OUTPUT_LEN = 512 -def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN): +def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN, trust_remote_code=False): """Returns the tokenizer from the model ckpt_path.""" print(f"Initializing tokenizer from {ckpt_path}") tokenizer = AutoTokenizer.from_pretrained( - ckpt_path, model_max_length=max_seq_len, padding_side="left", trust_remote_code=True + ckpt_path, + model_max_length=max_seq_len, + padding_side="left", + trust_remote_code=trust_remote_code, ) if type(tokenizer).__name__ == "QWenTokenizer": # qwen use token id 151643 as pad and eos tokens @@ -54,7 +51,7 @@ def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN): def _quantize_model_with_dataset( - lm, quant_cfg: Union[str, List[str]], calib_dataset, auto_quantize_bits=None, batch_size=1 + lm, quant_cfg: Union[str, list[str]], calib_dataset, auto_quantize_bits=None, batch_size=1 ): if hasattr(lm, "gpt2"): net = lm.gpt2 @@ -81,19 +78,9 @@ def _quantize_model_with_dataset( verbose=True, ) else: - atq_cfg = getattr(mtq, quant_cfg) # type: ignore [arg-type] + mtq_cfg = getattr(mtq, quant_cfg) # type: ignore [arg-type] - # The calibrate_loop is a custom defined method to run the model with the input data. - # The basic version looks like: - # - # def calibrate_loop(model, dataloader): - # for data in dataloader: - # model(**data) - # - # We also provided a util method to generate the forward_loop with additional error handlings. - calibrate_loop = create_forward_loop(dataloader=calib_dataset) - - def is_dynamic(atq_cfg): + def is_dynamic(mtq_cfg): def _not_dynamic(cfg): return ( cfg.get("enable", True) @@ -101,7 +88,10 @@ def _quantize_model_with_dataset( and cfg.get("*", {}).get("enable", True) ) - for _, cfg in atq_cfg.get("quant_cfg", {}).items(): + for name, cfg in mtq_cfg.get("quant_cfg", {}).items(): + if "weight_quantizer" in name: + # We don't calibrate weight quantizer + continue # quantization like W4A8 has a list of weight quantizers if isinstance(cfg, list): for config in cfg: @@ -113,18 +103,29 @@ def _quantize_model_with_dataset( return True - use_calibration = not is_dynamic(atq_cfg) + calibrate_loop = None + use_calibration = not (is_dynamic(mtq_cfg) and mtq_cfg.get("algorithm") == "max") if not use_calibration: print("Dynamic quantization. Calibration skipped.") + else: + # The calibrate_loop is a custom defined method to run the model with the input data. + # The basic version looks like: + # + # def calibrate_loop(model, dataloader): + # for data in dataloader: + # model(**data) + # + # We also provided a util method to generate the forward_loop with additional error handlings. + calibrate_loop = create_forward_loop(dataloader=calib_dataset) quantize_bmm_attention = False - for key in atq_cfg["quant_cfg"]: + for key in mtq_cfg["quant_cfg"]: if "bmm_quantizer" in key: quantize_bmm_attention = True if quantize_bmm_attention: register_hf_attentions_on_the_fly(net) - net = mtq.quantize(net, atq_cfg, calibrate_loop if use_calibration else None) + net = mtq.quantize(net, mtq_cfg, calibrate_loop) mtq.print_quant_summary(net) # Fold weights for faster evaluation. mtq.fold_weight(net) diff --git a/llm_eval/requirements.txt b/examples/llm_eval/requirements.txt similarity index 61% rename from llm_eval/requirements.txt rename to examples/llm_eval/requirements.txt index 4183bb0f83..db9e336d92 100644 --- a/llm_eval/requirements.txt +++ b/examples/llm_eval/requirements.txt @@ -1,6 +1,5 @@ -datasets>=2.14.5 fire>=0.5.0 -lm_eval[api]==0.4.5 +lm_eval[api]>=0.4.6 openai>=0.28.1 peft>=0.5.0 rwkv>=0.7.3 diff --git a/llm_eval/run_fastchat.sh b/examples/llm_eval/run_fastchat.sh similarity index 84% rename from llm_eval/run_fastchat.sh rename to examples/llm_eval/run_fastchat.sh index 4eb926c687..16aa54bc2e 100644 --- a/llm_eval/run_fastchat.sh +++ b/examples/llm_eval/run_fastchat.sh @@ -1,5 +1,20 @@ #!/bin/bash +# 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. + # This script can be used to generate responses for MT-Bench using FastChat. # Ensure that you are logged in to huggingface-cli using: huggingface-cli login. # If you are using NIM, ensure that you export the NIM API key using: diff --git a/llm_ptq/.gitignore b/examples/llm_ptq/.gitignore similarity index 100% rename from llm_ptq/.gitignore rename to examples/llm_ptq/.gitignore diff --git a/llm_ptq/README.md b/examples/llm_ptq/README.md similarity index 79% rename from llm_ptq/README.md rename to examples/llm_ptq/README.md index a8cfb351ae..928c2bc09e 100644 --- a/llm_ptq/README.md +++ b/examples/llm_ptq/README.md @@ -25,6 +25,8 @@ There are many quantization schemes supported in the example scripts: 1. The W4A8 AWQ is an extension of the INT4 AWQ quantization that it also uses FP8 for activation for more speed up and acceleration. +1. The [NVFP4](https://blogs.nvidia.com/blog/generative-ai-studio-ces-geforce-rtx-50-series/) is one of the new FP4 formats supported by NVIDIA Blackwell GPU and demonstrates good accuracy compared with other 4-bit alternatives. NVFP4 can be applied to both model weights as well as activations, providing the potential for both a significant increase in math throughput and reductions in memory footprint and memory bandwidth usage compared to the FP8 data format on Blackwell. + The following scripts provide an all-in-one and step-by-step model quantization example for Llama-3, NeMo Nemotron, and Megatron-LM models. The quantization format and the number of GPUs will be supplied as inputs to these scripts. By default, we build the engine for the fp8 format and 1 GPU. ```bash @@ -39,15 +41,19 @@ For LLM models like [Llama-3](https://huggingface.co/meta-llama): # Install model specific pip dependencies if needed export HF_PATH= -scripts/huggingface_example.sh --model $HF_PATH --quant [fp8|int8_sq|int4_awq|w4a8_awq] --tp [1|2|4|8] +scripts/huggingface_example.sh --model $HF_PATH --quant [fp8|nvfp4|int8_sq|int4_awq|w4a8_awq] --tp [1|2|4|8] ``` +> *By default `trust_remote_code` is set to false. Please turn it on if model calibration and eval requires it using `--trust_remote_code`.* + > *If the Huggingface model calibration fails on a multi-GPU system due to mismatched tensor placement, please try setting CUDA_VISIBLE_DEVICES to a smaller number.* > *FP8 calibration over a large model with limited GPU memory is not recommended but possible with the [accelerate](https://huggingface.co/docs/accelerate/en/usage_guides/big_modeling) package. Please tune the device_map setting in [`example_utils.py`](./example_utils.py) if needed for model loading and the calibration process can be slow.* > *Huggingface models trained with `modelopt.torch.speculative` can be used as regular Huggingface models in PTQ.* +> *Calibration by default uses left padding_side for the Huggingface tokenizer as it usually leads to lower accuracy loss. The exported tokenizer files restores the default padding_side.* + #### For NeMo models like [nemotron](https://huggingface.co/nvidia/nemotron-3-8b-base-4k): NeMo PTQ requires the NeMo package installed. It's recommended to start from the NeMo containers like `nvcr.io/nvidia/nemo:24.07` or latest `nvcr.io/nvidia/nemo:dev` directly. @@ -61,14 +67,14 @@ export GPT_MODEL_FILE=Nemotron-3-8B-Base-4k.nemo pip install -U "nvidia-modelopt[torch]" --extra-index-url https://pypi.nvidia.com python -c "import modelopt.torch.quantization.extensions as ext; ext.precompile()" -scripts/nemo_example.sh --type gpt --model $GPT_MODEL_FILE --quant [fp8|int8_sq|int4_awq] --tp [1|2|4|8] +scripts/nemo_example.sh --type gpt --model $GPT_MODEL_FILE --quant [fp8|nvfp4|int8_sq|int4_awq] --tp [1|2|4|8] ``` > *If the TensorRT-LLM version in the NeMo container is older than the supported version, please continue building TRT-LLM engine with the `docker.io/library/modelopt_examples:latest` container built in the ModelOpt docker build step. Additionally you would also need to `pip install megatron-core "nemo-toolkit[all]" --extra-index-url https://pypi.nvidia.com` to install required NeMo dependencies* #### For Megatron-LM models: -Megatron-LM framework PTQ and TensorRT-LLM deployment examples are maintained in the Megatron-LM GitHub repo. Please refer to the examples [here](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/inference). +Megatron-LM framework PTQ and TensorRT-LLM deployment examples are maintained in the Megatron-LM GitHub repo. Please refer to the examples [here](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/export). > *If GPU out-of-memory error is reported running the scripts, please try editing the scripts and reducing the max batch size of the TensorRT-LLM engine to save GPU memory.* @@ -78,45 +84,45 @@ Please refer to the `Technical Details` section below about the stage executed i ### Model Support List -Model | fp8 | int8_sq | int4_awq | w4a8_awq1 ---- | --- | --- | --- | --- -GPT2 | Yes | Yes | No | No -GPTJ | Yes | Yes | Yes | Yes -LLAMA 2 | Yes | Yes | Yes | Yes -LLAMA 3, 3.1 | Yes | No | Yes | No -LLAMA 2 (Nemo) | Yes | Yes | Yes | Yes -CodeLlama | Yes | Yes | Yes | No -Mistral | Yes | Yes | Yes | No -Mixtral 8x7B, 8x22B | Yes3 | No | Yes2 | No -Snowflake Arctic2 | Yes | No | Yes | No -Falcon 40B, 180B | Yes | Yes | Yes | Yes -Falcon 7B | Yes | Yes | No | No -MPT 7B, 30B | Yes | Yes | Yes | Yes -Baichuan 1, 2 | Yes | Yes | Yes | Yes -ChatGLM2, 3 6B | No | No | Yes | No -Bloom | Yes | Yes | Yes | Yes -Phi-1,2,3 | Yes | Yes | Yes | Yes4 -Phi-3.5 MOE | Yes | No | No | No -Nemotron 8B | Yes | No | Yes | No -Gemma 2B, 7B | Yes | No | Yes | Yes -Gemma 2 9B, 27B | Yes | No | Yes | No -RecurrentGemma 2B | Yes | Yes | Yes | No -StarCoder 2 | Yes | Yes | Yes | No -QWen5 | Yes | Yes | Yes | Yes -DBRX | Yes | No | No | No -InternLM2 | Yes | No | Yes | Yes4 -Exaone | Yes | Yes | Yes | Yes -Minitron | Yes | Yes | Yes | Yes2 +Model | fp8 | int8_sq | int4_awq | w4a8_awq1 | nvfp45 | +--- | --- | --- | --- | --- | --- +GPTJ | Yes | Yes | Yes | Yes | - +LLAMA 2 | Yes | Yes | Yes | Yes | - +LLAMA 3, 3.1, 3.3 | Yes | No | Yes | Yes3 | Yes +LLAMA 2 (Nemo) | Yes | Yes | Yes | Yes | - +CodeLlama | Yes | Yes | Yes | No | - +Mistral | Yes | Yes | Yes | No | Yes +Mixtral 8x7B, 8x22B | Yes | No | Yes2 | No | Yes +Snowflake Arctic2 | Yes | No | Yes | No | - +Falcon 40B, 180B | Yes | Yes | Yes | Yes | - +Falcon 7B | Yes | Yes | No | No | - +MPT 7B, 30B | Yes | Yes | Yes | Yes | - +Baichuan 1, 2 | Yes | Yes | Yes | Yes | - +ChatGLM2, 3 6B | No | No | Yes | No | - +Bloom | Yes | Yes | Yes | Yes | - +Phi-1,2,3,4 | Yes | Yes | Yes | Yes3 | +Phi-3.5 MOE | Yes | No | No | No | - +Nemotron 8B | Yes | No | Yes | No | - +Gemma 2B, 7B | Yes | No | Yes | Yes | - +Gemma 2 9B, 27B | Yes2 | No | Yes | No | - +RecurrentGemma 2B | Yes | Yes | Yes | No | - +StarCoder 2 | Yes | Yes | Yes | No | - +QWen 2, 2.5 4 | Yes | Yes | Yes | Yes | Yes +DBRX | Yes | No | No | No | - +InternLM2 | Yes | No | Yes | Yes3 | - +Exaone | Yes | Yes | Yes | Yes | - +Minitron | Yes | Yes | Yes | Yes2 | Yes +T5 | Yes | Yes | Yes | Yes | - -> *1.The w4a8_awq is an experimental quantization scheme that may result in a higher accuracy penalty. Only available on sm90 GPUs* +> *1.The w4a8_awq is an experimental quantization scheme that may result in a higher accuracy penalty.* > *2.For some models, there is only support for exporting quantized checkpoints.* -> *3.Mixtral FP8 only available on sm90 GPUs* +> *3.W4A8_AWQ is only available on some models but not all* -> *4.W4A8_AWQ is only available on some models but not all* +> *4.For some models, KV cache quantization may result in a higher accuracy penalty.* -> *5.For some models, KV cache quantization may result in a higher accuracy penalty.* +> *5.A selective set of the popular models are internally tested. The actual model support list may be longer. NVFP4 inference requires Blackwell GPUs and TensorRT-LLM v0.17 or later* > *The accuracy loss after PTQ may vary depending on the actual model and the quantization method. Different models may have different accuracy loss and usually the accuracy loss is more significant when the base model is small. If the accuracy after PTQ is not meeting the requirement, please try either modifying [hf_ptq.py](./hf_ptq.py) and disabling the KV cache quantization or using the [QAT](./../llm_qat/README.md) instead.* @@ -140,7 +146,7 @@ llm_fp8 = LLM(model="", quantization="modelopt") print(llm_fp8.generate(["What's the age of the earth? "])) ``` -> *1. Unified checkpoint export currently does not support sparsity, speculative decoding, KV cache quantization or AutoQuantize.* +> *1. Unified checkpoint export currently does not support sparsity, KV cache quantization or AutoQuantize. Speculative decoding is only supported in unified checkpoint export. The exported unified checkpoint then needs a TensorRT-LLM checkpoint converter (e.g., [this](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/eagle/convert_checkpoint.py)) to convert and build the TensorRT engine(s) for deployment.* > *2. Exported checkpoint can be deployed on vLLM with changes from this unmerged [PR](https://github.com/vllm-project/vllm/pull/6112/files#diff-b2645ce390db5e2ac2123144700f912fab9458314789aa8245c657cf42c6039e). Users can deploy to vLLM directly once the PR has been merged.* ### Model Support List @@ -157,32 +163,33 @@ CodeLlama | Yes [AutoQuantize](https://nvidia.github.io/TensorRT-Model-Optimizer/reference/generated/modelopt.torch.quantization.model_quant.html#modelopt.torch.quantization.model_quant.auto_quantize) is a PTQ algorithm from ModelOpt which quantizes a model by searching for the best quantization format per-layer while meeting the performance constraint specified by the user. This way, `AutoQuantize` enables to trade-off model accuracy for performance. -Currently `AutoQuantize` supports only `effective_bits` as the performance constraint (for both weight-only quantization and +Currently `AutoQuantize` supports only `auto_quantize_bits` as the performance constraint (for both weight-only quantization and weight & activation quantization). See [AutoQuantize documentation](https://nvidia.github.io/TensorRT-Model-Optimizer/reference/generated/modelopt.torch.quantization.model_quant.html#modelopt.torch.quantization.model_quant.auto_quantize) for more details. -#### AutoQuantize for Hugging Face models: +#### AutoQuantize for Hugging Face models `AutoQuantize` can be performed for Huggingface LLM models like [Llama-3](https://huggingface.co/meta-llama) as shown below: ```bash export HF_PATH= -# --effective_bits specifies the constraint for `AutoQuantize` -# --quant specifies the format to be searched for `AutoQuantize` -scripts/huggingface_example.sh --type llama --model $HF_PATH --quant w4a8_awq,fp8 --effective_bits 4.8 --tp [1|2|4|8] --calib_batch_size 4 +# --auto_quantize_bits specifies the constraint for `AutoQuantize` +# --quant specifies the formats to be searched for `AutoQuantize` +# NOTE: auto_quantize_bits cannot be lower than the number of bits for the smallest quantization format in --quant +scripts/huggingface_example.sh --type llama --model $HF_PATH --quant w4a8_awq,fp8 --auto_quantize_bits 4.8 --tp [1|2|4|8] --calib_batch_size 4 ``` The above example perform `AutoQuantize` where the less quantization sensitive layers are quantized with `w4a8_awq` (specified by `--quant w4a8_awq`) and the more sensitive layers -are kept un-quantized such that the effective bits is 4.8 (specified by `--effective_bits 4.8`). +are kept un-quantized such that the effective bits is 4.8 (specified by `--auto_quantize_bits 4.8`). -#### AutoQuantize for NeMo models: +#### AutoQuantize for NeMo models The usage is similar for NeMo models to perform `AutoQuantize`. Please refer to the earlier section on [NeMo models](#for-nemo-models-like-nemotron) for the full setup instructions. ```bash -# --effective_bits specifies the constraint for `AutoQuantize` +# --auto_quantize_bits specifies the constraint for `AutoQuantize` # --quant specifies the formats to be searched for `AutoQuantize`. Multiple formats can be searched over by passing them as comma separated values -scripts/nemo_example.sh --type gpt --model $GPT_MODEL_FILE --quant fp8,int4_awq --effective_bits 6.4 --tp [1|2|4|8] +scripts/nemo_example.sh --type gpt --model $GPT_MODEL_FILE --quant fp8,int4_awq --auto_quantize_bits 6.4 --tp [1|2|4|8] ``` ## Technical Details diff --git a/examples/llm_ptq/benchmarks/benchmark_suite.py b/examples/llm_ptq/benchmarks/benchmark_suite.py new file mode 100644 index 0000000000..fd767b95a5 --- /dev/null +++ b/examples/llm_ptq/benchmarks/benchmark_suite.py @@ -0,0 +1,21 @@ +# 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. + +# This import overrides the benchmark executor and disable chunked_context +import disable_bench_chunked_context # noqa +import tensorrt_llm.commands.bench as trtllm_bench + +if __name__ == "__main__": + trtllm_bench.main() diff --git a/examples/llm_ptq/benchmarks/disable_bench_chunked_context.py b/examples/llm_ptq/benchmarks/disable_bench_chunked_context.py new file mode 100644 index 0000000000..dfb61b20e4 --- /dev/null +++ b/examples/llm_ptq/benchmarks/disable_bench_chunked_context.py @@ -0,0 +1,36 @@ +# 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. + +from pathlib import Path +from typing import Union + +import tensorrt_llm.bench.benchmark.utils.general as general +from tensorrt_llm.bench.benchmark.utils.general import ( + get_settings_from_engine as get_settings_from_engine_original, +) + + +def get_settings_from_engine_override( + engine_path: Path, +) -> tuple[dict[str, Union[str, int]], dict[str, Union[str, int]]]: + exec_settings, build_cfg = get_settings_from_engine_original(engine_path) + # Do not enable chunking if use_paged_context_fmha if not enabled. + exec_settings["settings_config"]["chunking"] = build_cfg["plugin_config"][ + "use_paged_context_fmha" + ] + return exec_settings, build_cfg + + +general.get_settings_from_engine = get_settings_from_engine_override diff --git a/llm_ptq/config/megatron_quantization.yaml b/examples/llm_ptq/config/megatron_quantization.yaml similarity index 100% rename from llm_ptq/config/megatron_quantization.yaml rename to examples/llm_ptq/config/megatron_quantization.yaml diff --git a/llm_ptq/example_utils.py b/examples/llm_ptq/example_utils.py similarity index 75% rename from llm_ptq/example_utils.py rename to examples/llm_ptq/example_utils.py index 5156118f9b..7cafb0490c 100755 --- a/llm_ptq/example_utils.py +++ b/examples/llm_ptq/example_utils.py @@ -1,31 +1,25 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 import sys import torch -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer +from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer -MAX_SEQ_LEN = 2048 +from modelopt.torch.utils.image_processor import MllamaImageProcessor MODEL_NAME_PATTERN_MAP = { "GPT2": "gpt", @@ -57,6 +51,7 @@ MODEL_NAME_PATTERN_MAP = { "InternLM2ForCausalLM": "internlm", "ExaoneForCausalLM": "exaone", "Nemotron": "gpt", + "Deepseek": "deepseek", } @@ -77,7 +72,7 @@ def get_mode_type_from_engine_dir(engine_dir_str): return model_type -def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN): +def get_tokenizer(ckpt_path, trust_remote_code=False, **kwargs): print(f"Initializing tokenizer from {ckpt_path}") if "vila" in ckpt_path.lower(): @@ -92,10 +87,7 @@ def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN): else: tokenizer = AutoTokenizer.from_pretrained( - ckpt_path, - model_max_length=max_seq_len, - padding_side="left", - trust_remote_code=True, + ckpt_path, trust_remote_code=trust_remote_code, **kwargs ) if "qwen" in type(tokenizer).__name__.lower(): @@ -113,6 +105,22 @@ def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN): return tokenizer +def get_processor(ckpt_path, device=None, trust_remote_code=False): + """ + Returns a :class:`modelopt.torch.utils.image_processor.MllamaImageProcessor` object. + """ + processor = AutoProcessor.from_pretrained( + ckpt_path, + padding_side="left", + trust_remote_code=trust_remote_code, + ) + if processor.tokenizer.pad_token is None: + processor.tokenizer.pad_token = processor.tokenizer.eos_token + assert processor.tokenizer.pad_token is not None, f"Pad token for {ckpt_path} cannot be set!" + + return MllamaImageProcessor(processor, device) + + def get_dtype(dtype): if dtype == "bf16": dtype = torch.bfloat16 @@ -126,7 +134,7 @@ def get_dtype(dtype): return dtype -def get_model(ckpt_path, device="cuda", gpu_mem_percentage=0.8): +def get_model(ckpt_path, device="cuda", gpu_mem_percentage=0.8, trust_remote_code=False): print(f"Initializing model from {ckpt_path}") device_map = "auto" @@ -144,11 +152,11 @@ def get_model(ckpt_path, device="cuda", gpu_mem_percentage=0.8): hf_vila = AutoModel.from_pretrained( ckpt_path, device_map=device_map, - trust_remote_code=True, + trust_remote_code=trust_remote_code, ) model = hf_vila.llm else: - hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=True) + hf_config = AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code) if hf_config.model_type == "llava": from transformers import LlavaForConditionalGeneration @@ -158,16 +166,16 @@ def get_model(ckpt_path, device="cuda", gpu_mem_percentage=0.8): ) model = hf_llava.language_model elif hf_config.model_type == "t5": - from transformers import T5ForConditionalGeneration + from transformers import AutoModelForSeq2SeqLM - model = T5ForConditionalGeneration.from_pretrained( + model = AutoModelForSeq2SeqLM.from_pretrained( ckpt_path, device_map=device_map, **model_kwargs ) elif hf_config.model_type == "glm": from transformers import AutoModelForSeq2SeqLM model = AutoModelForSeq2SeqLM.from_pretrained( - ckpt_path, device_map="cuda", **model_kwargs, trust_remote_code=True + ckpt_path, device_map="cuda", **model_kwargs, trust_remote_code=trust_remote_code ) elif hf_config.model_type == "mllama": from transformers import MllamaForConditionalGeneration @@ -176,14 +184,9 @@ def get_model(ckpt_path, device="cuda", gpu_mem_percentage=0.8): ckpt_path, device_map=device_map, **model_kwargs, - trust_remote_code=True, + trust_remote_code=trust_remote_code, ) - # Add vision_output_dim back to the language model config for export. - vision_output_dim = model.config.vision_config.vision_output_dim - model = model.language_model - setattr(model.config, "vision_output_dim", vision_output_dim) - else: from accelerate import infer_auto_device_map, init_empty_weights from accelerate.utils import get_max_memory @@ -193,7 +196,7 @@ def get_model(ckpt_path, device="cuda", gpu_mem_percentage=0.8): # unless specified by the hf_config. torch_dtype = getattr(hf_config, "torch_dtype", torch.float16) model = AutoModelForCausalLM.from_config( - hf_config, torch_dtype=torch_dtype, trust_remote_code=True + hf_config, torch_dtype=torch_dtype, trust_remote_code=trust_remote_code ) max_memory = get_max_memory() @@ -218,7 +221,7 @@ def get_model(ckpt_path, device="cuda", gpu_mem_percentage=0.8): ckpt_path, device_map=device_map, **model_kwargs, - trust_remote_code=True, + trust_remote_code=trust_remote_code, ) model.eval() if device == "cuda": @@ -231,3 +234,8 @@ def get_model(ckpt_path, device="cuda", gpu_mem_percentage=0.8): def is_model_on_gpu(model) -> bool: """Returns if the model is fully loaded on GPUs.""" return all("cuda" in str(param.device) for param in model.parameters()) + + +def is_enc_dec(model_type) -> bool: + """Return if the model is a encoder-decoder model.""" + return model_type in ["t5", "bart"] diff --git a/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py similarity index 69% rename from llm_ptq/hf_ptq.py rename to examples/llm_ptq/hf_ptq.py index 17adb57b4b..c9936de4e7 100755 --- a/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -1,43 +1,48 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 random -import sys import time -from pathlib import Path +import warnings import numpy as np import torch from accelerate.hooks import remove_hook_from_module -from example_utils import get_model, get_model_type, get_tokenizer, is_model_on_gpu +from example_utils import ( + get_model, + get_model_type, + get_processor, + get_tokenizer, + is_enc_dec, + is_model_on_gpu, +) +from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq import modelopt.torch.sparsity as mts from modelopt.torch.export import export_hf_checkpoint, export_tensorrt_llm_checkpoint - -sys.path.append(str(Path(__file__).resolve().parent / "../common")) -from dataset_utils import create_forward_loop, get_dataset_dataloader, get_max_batch_size +from modelopt.torch.utils.dataset_utils import ( + create_forward_loop, + get_dataset_dataloader, + get_max_batch_size, +) +from modelopt.torch.utils.image_processor import MllamaImageProcessor +from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader RAND_SEED = 1234 @@ -47,6 +52,8 @@ QUANT_CFG_CHOICES = { "fp8": "FP8_DEFAULT_CFG", "int4_awq": "INT4_AWQ_CFG", "w4a8_awq": "W4A8_AWQ_BETA_CFG", + "nvfp4": "NVFP4_DEFAULT_CFG", + "nvfp4_awq": "NVFP4_AWQ_LITE_CFG", } mto.enable_huggingface_checkpointing() @@ -59,11 +66,12 @@ def auto_quantize( # Check if all provided quantization formats are supported if args.export_fmt == "hf": assert all( - qformat in ["fp8", "int4_awq"] for qformat in qformat_list + qformat in ["fp8", "int4_awq", "nvfp4", "nvfp4_awq"] for qformat in qformat_list ), "One or more quantization formats provided are not supported for unified checkpoint export" else: assert all( - qformat in ["fp8", "int8_sq", "int4_awq", "w4a8_awq"] for qformat in qformat_list + qformat in ["fp8", "int8_sq", "int4_awq", "w4a8_awq", "nvfp4", "nvfp4_awq"] + for qformat in qformat_list ), "One or more quantization formats provided are not supported for tensorrt llm export" model, _ = mtq.auto_quantize( @@ -79,10 +87,11 @@ def auto_quantize( len(calib_dataloader), 128 // batch_size ), # Limit the number of score steps to avoid long calibration time verbose=True, + disabled_layers=["*lm_head*"], ) # We need to explicitly calibrate for kv cache quantization - enable_kv_cache_quantization = "int8" not in args.qformat and "hf" not in args.export_fmt + enable_kv_cache_quantization = "int8" not in args.qformat if enable_kv_cache_quantization: mtq.set_quantizer_by_cfg( model, @@ -163,14 +172,39 @@ def main(args): assert args.qformat in [ "int4_awq", "fp8", + "nvfp4", + "nvfp4_awq", + "w4a8_awq", ], f"Quantization format {args.qformat} not supported for HF export path" - model = get_model(args.pyt_ckpt_path, args.device) - tokenizer = get_tokenizer(args.pyt_ckpt_path) + model = get_model(args.pyt_ckpt_path, args.device, trust_remote_code=args.trust_remote_code) model_type = get_model_type(model) + device = model.device if hasattr(model, "model"): device = model.model.device + processor = None + tokenizer = None + if model_type == "mllama": + if args.dataset is None: + args.dataset = "scienceqa" + warnings.warn( + "Currently only the scienceqa dataset is supported for the mllama model. " + "Overriding dataset to scienceqa." + ) + elif args.dataset != "scienceqa": + raise ValueError("Only the scienceqa dataset is supported for the mllama model.") + processor = get_processor( + args.pyt_ckpt_path, device, trust_remote_code=args.trust_remote_code + ) + else: + if args.dataset is None: + args.dataset = "cnn_dailymail" + warnings.warn("No dataset specified. Defaulting to cnn_dailymail.") + tokenizer = get_tokenizer(args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code) + default_padding_side = tokenizer.padding_side + # Left padding usually provides better calibration result. + tokenizer.padding_side = "left" if args.sparsity_fmt != "dense": if args.batch_size == 0: @@ -183,8 +217,11 @@ def main(args): # Different calibration datasets are also available, e.g., "pile" and "wikipedia" # Please also check the docstring for the datasets available + assert tokenizer is not None and isinstance( + tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast) + ), "The PreTrainedTokenizer must be set" calib_dataloader = get_dataset_dataloader( - dataset_name=args.dataset_name, + dataset_name=args.dataset, tokenizer=tokenizer, batch_size=args.batch_size, num_samples=args.calib_size, @@ -199,11 +236,11 @@ def main(args): if ( not args.auto_quantize_bits - and args.qformat in ["fp8", "int8_sq", "int4_awq", "w4a8_awq"] + and args.qformat in ["fp8", "int8_sq", "int4_awq", "w4a8_awq", "nvfp4", "nvfp4_awq"] and not args.naive_quantization ) or args.auto_quantize_bits: # If any qformat provided is not fp8, assert model is on GPU - if args.qformat not in ["fp8"]: + if args.qformat not in ["fp8", "nvfp4"]: assert is_model_on_gpu(model), ( f"Model must be fully loaded onto GPUs for {args.qformat} calibration. " "Please make sure the system has enough GPU memory to load the model." @@ -232,14 +269,29 @@ def main(args): print(f"Use calib batch_size {args.batch_size}") - calib_dataloader = get_dataset_dataloader( - dataset_name=args.dataset_name, - tokenizer=tokenizer, - batch_size=args.batch_size, - num_samples=args.calib_size, - device=device, - include_labels=args.auto_quantize_bits is not None, - ) + calib_dataloader = None + if model_type == "mllama": + assert processor is not None and isinstance( + processor, MllamaImageProcessor + ), "The MllamaImageProcessor must be set." + calib_dataloader = get_vlm_dataset_dataloader( + dataset_name=args.dataset, + processor=processor, + batch_size=args.batch_size, + num_samples=args.calib_size, + ) + else: + assert tokenizer is not None and isinstance( + tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast) + ), "The PreTrainedTokenizer must be set" + calib_dataloader = get_dataset_dataloader( + dataset_name=args.dataset, + tokenizer=tokenizer, + batch_size=args.batch_size, + num_samples=args.calib_size, + device=device, + include_labels=args.auto_quantize_bits is not None, + ) quant_cfg = None if not args.auto_quantize_bits: @@ -264,8 +316,8 @@ def main(args): # Always turn on FP8 kv cache to save memory footprint. # For int8_sq, we do not quantize kv cache to preserve accuracy. # We turn off FP8 kv cache for unified_hf checkpoint - enable_quant_kv_cache = "int8_sq" not in args.qformat and "hf" not in args.export_fmt - print(f'{"Enable" if enable_quant_kv_cache else "Disable"} KV cache quantization') + enable_quant_kv_cache = "int8_sq" not in args.qformat + print(f"{'Enable' if enable_quant_kv_cache else 'Disable'} KV cache quantization") quant_cfg["quant_cfg"]["*output_quantizer"] = { "num_bits": 8 if args.qformat == "int8_sq" else (4, 3), "axis": None, @@ -287,17 +339,35 @@ def main(args): # Run some samples generated_ids_after_ptq = model.generate(input_ids, max_new_tokens=100) + def input_decode(input_ids): + if processor is not None and isinstance(processor, MllamaImageProcessor): + return processor.tokenizer.batch_decode(input_ids) + elif tokenizer is not None: + return tokenizer.batch_decode(input_ids) + else: + raise ValueError("The processor or tokenizer must be set") + + def output_decode(generated_ids, input_shape): + if tokenizer is not None and is_enc_dec(model_type): + return tokenizer.batch_decode(generated_ids, skip_special_tokens=True) + elif processor is not None and isinstance(processor, MllamaImageProcessor): + return processor.tokenizer.batch_decode(generated_ids[:, input_shape:]) + elif tokenizer is not None: + return tokenizer.batch_decode(generated_ids[:, input_shape:]) + else: + raise ValueError("The processor or tokenizer must be set") + print("--------") - print(f"example test input: {tokenizer.batch_decode(input_ids)}") + print(f"example test input: {input_decode(input_ids)}") print("--------") print( - "example outputs before ptq: " - f"{tokenizer.batch_decode(generated_ids_before_ptq[:, input_ids.shape[1]:])}" + f"example outputs before ptq: {output_decode(generated_ids_before_ptq, input_ids.shape[1])}" ) print("--------") print( - f"example outputs after ptq: {tokenizer.batch_decode(generated_ids_after_ptq[:, input_ids.shape[1]:])}" + f"example outputs after ptq: {output_decode(generated_ids_after_ptq, input_ids.shape[1])}" ) + else: assert model_type != "dbrx", f"Does not support export {model_type} without quantizaton" print(f"No quantization applied, export {device} model") @@ -309,18 +379,28 @@ def main(args): export_path = args.export_path + if model_type == "mllama": + full_model_config = model.config + model = model.language_model + # TRT-LLM expects both the vision_config and text_config to be set for export. + setattr(model.config, "vision_config", full_model_config.vision_config) + setattr(model.config, "text_config", full_model_config.text_config) + setattr(model.config, "architectures", full_model_config.architectures) + start_time = time.time() if args.export_fmt == "tensorrt_llm": # Move meta tensor back to device before exporting. remove_hook_from_module(model, recurse=True) + dtype = None if "w4a8_awq" in args.qformat: # TensorRT-LLM w4a8 only support fp16 as the dtype. - model = model.to(torch.float16) + dtype = torch.float16 export_tensorrt_llm_checkpoint( model, model_type, + dtype=dtype, export_dir=export_path, inference_tensor_parallel=args.inference_tensor_parallel, inference_pipeline_parallel=args.inference_pipeline_parallel, @@ -334,8 +414,10 @@ def main(args): else: raise NotImplementedError(f"{args.export_fmt} not supported") - # Export the tokenizer as well. - tokenizer.save_pretrained(export_path) + # Restore default padding and export the tokenizer as well. + if tokenizer is not None: + tokenizer.padding_side = default_padding_side + tokenizer.save_pretrained(export_path) end_time = time.time() print( @@ -369,9 +451,7 @@ if __name__ == "__main__": "--calib_size", help="Number of samples for calibration.", type=int, default=512 ) parser.add_argument("--export_path", default="exported_model") - parser.add_argument( - "--dataset_name", help="name of dataset.", type=str, default="cnn_dailymail" - ) + parser.add_argument("--dataset", help="name of dataset.", type=str, default=None) parser.add_argument("--inference_tensor_parallel", type=int, default=1) parser.add_argument("--inference_pipeline_parallel", type=int, default=1) parser.add_argument("--awq_block_size", default=0, type=int) @@ -404,6 +484,12 @@ if __name__ == "__main__": choices=["tensorrt_llm", "hf"], help=("Checkpoint export format"), ) + parser.add_argument( + "--trust_remote_code", + help="Set trust_remote_code for Huggingface models and tokenizers", + default=False, + action="store_true", + ) args = parser.parse_args() main(args) diff --git a/llm_ptq/modelopt_to_tensorrt_llm.py b/examples/llm_ptq/modelopt_to_tensorrt_llm.py similarity index 52% rename from llm_ptq/modelopt_to_tensorrt_llm.py rename to examples/llm_ptq/modelopt_to_tensorrt_llm.py index 6de4f1f21f..d126378c8b 100644 --- a/llm_ptq/modelopt_to_tensorrt_llm.py +++ b/examples/llm_ptq/modelopt_to_tensorrt_llm.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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. """An example to convert an Model Optimizer exported model to tensorrt_llm.""" @@ -30,7 +24,8 @@ from typing import Optional, Union import tensorrt_llm import torch from packaging.version import parse -from tensorrt_llm.builder import BuildConfig +from tensorrt_llm.llmapi import BuildConfig +from tensorrt_llm.llmapi.llm import LLM from tensorrt_llm.models import PretrainedConfig from transformers import AutoTokenizer @@ -69,6 +64,8 @@ def build_tensorrt_llm( max_prompt_embedding_table_size: int = BuildConfig.max_prompt_embedding_table_size, max_encoder_input_len: int = BuildConfig.max_encoder_input_len, perf_mode: bool = False, + checkpoint_format: str = "tensorrt_llm", + tp: int = 1, ): """The API to convert the TensorRT-LLM checkpoint to engines. @@ -100,13 +97,14 @@ def build_tensorrt_llm( feature embeddings or prompt tuning embeddings) to the LLM input embeddings. max_encoder_input_len: Maximum encoder input length for enc-dec models. perf_mode: Whether build the engine with max perf at a cost of longer build time and less flexibility. + checkpoint_format: The model checkpoint format. Choose between [tensorrt_llm, hf]. + tp: tensor_parallel_size. Effective for hf checkpoint_format only. """ engine_dir = Path(engine_dir) engine_dir.mkdir(parents=True, exist_ok=True) pretrained_config_path = Path(pretrained_config) assert pretrained_config_path.exists() - config = PretrainedConfig.from_json_file(pretrained_config_path) ckpt_dir = pretrained_config_path.parent timing_cache_file = ( @@ -117,90 +115,125 @@ def build_tensorrt_llm( ) timing_cache_path = engine_dir / timing_cache_file - log_level = "warning" + if not max_num_tokens: + # tensorrt-llm recommends max max_num_tokens to be 16384 + max_num_tokens = min(max_batch_size * max_input_len, 16384) - if max_batch_size < 4: - warnings.warn( - "TensorRT LLM may hit a runtime issue with batch size is smaller than 4 on some models. Force setting to 4" + # We use llmapi to build the HF checkpoints while trtllm-build the tensorrt_llm checkpoints. + if checkpoint_format == "hf": + build_config = BuildConfig( + max_input_len=max_input_len, + max_seq_len=max_input_len + max_output_len, + max_batch_size=max_batch_size, + max_num_tokens=max_num_tokens, + max_prompt_embedding_table_size=max_prompt_embedding_table_size, + gather_context_logits=not perf_mode, + input_timing_cache=str(timing_cache_path), + output_timing_cache=str(timing_cache_path), + max_encoder_input_len=max_encoder_input_len, ) - max_batch_size = 4 + print(f"trtllm build_config:\n{build_config}") + llm = LLM(ckpt_dir, tensor_parallel_size=tp, build_config=build_config) + llm.save(engine_dir) - is_no_quant_or_fp8 = config.quantization.quant_algo in [ - "FP8", - None, - ] + elif checkpoint_format == "tensorrt_llm": + config = PretrainedConfig.from_json_file(pretrained_config_path) - use_fused_mlp = "RecurrentGemma" not in config.architecture - if config.quantization.exclude_modules: - for module_name in config.quantization.exclude_modules: - # fp8_context_fhma requires all attention.dense to be quantized - if "attention.dense" in module_name: - is_no_quant_or_fp8 = False - # For AutoQuant, fc and gate might not be quantized at the same time - # TODO: relax this limitation on the TRT-LLM side - if "gate" in module_name or "fc" in module_name: - use_fused_mlp = False + log_level = "warning" - quant_algo = config.quantization.quant_algo - use_qdq = quant_algo in ["FP8", "W8A8_SQ_PER_CHANNEL"] + use_paged_context_fmha = config.quantization.quant_algo in [ + "FP8", + "W4A8_AWQ", + "NVFP4", + None, + ] - speculative_decoding_mode = "medusa" if "Medusa" in config.architecture else None + use_fused_mlp = "RecurrentGemma" not in config.architecture + if config.quantization.exclude_modules: + for module_name in config.quantization.exclude_modules: + # fp8_context_fhma requires all attention.dense to be quantized + if "attention.dense" in module_name: + use_paged_context_fmha = False + # For AutoQuant, fc and gate might not be quantized at the same time + # TODO: relax this limitation on the TRT-LLM side + if "gate" in module_name or "fc" in module_name: + use_fused_mlp = False - build_cmd = "trtllm-build " - build_cmd += f"--checkpoint_dir {ckpt_dir} " - build_cmd += f"--input_timing_cache {timing_cache_path} " - build_cmd += f"--output_timing_cache {timing_cache_path} " - build_cmd += f"--log_level {log_level} " - build_cmd += f"--output_dir {engine_dir} " - build_cmd += f"--workers {num_build_workers} " - build_cmd += f"--max_batch_size {max_batch_size} " - build_cmd += f"--max_input_len {max_input_len} " - build_cmd += f"--max_seq_len {max_output_len + max_input_len} " - build_cmd += f"--max_beam_width {max_beam_width} " - build_cmd += f"--max_prompt_embedding_table_size {max_prompt_embedding_table_size} " - build_cmd += f"--max_encoder_input_len {max_encoder_input_len} " - build_cmd += f"--max_num_tokens {max_batch_size * max_input_len} " - build_cmd += ( - "--reduce_fusion enable " - if config.mapping.pp_size == 1 - and config.architecture - not in ["DbrxForCausalLM", "BaichuanForCausalLM", "QWenForCausalLM", "GPTForCausalLM"] - else "" - ) + quant_algo = config.quantization.quant_algo + use_qdq = quant_algo in ["FP8", "W8A8_SQ_PER_CHANNEL"] - if use_fused_mlp: - build_cmd += "--use_fused_mlp enable " - else: - build_cmd += "--use_fused_mlp disable " + speculative_decoding_mode = "medusa" if "Medusa" in config.architecture else None - if enable_sparsity: - build_cmd += "--weight_sparsity " + if num_build_workers > torch.cuda.device_count(): + num_build_workers = torch.cuda.device_count() + print(f"Cap num_build_workers to num gpus: ${num_build_workers}") - if not use_qdq: - build_cmd += "--gemm_plugin auto " + build_cmd = "trtllm-build " + build_cmd += f"--checkpoint_dir {ckpt_dir} " + build_cmd += f"--input_timing_cache {timing_cache_path} " + build_cmd += f"--output_timing_cache {timing_cache_path} " + build_cmd += f"--log_level {log_level} " + build_cmd += f"--output_dir {engine_dir} " + build_cmd += f"--workers {num_build_workers} " + build_cmd += f"--max_batch_size {max_batch_size} " + build_cmd += f"--max_input_len {max_input_len} " + build_cmd += f"--max_seq_len {max_output_len + max_input_len} " + build_cmd += f"--max_beam_width {max_beam_width} " + build_cmd += f"--max_prompt_embedding_table_size {max_prompt_embedding_table_size} " + build_cmd += f"--max_encoder_input_len {max_encoder_input_len} " + build_cmd += ( + "--reduce_fusion enable " + if config.mapping.pp_size == 1 + and config.architecture + not in [ + "DbrxForCausalLM", + "BaichuanForCausalLM", + "QWenForCausalLM", + "GPTForCausalLM", + ] + else "" + ) + + if use_fused_mlp: + build_cmd += "--use_fused_mlp enable " + else: + build_cmd += "--use_fused_mlp disable " + + if enable_sparsity: + build_cmd += "--weight_sparsity " + + # Low batch size scenario + if max_batch_size <= 4 and quant_algo == "FP8": + build_cmd += "--gemm_plugin fp8 " + if quant_algo == "NVFP4": + build_cmd += "--gemm_plugin nvfp4 " + elif not use_qdq: + build_cmd += "--gemm_plugin auto " - if max_num_tokens: build_cmd += f"--max_num_tokens {max_num_tokens} " - if speculative_decoding_mode: - build_cmd += f"--speculative_decoding_mode {speculative_decoding_mode} " + if speculative_decoding_mode: + build_cmd += f"--speculative_decoding_mode {speculative_decoding_mode} " - if is_no_quant_or_fp8: - build_cmd += "--use_paged_context_fmha enable " + if use_paged_context_fmha: + build_cmd += "--use_paged_context_fmha enable " - if perf_mode: - build_cmd += "--multiple_profiles enable" - elif not speculative_decoding_mode: - build_cmd += "--gather_context_logits " # for evaluation benchmarking purpose + if perf_mode: + build_cmd += "--multiple_profiles enable" + elif not speculative_decoding_mode: + build_cmd += "--gather_context_logits " # for evaluation benchmarking purpose - print(f"trtllm-build command:\n{build_cmd}") + print(f"trtllm-build command:\n{build_cmd}") - assert parse(tensorrt_llm.__version__) >= parse(MIN_TENSORRT_LLM_VERSION), ( - f"Detected lower version of tensorrt_llm installed instead of {MIN_TENSORRT_LLM_VERSION}. " - f"Please build the tensorrt_llm engines with tensorrt_llm version {MIN_TENSORRT_LLM_VERSION} or higher instead." - f"\n\n Build command: {build_cmd}" - ) - subprocess.run(build_cmd, shell=True, check=True) + assert parse(tensorrt_llm.__version__) >= parse(MIN_TENSORRT_LLM_VERSION), ( + f"Detected lower version of tensorrt_llm installed instead of {MIN_TENSORRT_LLM_VERSION}. " + f"Please build the tensorrt_llm engines with tensorrt_llm version {MIN_TENSORRT_LLM_VERSION} " + " or higher instead.\n\n Build command: {build_cmd}" + ) + subprocess.run(build_cmd, shell=True, check=True) + + else: + raise NotImplementedError(f"checkpoint_format: {checkpoint_format} not supported") try: tokenizer = AutoTokenizer.from_pretrained(ckpt_dir) @@ -247,6 +280,25 @@ def parse_arguments(): default=BuildConfig.max_encoder_input_len, help="Maximum encoder input length for enc-dec models.", ) + parser.add_argument( + "--checkpoint_format", + type=str, + help="The exported checkpoint format.", + choices=["tensorrt_llm", "hf"], + default="tensorrt_llm", + ) + parser.add_argument( + "--tp", + type=int, + default=1, + help="The tensor_parallel_size for the HuggingFace checkpoint engine build.", + ) + parser.add_argument( + "--trust_remote_code", + help="Set trust_remote_code for Huggingface models and tokenizers", + default=False, + action="store_true", + ) return parser.parse_args() @@ -264,6 +316,8 @@ def main(args): max_prompt_embedding_table_size=args.max_prompt_embedding_table_size, max_encoder_input_len=args.max_encoder_input_len, perf_mode=args.perf, + checkpoint_format=args.checkpoint_format, + tp=args.tp, ) if ( diff --git a/llm_ptq/nemo_ptq.py b/examples/llm_ptq/nemo_ptq.py similarity index 92% rename from llm_ptq/nemo_ptq.py rename to examples/llm_ptq/nemo_ptq.py index bbbe7df986..b6fa911300 100644 --- a/llm_ptq/nemo_ptq.py +++ b/examples/llm_ptq/nemo_ptq.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 diff --git a/llm_ptq/requirements.txt b/examples/llm_ptq/requirements.txt similarity index 75% rename from llm_ptq/requirements.txt rename to examples/llm_ptq/requirements.txt index 8999f8d840..d3d2dd8027 100644 --- a/llm_ptq/requirements.txt +++ b/examples/llm_ptq/requirements.txt @@ -1,10 +1,9 @@ bitsandbytes -datasets>=2.14.4 fire flash-attn>=2.6.0 pynvml>=11.5.0 rouge_score>=0.1.2 tiktoken -transformers>=4.43.3 +transformers>=4.45.0 transformers_stream_generator zstandard diff --git a/llm_ptq/run_tensorrt_llm.py b/examples/llm_ptq/run_tensorrt_llm.py similarity index 63% rename from llm_ptq/run_tensorrt_llm.py rename to examples/llm_ptq/run_tensorrt_llm.py index f537dbd0d4..04071febf0 100644 --- a/llm_ptq/run_tensorrt_llm.py +++ b/examples/llm_ptq/run_tensorrt_llm.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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. """An example script to run the tensorrt_llm engine.""" @@ -43,6 +37,12 @@ def parse_arguments(): ), help="Input texts. Please use | to separate different batches.", ) + parser.add_argument( + "--trust_remote_code", + help="Set trust_remote_code for Huggingface models and tokenizers", + default=False, + action="store_true", + ) return parser.parse_args() @@ -55,7 +55,9 @@ def run(args): if isinstance(args.tokenizer, PreTrainedTokenizerBase): tokenizer = args.tokenizer else: - tokenizer = get_tokenizer(ckpt_path=args.tokenizer) + tokenizer = get_tokenizer( + ckpt_path=args.tokenizer, trust_remote_code=args.trust_remote_code + ) input_texts = args.input_texts.split("|") assert input_texts, "input_text not specified" diff --git a/llm_ptq/scripts/huggingface_example.sh b/examples/llm_ptq/scripts/huggingface_example.sh similarity index 80% rename from llm_ptq/scripts/huggingface_example.sh rename to examples/llm_ptq/scripts/huggingface_example.sh index d486d5f896..b09e3126b8 100755 --- a/llm_ptq/scripts/huggingface_example.sh +++ b/examples/llm_ptq/scripts/huggingface_example.sh @@ -1,4 +1,19 @@ #!/bin/bash +# 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. + set -e script_dir="$(dirname "$(readlink -f "$0")")" @@ -30,10 +45,10 @@ if [ "$EXPORT_FORMAT" = "hf" ]; then IFS="," for qformat in $QFORMAT; do case $qformat in - fp8|int4_awq) + fp8|int4_awq|nvfp4|nvfp4_awq|w4a8_awq) ;; *) - echo "Unsupported quant argument: Expected one of: [fp8, int4_awq]" >&2 + echo "Unsupported quant argument: Expected one of: [fp8, int4_awq, nvfp4, nvfp4_awq, w4a8_awq]" >&2 exit 1 esac done @@ -59,10 +74,10 @@ esac IFS="," for qformat in $QFORMAT; do case $qformat in - fp8|fp8_naive|int8_sq|int4_awq|w4a8_awq|fp16|bf16) + fp8|fp8_naive|int8_sq|int4_awq|w4a8_awq|fp16|bf16|nvfp4|nvfp4_awq) ;; *) - echo "Unknown quant argument: Expected one of: [fp8, fp8_naive, int8_sq, int4_awq, w4a8_awq, fp16, bf16]" >&2 + echo "Unknown quant argument: Expected one of: [fp8, fp8_naive, int8_sq, int4_awq, w4a8_awq, fp16, bf16, nvfp4, nvfp4_awq]" >&2 exit 1 esac done @@ -102,7 +117,7 @@ MODEL_NAME=$(basename $MODEL_PATH | sed 's/[^0-9a-zA-Z\-]/_/g') SAVE_PATH=${ROOT_SAVE_PATH}/saved_models_${MODEL_NAME}_${SPARSITY_FMT}_${QFORMAT_MODIFIED}_tp${TP}_pp${PP} if [ $EXPORT_FORMAT != "tensorrt_llm" ]; then - SAVE_PATH=${SAVE_PATH}_${EXPORT_FORMAT} + SAVE_PATH=${ROOT_SAVE_PATH}/saved_models_${MODEL_NAME}_${QFORMAT_MODIFIED}_${EXPORT_FORMAT} fi MODEL_CONFIG=${SAVE_PATH}/config.json @@ -115,11 +130,15 @@ fi PTQ_ARGS="" if [ $QFORMAT == "fp8_naive" ]; then QFORMAT=fp8 - PTQ_ARGS+="--naive_quantization" + PTQ_ARGS+=" --naive_quantization " fi if [ -n "$AUTO_QUANTIZE_BITS" ]; then - PTQ_ARGS+="--auto_quantize_bits=$AUTO_QUANTIZE_BITS" + PTQ_ARGS+=" --auto_quantize_bits=$AUTO_QUANTIZE_BITS " +fi + +if $TRUST_REMOTE_CODE; then + PTQ_ARGS+=" --trust_remote_code " fi AWQ_ARGS="" @@ -152,11 +171,6 @@ if [[ $TASKS =~ "build" ]] || [[ ! -d "$ENGINE_DIR" ]] || [[ ! $(ls -A $ENGINE_D echo "Quantized model config $MODEL_CONFIG exists, skipping the quantization stage" fi - if [[ $EXPORT_FORMAT != "tensorrt_llm" ]]; then - echo "Please continue to deployment with $EXPORT_FORMAT. Checkpoint export_path: $SAVE_PATH" - exit 0 - fi - # for enc-dec model, users need to refer TRT-LLM example to build engines and deployment if [[ -f "$SAVE_PATH/encoder/config.json" && -f "$SAVE_PATH/decoder/config.json" && ! -f $MODEL_CONFIG ]]; then echo "Please continue to deployment with the TRT-LLM enc_dec example, https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/enc_dec. Checkpoint export_path: $SAVE_PATH" @@ -164,15 +178,17 @@ if [[ $TASKS =~ "build" ]] || [[ ! -d "$ENGINE_DIR" ]] || [[ ! $(ls -A $ENGINE_D fi if [[ "$QFORMAT" == *"nvfp4"* ]]; then - echo "Please build tensorrt_llm engine with this model from the tensorrt_llm repo for deployment. Checkpoint export_path: $SAVE_PATH" - exit 0 + compute_capability=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader,nounits | head -n 1) + if awk "BEGIN {exit !($compute_capability < 10.0)}"; then + echo "Please build tensorrt_llm engine with this model from the tensorrt_llm repo for deployment. Checkpoint export_path: $SAVE_PATH" + fi fi echo "Building tensorrt_llm engine from Model Optimizer-quantized model..." - perf="" + BUILD_ARGS="" if [[ $TASKS =~ "benchmark" && ! $TASKS =~ "lm_eval" ]]; then - perf=" --perf " + BUILD_ARGS+=" --perf " fi python modelopt_to_tensorrt_llm.py \ @@ -184,7 +200,9 @@ if [[ $TASKS =~ "build" ]] || [[ ! -d "$ENGINE_DIR" ]] || [[ ! $(ls -A $ENGINE_D --max_batch_size=$BUILD_MAX_BATCH_SIZE \ --num_build_workers=$GPUS \ --enable_sparsity=$ENABLE_SPARSITY \ - $perf + --checkpoint_format=$EXPORT_FORMAT \ + --tp=$TP \ + $BUILD_ARGS fi if [[ -d "${MODEL_PATH}" ]]; then @@ -204,7 +222,11 @@ fi lm_eval_flags="" if [ -n "$LM_EVAL_LIMIT" ]; then - lm_eval_flags=" --limit $LM_EVAL_LIMIT " + lm_eval_flags+=" --limit $LM_EVAL_LIMIT " +fi + +if $TRUST_REMOTE_CODE; then + lm_eval_flags+=" --trust_remote_code " fi LM_EVAL_RESULT=${ENGINE_DIR}/lm_eval.txt @@ -322,11 +344,11 @@ fi TRT_LLM_PREPARE_DATASET=$TRT_LLM_CODE_PATH/benchmarks/cpp/prepare_dataset.py # Align with the official benchmark -BENCHMARK_INPUT_LEN=2048 -BENCHMARK_OUTPUT_LEN=128 -BENCHMARK_NUM_REQUESTS=200 +BENCHMARK_INPUT_LEN=$BUILD_MAX_INPUT_LEN +BENCHMARK_OUTPUT_LEN=$BUILD_MAX_OUTPUT_LEN +BENCHMARK_NUM_REQUESTS=256 -DATASET_TXT=${SAVE_PATH}/synthetic_${BENCHMARK_INPUT_LEN}_${BENCHMARK_OUTPUT_LEN}.txt +DATASET_TXT=${SAVE_PATH}/synthetic_${BENCHMARK_INPUT_LEN}_${BENCHMARK_OUTPUT_LEN}_${BENCHMARK_NUM_REQUESTS}.txt if [ -z "$TRT_LLM_PREPARE_DATASET" ]; then echo "Unable to prepare dataset for benchmarking. Please set TRT_LLM_CODE_PATH to the TRT-LLM code path." @@ -340,16 +362,20 @@ else fi fi -python benchmarks/benchmark_suite.py --model $MODEL_PATH latency --engine_dir $ENGINE_DIR --dataset $DATASET_TXT | tee $BENCHMARK_RESULT -python benchmarks/benchmark_suite.py --model $MODEL_PATH throughput --engine_dir $ENGINE_DIR --dataset $DATASET_TXT | tee -a $BENCHMARK_RESULT +if [ "$BUILD_MAX_BATCH_SIZE" -gt 1 ]; then + python benchmarks/benchmark_suite.py --model $MODEL_PATH throughput --engine_dir $ENGINE_DIR --dataset $DATASET_TXT | tee -a $BENCHMARK_RESULT +else + python benchmarks/benchmark_suite.py --model $MODEL_PATH latency --engine_dir $ENGINE_DIR --dataset $DATASET_TXT | tee $BENCHMARK_RESULT +fi fi if [ -n "$FREE_SPACE" ]; then rm -f $SAVE_PATH/*.json rm -f $SAVE_PATH/*.safetensors - rm -f $ENGINE_DIR/*.json - rm -f $ENGINE_DIR/*.engine + rm -f $SAVE_PATH/*/*.json + rm -f $SAVE_PATH/*/*.engine + rm -f $SAVE_PATH/*/*.cache fi popd diff --git a/llm_ptq/scripts/nemo_example.sh b/examples/llm_ptq/scripts/nemo_example.sh similarity index 90% rename from llm_ptq/scripts/nemo_example.sh rename to examples/llm_ptq/scripts/nemo_example.sh index 868ed506b1..484455e7ba 100755 --- a/llm_ptq/scripts/nemo_example.sh +++ b/examples/llm_ptq/scripts/nemo_example.sh @@ -1,4 +1,19 @@ #!/bin/bash +# 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. + set -e script_dir="$(dirname "$(readlink -f "$0")")" @@ -37,10 +52,10 @@ fi IFS=',' for qformat in $QFORMAT; do case $qformat in - bf16|fp16|fp8|int8_sq|int4_awq|w4a8_awq) + bf16|fp16|fp8|int8_sq|int4_awq|w4a8_awq|nvfp4) ;; *) - echo "Unknown quant argument: Expected one of: [bf16, fp16, fp8, int8_sq, int4_awq, w4a8_awq]" >&2 + echo "Unknown quant argument: Expected one of: [bf16, fp16, fp8, int8_sq, int4_awq, w4a8_awq, nvfp4]" >&2 exit 1 esac done diff --git a/llm_ptq/scripts/parser.sh b/examples/llm_ptq/scripts/parser.sh similarity index 83% rename from llm_ptq/scripts/parser.sh rename to examples/llm_ptq/scripts/parser.sh index d8ddb82d17..68f62cad56 100644 --- a/llm_ptq/scripts/parser.sh +++ b/examples/llm_ptq/scripts/parser.sh @@ -1,5 +1,20 @@ #!/bin/bash +# 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. + # Define a function to parse command-line options parse_options() { # Default values @@ -17,8 +32,10 @@ parse_options() { TASKS="build" + TRUST_REMOTE_CODE=false + # Parse command-line options - ARGS=$(getopt -o "" -l "type:,model:,quant:,tp:,calib_tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,effective_bits:,input:,output:,batch:,tasks:,export_fmt:,lm_eval_tasks:,lm_eval_limit:" -n "$0" -- "$@") + ARGS=$(getopt -o "" -l "type:,model:,quant:,tp:,calib_tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,effective_bits:,input:,output:,batch:,tasks:,export_fmt:,lm_eval_tasks:,lm_eval_limit:,trust_remote_code" -n "$0" -- "$@") eval set -- "$ARGS" while true; do case "$1" in @@ -41,6 +58,7 @@ parse_options() { --lm_eval_tasks ) LM_EVAL_TASKS="$2"; shift 2;; --lm_eval_limit ) LM_EVAL_LIMIT="$2"; shift 2;; --num_samples ) NUM_SAMPLES="$2"; shift 2;; + --trust_remote_code ) TRUST_REMOTE_CODE=true; shift;; -- ) shift; break ;; * ) break ;; esac @@ -79,7 +97,7 @@ parse_options() { exit 1 fi - VALID_TASKS=("build" "mmlu" "humaneval" "mtbench" "benchmark" "lm_eval") + VALID_TASKS=("build" "mmlu" "humaneval" "mtbench" "benchmark" "lm_eval" "gqa") for task in $(echo $TASKS | tr ',' ' '); do if [[ ! " ${VALID_TASKS[@]} " =~ " $task " ]]; then diff --git a/llm_qat/README.md b/examples/llm_qat/README.md similarity index 100% rename from llm_qat/README.md rename to examples/llm_qat/README.md diff --git a/llm_qat/launch.sh b/examples/llm_qat/launch.sh similarity index 83% rename from llm_qat/launch.sh rename to examples/llm_qat/launch.sh index 9db4f806f2..86b8eb9e03 100755 --- a/llm_qat/launch.sh +++ b/examples/llm_qat/launch.sh @@ -1,4 +1,19 @@ #!/bin/bash +# 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. + set -e set -x set -o pipefail diff --git a/llm_qat/main.py b/examples/llm_qat/main.py similarity index 87% rename from llm_qat/main.py rename to examples/llm_qat/main.py index 923fceeed5..53bfd92d1a 100644 --- a/llm_qat/main.py +++ b/examples/llm_qat/main.py @@ -15,25 +15,19 @@ # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 dataclasses import dataclass, field diff --git a/llm_qat/requirements.txt b/examples/llm_qat/requirements.txt similarity index 77% rename from llm_qat/requirements.txt rename to examples/llm_qat/requirements.txt index 9275c19e74..ac8cd39eea 100644 --- a/llm_qat/requirements.txt +++ b/examples/llm_qat/requirements.txt @@ -1,9 +1,7 @@ accelerate>=0.27.2,!=0.34.* -datasets>=2.16.1 flash-attn peft py7zr sentencepiece>=0.2.0 tensorboardX -torch>=2.1.0 transformers>=4.47.0 diff --git a/llm_qat/utils.py b/examples/llm_qat/utils.py similarity index 84% rename from llm_qat/utils.py rename to examples/llm_qat/utils.py index 7339490d97..fe0a2014c3 100644 --- a/llm_qat/utils.py +++ b/examples/llm_qat/utils.py @@ -1,28 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 gc import types from functools import partial -from typing import Dict import datasets import transformers @@ -140,7 +133,7 @@ def get_daring_anteater(tokenizer: transformers.AutoTokenizer, split="train", ma def make_supervised_data_module( dataset="samsum", tokenizer: transformers.PreTrainedTokenizer = None -) -> Dict: +) -> dict: """Make dataset and collmtor for supervised fine-tuning.""" if dataset == "samsum": train_dataset = get_preprocessed_samsum(tokenizer, "train", tokenizer.model_max_length) diff --git a/llm_sparsity/README.md b/examples/llm_sparsity/README.md similarity index 100% rename from llm_sparsity/README.md rename to examples/llm_sparsity/README.md diff --git a/llm_sparsity/data_prep.py b/examples/llm_sparsity/data_prep.py similarity index 60% rename from llm_sparsity/data_prep.py rename to examples/llm_sparsity/data_prep.py index fe0176e70c..d92f17d2e7 100644 --- a/llm_sparsity/data_prep.py +++ b/examples/llm_sparsity/data_prep.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 json diff --git a/llm_sparsity/eval.py b/examples/llm_sparsity/eval.py similarity index 87% rename from llm_sparsity/eval.py rename to examples/llm_sparsity/eval.py index 6f55160b72..b78d23366f 100644 --- a/llm_sparsity/eval.py +++ b/examples/llm_sparsity/eval.py @@ -1,29 +1,23 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 json import os from dataclasses import dataclass -from typing import Dict, Sequence +from typing import Sequence import evaluate import nltk @@ -49,7 +43,7 @@ G_PROMPT_INPUT = ( def smart_tokenizer_and_embedding_resize( - special_tokens_dict: Dict, + special_tokens_dict: dict, tokenizer: transformers.PreTrainedTokenizer, model: transformers.PreTrainedModel, ): @@ -116,7 +110,7 @@ class CNNDailymailDataset(Dataset): self.sources, self.labels = preprocess_cnndailymail(accelerator, data_path, calib=False) self.tokenizer = tokenizer - def __getitem__(self, index) -> Dict[str, torch.Tensor]: + def __getitem__(self, index) -> dict[str, torch.Tensor]: return dict(src_idx=self.sources[index], label_idx=self.labels[index]) def __len__(self): @@ -130,7 +124,7 @@ class DataCollator(object): tokenizer: transformers.PreTrainedTokenizer model_max_length: int - def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]: + def __call__(self, instances: Sequence[dict]) -> dict[str, torch.Tensor]: sources, labels = tuple( [instance[key] for instance in instances] for key in ("src_idx", "label_idx") ) diff --git a/llm_sparsity/export_trtllm_ckpt.py b/examples/llm_sparsity/export_trtllm_ckpt.py similarity index 77% rename from llm_sparsity/export_trtllm_ckpt.py rename to examples/llm_sparsity/export_trtllm_ckpt.py index ebdcadc83c..7abdea71e0 100644 --- a/llm_sparsity/export_trtllm_ckpt.py +++ b/examples/llm_sparsity/export_trtllm_ckpt.py @@ -1,30 +1,23 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import argparse import os import random import sys from pathlib import Path -from typing import Dict import numpy as np import torch @@ -41,7 +34,7 @@ DEFAULT_PAD_TOKEN = "[PAD]" def smart_tokenizer_and_embedding_resize( - special_tokens_dict: Dict, + special_tokens_dict: dict, tokenizer: PreTrainedTokenizer, model: PreTrainedModel, ): @@ -63,19 +56,20 @@ def smart_tokenizer_and_embedding_resize( output_embeddings[-num_new_tokens:] = output_embeddings_avg -def get_tokenizer(ckpt_path: str, model_max_length: int): +def get_tokenizer(ckpt_path: str, model_max_length: int, trust_remote_code: bool = False): print(f"Initializing tokenizer from {ckpt_path}") tokenizer = AutoTokenizer.from_pretrained( ckpt_path, model_max_length=model_max_length, padding_side="right", use_fast=False, + trust_remote_code=trust_remote_code, ) return tokenizer -def get_model(ckpt_path, dtype="fp16"): +def get_model(ckpt_path, dtype="fp16", trust_remote_code=False): print(f"Initializing model from {ckpt_path} using dtype {dtype}.") if dtype == "bf16": dtype = torch.bfloat16 @@ -88,7 +82,7 @@ def get_model(ckpt_path, dtype="fp16"): model_kwargs = {"torch_dtype": dtype} model = AutoModelForCausalLM.from_pretrained( - ckpt_path, device_map="auto", **model_kwargs, trust_remote_code=True + ckpt_path, device_map="auto", **model_kwargs, trust_remote_code=trust_remote_code ) model.eval() @@ -102,8 +96,10 @@ def main(args): random.seed(1234) np.random.seed(1234) - model = get_model(args.model_name_or_path, args.dtype) - tokenizer = get_tokenizer(args.model_name_or_path, args.model_max_length) + model = get_model(args.model_name_or_path, args.dtype, trust_remote_code=args.trust_remote_code) + tokenizer = get_tokenizer( + args.model_name_or_path, args.model_max_length, trust_remote_code=args.trust_remote_code + ) if tokenizer.pad_token is None: smart_tokenizer_and_embedding_resize( @@ -167,6 +163,12 @@ if __name__ == "__main__": default=1, ) parser.add_argument("--naive_quantization", default=False, action="store_true") + parser.add_argument( + "--trust_remote_code", + help="Set trust_remote_code for Huggingface models and tokenizers", + default=False, + action="store_true", + ) args = parser.parse_args() diff --git a/llm_sparsity/finetune.py b/examples/llm_sparsity/finetune.py similarity index 91% rename from llm_sparsity/finetune.py rename to examples/llm_sparsity/finetune.py index b57d90f81b..177c12d1cf 100644 --- a/llm_sparsity/finetune.py +++ b/examples/llm_sparsity/finetune.py @@ -15,32 +15,26 @@ # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 pickle from dataclasses import dataclass, field -from typing import Dict, Optional, Sequence +from typing import Optional, Sequence import torch import transformers @@ -83,6 +77,10 @@ class ModelArguments: default=False, metadata={"help": "Enables Flash attention for training."}, ) + trust_remote_code: bool = field( + default=False, + metadata={"help": "Set trust_remote_code for Huggingface models and tokenizers."}, + ) @dataclass @@ -146,7 +144,7 @@ class ModelOptArguments: def smart_tokenizer_and_embedding_resize( - special_tokens_dict: Dict, + special_tokens_dict: dict, tokenizer: transformers.PreTrainedTokenizer, model: transformers.PreTrainedModel, ): @@ -168,7 +166,7 @@ def smart_tokenizer_and_embedding_resize( output_embeddings[-num_new_tokens:] = output_embeddings_avg -def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict: +def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> dict: """Tokenize a list of strings.""" tokenized_list = [ tokenizer( @@ -201,7 +199,7 @@ def preprocess( sources: Sequence[str], targets: Sequence[str], tokenizer: transformers.PreTrainedTokenizer, -) -> Dict: +) -> dict: """Preprocess the data by tokenizing.""" examples = [s + t for s, t in zip(sources, targets)] examples_tokenized, sources_tokenized = [ @@ -262,7 +260,7 @@ class SupervisedDataset(Dataset): def __len__(self): return len(self.input_ids) - def __getitem__(self, i) -> Dict[str, torch.Tensor]: + def __getitem__(self, i) -> dict[str, torch.Tensor]: return dict(input_ids=self.input_ids[i], labels=self.labels[i]) @@ -272,7 +270,7 @@ class DataCollatorForSupervisedDataset(object): tokenizer: transformers.PreTrainedTokenizer - def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]: + def __call__(self, instances: Sequence[dict]) -> dict[str, torch.Tensor]: input_ids, labels = tuple( [instance[key] for instance in instances] for key in ("input_ids", "labels") ) @@ -294,7 +292,7 @@ def make_supervised_data_module( train_datapath: str, val_datapath: str, tokenizer: transformers.PreTrainedTokenizer, -) -> Dict: +) -> dict: """Make dataset and collator for supervised fine-tuning.""" train_dataset = SupervisedDataset(training_args, train_datapath, tokenizer, "train") val_dataset = SupervisedDataset(training_args, val_datapath, tokenizer, "val") @@ -323,7 +321,7 @@ def train(): model = transformers.AutoModelForCausalLM.from_pretrained( model_args.model_name_or_path, - trust_remote_code=True, + trust_remote_code=args.trust_remote_code, cache_dir=args.cache_dir, attn_implementation="flash_attention_2" if model_args.use_flash_attn else "eager", ).to(torch.device("cuda")) @@ -334,6 +332,7 @@ def train(): model_max_length=args.model_max_length, padding_side="right", use_fast=True, + trust_remote_code=args.trust_remote_code, ) if tokenizer.pad_token is None: diff --git a/llm_sparsity/hf_pts.py b/examples/llm_sparsity/hf_pts.py similarity index 79% rename from llm_sparsity/hf_pts.py rename to examples/llm_sparsity/hf_pts.py index 3906c1ea16..7aca669268 100644 --- a/llm_sparsity/hf_pts.py +++ b/examples/llm_sparsity/hf_pts.py @@ -1,29 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import argparse import os import random import time -from typing import Dict import numpy as np import torch @@ -65,7 +58,7 @@ def get_calib_dataloader( def smart_tokenizer_and_embedding_resize( - special_tokens_dict: Dict, + special_tokens_dict: dict, tokenizer: PreTrainedTokenizer, model: PreTrainedModel, ): @@ -87,20 +80,21 @@ def smart_tokenizer_and_embedding_resize( output_embeddings[-num_new_tokens:] = output_embeddings_avg -def get_tokenizer(ckpt_path: str, model_max_length: int): +def get_tokenizer(ckpt_path: str, model_max_length: int, trust_remote_code: bool = False): print(f"Initializing tokenizer from {ckpt_path}") tokenizer = AutoTokenizer.from_pretrained( ckpt_path, model_max_length=model_max_length, padding_side="right", use_fast=False, + trust_remote_code=trust_remote_code, ) return tokenizer -def get_model(ckpt_path, dtype="fp16"): - print(f"Initializing model from {ckpt_path}") +def get_model(ckpt_path, dtype="fp16", trust_remote_code=False): + print(f"Initializing model from {ckpt_path} using dtype {dtype}.") if dtype == "bf16": dtype = torch.bfloat16 elif dtype == "fp16": @@ -112,7 +106,7 @@ def get_model(ckpt_path, dtype="fp16"): model_kwargs = {"torch_dtype": dtype} model = AutoModelForCausalLM.from_pretrained( - ckpt_path, device_map="auto", **model_kwargs, trust_remote_code=True + ckpt_path, device_map="auto", **model_kwargs, trust_remote_code=trust_remote_code ) model.eval() @@ -126,8 +120,10 @@ def main(args): random.seed(1234) np.random.seed(1234) - model = get_model(args.model_name_or_path, args.dtype) - tokenizer = get_tokenizer(args.model_name_or_path, args.model_max_length) + model = get_model(args.model_name_or_path, args.dtype, trust_remote_code=args.trust_remote_code) + tokenizer = get_tokenizer( + args.model_name_or_path, args.model_max_length, trust_remote_code=args.trust_remote_code + ) smart_tokenizer_and_embedding_resize( special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN), tokenizer=tokenizer, @@ -195,6 +191,12 @@ if __name__ == "__main__": default=1, ) parser.add_argument("--naive_quantization", default=False, action="store_true") + parser.add_argument( + "--trust_remote_code", + help="Set trust_remote_code for Huggingface models and tokenizers", + default=False, + action="store_true", + ) args = parser.parse_args() diff --git a/llm_sparsity/launch_finetune.sh b/examples/llm_sparsity/launch_finetune.sh similarity index 76% rename from llm_sparsity/launch_finetune.sh rename to examples/llm_sparsity/launch_finetune.sh index 5b2f536c70..aaf7e899ae 100755 --- a/llm_sparsity/launch_finetune.sh +++ b/examples/llm_sparsity/launch_finetune.sh @@ -1,5 +1,20 @@ #!/bin/bash +# 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. + set -x while [ $# -gt 0 ]; do diff --git a/llm_sparsity/requirements.txt b/examples/llm_sparsity/requirements.txt similarity index 77% rename from llm_sparsity/requirements.txt rename to examples/llm_sparsity/requirements.txt index e488951122..82ac89bf8d 100644 --- a/llm_sparsity/requirements.txt +++ b/examples/llm_sparsity/requirements.txt @@ -1,4 +1,3 @@ -datasets>=2.16.1 flash-attn sentencepiece>=0.2.0 tensorboardX diff --git a/examples/llm_sparsity/utils.py b/examples/llm_sparsity/utils.py new file mode 100644 index 0000000000..132a6c1884 --- /dev/null +++ b/examples/llm_sparsity/utils.py @@ -0,0 +1,31 @@ +# 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 io +import json + + +def _make_r_io_base(f, mode: str): + if not isinstance(f, io.IOBase): + f = open(f, mode=mode) + return f + + +def jload(f, mode="r"): + """Load a .json file into a dictionary.""" + f = _make_r_io_base(f, mode) + jdict = json.load(f) + f.close() + return jdict diff --git a/model_hub/README.md b/examples/model_hub/README.md similarity index 100% rename from model_hub/README.md rename to examples/model_hub/README.md diff --git a/examples/model_hub/download_hf_ckpt.py b/examples/model_hub/download_hf_ckpt.py new file mode 100644 index 0000000000..cc06b57485 --- /dev/null +++ b/examples/model_hub/download_hf_ckpt.py @@ -0,0 +1,31 @@ +# 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 os + +from huggingface_hub import login, snapshot_download + +login(token="", add_to_git_credential=True) +model_names = [ + "Llama-3.1-8B-Instruct-FP8", + "Llama-3.1-70B-Instruct-FP8", + "Llama-3.1-405B-Instruct-FP8", +] +for model_name in model_names: + hf_repo = "nvidia/" + model_name + local_dir = "quantized_ckpts/" + model_name + os.makedirs(local_dir, exist_ok=True) + print(f"downloading {model_name}...") + snapshot_download(repo_id=hf_repo, local_dir=local_dir) diff --git a/model_hub/llama_fp8_deploy_trtllm.sh b/examples/model_hub/llama_fp8_deploy_trtllm.sh similarity index 70% rename from model_hub/llama_fp8_deploy_trtllm.sh rename to examples/model_hub/llama_fp8_deploy_trtllm.sh index 51344322c7..1617527f46 100644 --- a/model_hub/llama_fp8_deploy_trtllm.sh +++ b/examples/model_hub/llama_fp8_deploy_trtllm.sh @@ -1,5 +1,20 @@ #!/bin/bash +# 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. + # Check if enough arguments are provided if [ "$#" -lt 2 ]; then echo "Usage: $0 " diff --git a/model_hub/llama_fp8_deploy_vllm.sh b/examples/model_hub/llama_fp8_deploy_vllm.sh similarity index 67% rename from model_hub/llama_fp8_deploy_vllm.sh rename to examples/model_hub/llama_fp8_deploy_vllm.sh index f5f8e900d0..4b4f960362 100644 --- a/model_hub/llama_fp8_deploy_vllm.sh +++ b/examples/model_hub/llama_fp8_deploy_vllm.sh @@ -1,5 +1,20 @@ #!/bin/bash +# 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. + # Check if enough arguments are provided if [ "$#" -lt 1 ]; then echo "Usage: $0 " diff --git a/onnx_ptq/README.md b/examples/onnx_ptq/README.md similarity index 94% rename from onnx_ptq/README.md rename to examples/onnx_ptq/README.md index c79a719787..e78b8042f9 100644 --- a/onnx_ptq/README.md +++ b/examples/onnx_ptq/README.md @@ -4,7 +4,7 @@ This ONNX PTQ Toolkit provides a comprehensive suite of tools designed to optimi Quantization is a technique that converts a model from using floating-point numbers to lower-precision formats like integers, which are computationally less expensive. This process can significantly speed up execution on compatible hardware and reduce memory and bandwidth requirements. -Note that this example is for ONNX model quantization. For end to end quantization examples with Large Language models, please refer to: [llm_ptq](https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/llm_ptq) +Note that this example is for ONNX model quantization. For end to end quantization examples with Large Language models, please refer to: [llm_ptq](https://github.com/NVIDIA/TensorRT-Model-Optimizer/tree/main/examples/llm_ptq) ## Environment setup @@ -74,7 +74,7 @@ Inference latency of the model is ms ## Quantize an ONNX model with custom op -This feature requires TensorRT 10+ and `ORT>=1.18`. For proper usage, please make sure that the path to `libcudnn*.so` is in the `PATH` or `LD_LIBRARY_PATH` variables and that the `tensorrt` python package is installed. +This feature requires TensorRT 10+ and `ORT>=1.20`. For proper usage, please make sure that the path to `libcudnn*.so` is in the `PATH` or `LD_LIBRARY_PATH` variables and that the `tensorrt` python package is installed. Please see the sample example below. @@ -110,17 +110,17 @@ python scripts/create_identity_neural_network.py This generates the identity_neural_network.onnx model in `TensorRT-Custom-Plugin-Example/data/identity_neural_network.onnx` -**Step 2**: Quantize the ONNX model. We will be using the `libidentity_conv_iplugin_v2.so` plugin for this example. +**Step 2**: Quantize the ONNX model. We will be using the `libidentity_conv_iplugin_v2_io_ext.so` plugin for this example. ```bash python -m modelopt.onnx.quantization \ --onnx_path=/path/to/identity_neural_network.onnx \ - --trt_plugins=/path/to/libidentity_conv_iplugin_v2.so + --trt_plugins=/path/to/libidentity_conv_iplugin_v2_io_ext.so ``` **Step 3**: Deploy the quantized model with TensorRT. ```bash trtexec --onnx=/path/to/identity_neural_network.quant.onnx \ - --staticPlugins=/path/to/libidentity_conv_iplugin_v2.so + --staticPlugins=/path/to/libidentity_conv_iplugin_v2_io_ext.so ``` diff --git a/onnx_ptq/download_example_onnx.py b/examples/onnx_ptq/download_example_onnx.py similarity index 67% rename from onnx_ptq/download_example_onnx.py rename to examples/onnx_ptq/download_example_onnx.py index 481a09220c..fb8476f51f 100644 --- a/onnx_ptq/download_example_onnx.py +++ b/examples/onnx_ptq/download_example_onnx.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. import argparse import os diff --git a/onnx_ptq/evaluate.py b/examples/onnx_ptq/evaluate.py similarity index 79% rename from onnx_ptq/evaluate.py rename to examples/onnx_ptq/evaluate.py index fcd9367dfd..c8ec123775 100644 --- a/onnx_ptq/evaluate.py +++ b/examples/onnx_ptq/evaluate.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 csv diff --git a/onnx_ptq/evaluation.py b/examples/onnx_ptq/evaluation.py similarity index 86% rename from onnx_ptq/evaluation.py rename to examples/onnx_ptq/evaluation.py index 9e6fd22e94..ebc4296dd2 100644 --- a/onnx_ptq/evaluation.py +++ b/examples/onnx_ptq/evaluation.py @@ -1,28 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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. """Module to evaluate a device model for the specified task.""" import random -from typing import Any, Final, Tuple, Union +from typing import Any, Final, Union import torch import torchvision.transforms as transforms @@ -47,7 +41,7 @@ TORCH_MODEL: Final[str] = "torch_model" def evaluate( model: Union[torch.nn.Module, bytes], - dummy_input: Union[Any, Tuple], + dummy_input: Union[Any, tuple], dataset_path: str, evaluation_type: str = ACCURACY, batch_size=1, diff --git a/onnx_ptq/image_prep.py b/examples/onnx_ptq/image_prep.py similarity index 56% rename from onnx_ptq/image_prep.py rename to examples/onnx_ptq/image_prep.py index f540df1427..e075a7f89b 100644 --- a/onnx_ptq/image_prep.py +++ b/examples/onnx_ptq/image_prep.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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. """Utility to dump imagenet data for calibration.""" diff --git a/onnx_ptq/requirements.txt b/examples/onnx_ptq/requirements.txt similarity index 77% rename from onnx_ptq/requirements.txt rename to examples/onnx_ptq/requirements.txt index 9285b5f2bc..22b1e83cbe 100644 --- a/onnx_ptq/requirements.txt +++ b/examples/onnx_ptq/requirements.txt @@ -1,4 +1,3 @@ -cupy-cuda12x datasets>=2.14.4 optimum sentencepiece diff --git a/pruning/README.md b/examples/pruning/README.md similarity index 100% rename from pruning/README.md rename to examples/pruning/README.md diff --git a/pruning/cifar_resnet.ipynb b/examples/pruning/cifar_resnet.ipynb similarity index 100% rename from pruning/cifar_resnet.ipynb rename to examples/pruning/cifar_resnet.ipynb diff --git a/speculative_decoding/README.md b/examples/speculative_decoding/README.md similarity index 93% rename from speculative_decoding/README.md rename to examples/speculative_decoding/README.md index 0a30593915..3d4a94f8be 100644 --- a/speculative_decoding/README.md +++ b/examples/speculative_decoding/README.md @@ -44,11 +44,13 @@ if mode == "medusa": config = { "medusa_num_heads": 2, "medusa_num_layers": 1, - } + } elif mode == "eagle": config = { - "eagle_num_layers": 1 - } + "eagle_num_layers": 1, + "use_input_layernorm_in_first_layer": True, + "use_last_layernorm": False + } mtsp.convert(model, [(mode, config)]) tokenizer = transformers.AutoTokenizer.from_pretrained(ckpt_path) @@ -82,17 +84,19 @@ Next, we fine-tune the speculative decoding models with the base model frozen. H ```sh ./launch.sh --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \ - --data finetuning/data.jsonl \ + --data finetune/data.jsonl \ --mode medusa \ --num_epochs 1 --lr 1e-5 --save_steps 1000 \ --output_dir medusa-tinyllama \ + --fsdp_transformer_layer_cls_to_wrap LlamaDecoderLayer \ --medusa_num_heads 2 --medusa_num_layers 1 ./launch.sh --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \ - --data finetuning/data.jsonl \ + --data finetune/data.jsonl \ --mode eagle \ --num_epochs 1 --lr 1e-5 --save_steps 1000 \ --output_dir eagle-tinyllama \ + --fsdp_transformer_layer_cls_to_wrap LlamaDecoderLayer \ --eagle_num_layers 1 ``` diff --git a/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py similarity index 84% rename from speculative_decoding/eagle_utils.py rename to examples/speculative_decoding/eagle_utils.py index c804633630..0df8045d02 100644 --- a/speculative_decoding/eagle_utils.py +++ b/examples/speculative_decoding/eagle_utils.py @@ -1,26 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 Any, Dict, List +from typing import Any import torch import transformers @@ -116,7 +110,7 @@ class SupervisedDataset(Dataset): def __len__(self): return len(self.input_ids) - def __getitem__(self, i) -> Dict[str, torch.Tensor]: + def __getitem__(self, i) -> dict[str, torch.Tensor]: return dict( input_ids=self.input_ids[i], labels=self.labels[i], @@ -147,7 +141,7 @@ class LazySupervisedDataset(Dataset): def __len__(self): return len(self.raw_data) - def __getitem__(self, i) -> Dict[str, torch.Tensor]: + def __getitem__(self, i) -> dict[str, torch.Tensor]: if i in self.cached_data_dict: return self.cached_data_dict[i] @@ -165,7 +159,7 @@ class LazySupervisedDataset(Dataset): def make_eagle_supervised_data_module( tokenizer: transformers.PreTrainedTokenizer, data_args -) -> Dict: +) -> dict: """Make dataset and collator for supervised fine-tuning. Args: @@ -203,7 +197,7 @@ class DataCollatorWithPadding: outtensors = torch.cat((intensors, padding_tensor)) return outtensors - def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]: + def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: max_length = max(item["input_ids"].shape[0] for item in features) batch_input_ids = torch.stack( [self.paddingtensor(item["input_ids"], max_length) for item in features] diff --git a/speculative_decoding/example.ipynb b/examples/speculative_decoding/example.ipynb similarity index 82% rename from speculative_decoding/example.ipynb rename to examples/speculative_decoding/example.ipynb index d743779136..2aa10501b1 100644 --- a/speculative_decoding/example.ipynb +++ b/examples/speculative_decoding/example.ipynb @@ -55,7 +55,7 @@ "metadata": {}, "outputs": [], "source": [ - "!vllm serve meta-llama/Llama-3.2-1B-Instruct --api-key token-abc123 --port 8000 --tensor-parallel-size 1 --quantization=modelopt" + "!vllm serve /tmp/llama3.2_1B_fp8 --api-key token-abc123 --port 8000 --tensor-parallel-size 1 --quantization=modelopt" ] }, { @@ -80,7 +80,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Let's load the base model and convert it to Medusa Model" + "Let's load the base model and convert it to EAGLE Model" ] }, { @@ -98,10 +98,9 @@ "\n", "model = transformers.AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-3.2-1B-Instruct\")\n", "config = {\n", - " \"medusa_num_heads\": 3,\n", - " \"medusa_num_layers\": 1,\n", + " \"eagle_num_layers\": 1,\n", "}\n", - "mtsp.convert(model, [(\"medusa\", config)])\n", + "mtsp.convert(model, [(\"eagle\", config)])\n", "\n", "tokenizer = transformers.AutoTokenizer.from_pretrained(\"meta-llama/Llama-3.2-1B-Instruct\")\n", "tokenizer.pad_token_id = tokenizer.eos_token_id" @@ -111,7 +110,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Once synthesized data is ready, we can start training the medusa heads." + "Once synthesized data is ready, we can start training the eagle model." ] }, { @@ -124,7 +123,7 @@ "from dataclasses import dataclass, field\n", "from typing import Optional\n", "\n", - "from speculative_decoding.medusa_utils import LazySupervisedDataset\n", + "from speculative_decoding.eagle_utils import LazySupervisedDataset\n", "from transformers import Trainer\n", "\n", "with open(\"/tmp/finetune/data.jsonl\", \"r\") as f:\n", @@ -149,7 +148,7 @@ "\n", "\n", "training_args = TrainingArguments(\n", - " output_dir=\"/tmp/medusa_bf16\",\n", + " output_dir=\"/tmp/eagle_bf16\",\n", " num_train_epochs=1.0,\n", " per_device_train_batch_size=1,\n", " per_device_eval_batch_size=1,\n", @@ -181,7 +180,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Now we have a Medusa model in BF16 format. Next, we quantize this model into FP8 (PTQ)." + "Now we have a EAGLE model in BF16 format. Next, we quantize this model into FP8 (PTQ)." ] }, { @@ -195,8 +194,8 @@ "\n", "mto.enable_huggingface_checkpointing()\n", "\n", - "model = transformers.AutoModelForCausalLM.from_pretrained(\"/tmp/medusa_bf16\")\n", - "tokenizer = transformers.AutoTokenizer.from_pretrained(\"/tmp/medusa_bf16\")\n", + "model = transformers.AutoModelForCausalLM.from_pretrained(\"/tmp/eagle_bf16\")\n", + "tokenizer = transformers.AutoTokenizer.from_pretrained(\"/tmp/eagle_bf16\")\n", "\n", "calib_dataloader = dataset_utils.get_dataset_dataloader(\n", " dataset_name=\"cnn_dailymail\",\n", @@ -218,7 +217,7 @@ "model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)\n", "mtq.print_quant_summary(model)\n", "\n", - "model.save_pretrained(\"/tmp/medusa_fp8_ptq\")" + "model.save_pretrained(\"/tmp/eagle_fp8_ptq\")" ] }, { @@ -234,7 +233,7 @@ "metadata": {}, "outputs": [], "source": [ - "training_args.output_dir = \"/tmp/medusa_fp8_qat\"\n", + "training_args.output_dir = \"/tmp/eagle_fp8_qat\"\n", "trainer = Trainer(\n", " model=model,\n", " tokenizer=tokenizer,\n", @@ -262,7 +261,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "To deploy this model, we need to first export it to a TensorRT-LLM checkpoint." + "To deploy this model, we need to first export it to a Unified checkpoint." ] }, { @@ -271,22 +270,16 @@ "metadata": {}, "outputs": [], "source": [ - "import torch\n", "from accelerate.hooks import remove_hook_from_module\n", "\n", - "from modelopt.torch.export import export_tensorrt_llm_checkpoint\n", + "from modelopt.torch.export import export_hf_checkpoint\n", "\n", "# Move meta tensor back to device before exporting.\n", "remove_hook_from_module(model, recurse=True)\n", "\n", - "export_tensorrt_llm_checkpoint(\n", + "export_hf_checkpoint(\n", " model,\n", - " \"llama\",\n", - " torch.float16,\n", - " export_dir=\"/tmp/trtllm_ckpt\",\n", - " inference_tensor_parallel=1,\n", - " inference_pipeline_parallel=1,\n", - " naive_fp8_quantization=False,\n", + " export_dir=\"/tmp/hf_ckpt\",\n", ")" ] }, @@ -294,7 +287,23 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Then build a TensorRT-LLM engine." + "Then convert the Unified ckeckpoint to TRTLLM checkpoint." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!python TensorRT-LLM/examples/eagle/convert_checkpoint.py --model_dir /tmp/hf_ckpt --output_dir /tmp/trtllm_ckpt --num_eagle_layers 5 --max_non_leaves_per_layer 4 --max_draft_len 25" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Last, build a TensorRT-LLM engine." ] }, { @@ -310,12 +319,12 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "To run the Medusa engine, please refer to https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/medusa\n", - "Inside TensorRT-LLM/examples/medusa:\n", + "To run the EAGLE engine, please refer to https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/eagle\n", + "Inside TensorRT-LLM/examples/eagle:\n", "python ../run.py --engine_dir ./trtllm_engine \\\n", - " --tokenizer_dir ./medusa_fp8_qat \\\n", + " --tokenizer_dir ./eagle_fp8_qat \\\n", " --max_output_len=100 \\\n", - " --medusa_choices=\"[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1]]\" \\\n", + " --eagle_choices=\"[[0],[1],[2],[3],[0,0],[0,1],[0,2],[1,0],[1,1],[2,0],[2,1],[3,0],[0,0,0],[0,0,1],[0,0,2],[0,1,0],[0,1,1],[0,2,0],[0,2,1],[1,0,0],[0,0,0,0],[0,0,0,1],[0,0,0,2],[0,0,0,0,0],[0,0,0,0,1]]\" \\\n", " --temperature 1.0 \\\n", " --input_text \"Once upon\"" ] diff --git a/speculative_decoding/launch.sh b/examples/speculative_decoding/launch.sh similarity index 84% rename from speculative_decoding/launch.sh rename to examples/speculative_decoding/launch.sh index 23952c0926..1b55769d18 100755 --- a/speculative_decoding/launch.sh +++ b/examples/speculative_decoding/launch.sh @@ -1,4 +1,19 @@ #!/bin/bash +# 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. + set -e set -x set -o pipefail diff --git a/speculative_decoding/main.py b/examples/speculative_decoding/main.py similarity index 84% rename from speculative_decoding/main.py rename to examples/speculative_decoding/main.py index 4b459ecdf7..1bf8cf86f9 100644 --- a/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -15,25 +15,19 @@ # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 dataclasses import dataclass, field @@ -93,6 +87,8 @@ class MedusaArguments: @dataclass class EagleArguments: eagle_num_layers: Optional[int] = field(default=1) + use_input_layernorm_in_first_layer: Optional[bool] = field(default=True) + use_last_layernorm: Optional[bool] = field(default=False) @dataclass @@ -156,7 +152,11 @@ def train(): } mtsp.convert(model, [("medusa", config)]) elif training_args.mode == "eagle": - config = {"eagle_num_layers": eagle_args.eagle_num_layers} + config = { + "eagle_num_layers": eagle_args.eagle_num_layers, + "use_input_layernorm_in_first_layer": eagle_args.use_input_layernorm_in_first_layer, + "use_last_layernorm": eagle_args.use_last_layernorm, + } mtsp.convert(model, [("eagle", config)]) elif training_args.mode == "redrafter": config = { diff --git a/speculative_decoding/medusa_utils.py b/examples/speculative_decoding/medusa_utils.py similarity index 84% rename from speculative_decoding/medusa_utils.py rename to examples/speculative_decoding/medusa_utils.py index 8f3986816c..b3aac66718 100644 --- a/speculative_decoding/medusa_utils.py +++ b/examples/speculative_decoding/medusa_utils.py @@ -16,28 +16,21 @@ # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 Dict import torch import transformers @@ -60,7 +53,7 @@ def change_format(conversations): def preprocess( sources, tokenizer: transformers.PreTrainedTokenizer, -) -> Dict: +) -> dict: """ Preprocesses conversation data and tokenizes it for model input. @@ -146,7 +139,7 @@ class SupervisedDataset(Dataset): def __len__(self): return len(self.input_ids) - def __getitem__(self, i) -> Dict[str, torch.Tensor]: + def __getitem__(self, i) -> dict[str, torch.Tensor]: return dict( input_ids=self.input_ids[i], labels=self.labels[i], @@ -176,7 +169,7 @@ class LazySupervisedDataset(Dataset): def __len__(self): return len(self.raw_data) - def __getitem__(self, i) -> Dict[str, torch.Tensor]: + def __getitem__(self, i) -> dict[str, torch.Tensor]: if i in self.cached_data_dict: return self.cached_data_dict[i] @@ -193,7 +186,7 @@ class LazySupervisedDataset(Dataset): def make_medusa_supervised_data_module( tokenizer: transformers.PreTrainedTokenizer, data_args -) -> Dict: +) -> dict: """Make dataset and collator for supervised fine-tuning. Args: diff --git a/examples/speculative_decoding/prepare_data.sh b/examples/speculative_decoding/prepare_data.sh new file mode 100644 index 0000000000..4a37e0dd8b --- /dev/null +++ b/examples/speculative_decoding/prepare_data.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# 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. + +set -e +set -x +set -o pipefail + +while [ $# -gt 0 ]; do + case "$1" in + --data*) + if [[ "$1" != *=* ]]; then shift; fi + DATA="${1#*=}" + ;; + --output_path*) + if [[ "$1" != *=* ]]; then shift; fi + OUTPUT_PATH="${1#*=}" + ;; + --max_token*) + if [[ "$1" != *=* ]]; then shift; fi + MAX_TOKEN="${1#*=}" + ;; + *) + >&2 printf "Error: Invalid argument\n" + exit 1 + ;; + esac + shift +done + +OUTPUT_DIR="$(dirname "${OUTPUT_PATH}")" +mkdir -p $OUTPUT_DIR + +CMD="python vllm_generate.py --data_path $DATA --output_path $OUTPUT_PATH --num_threads 8 --max_tokens $MAX_TOKEN --temperature 0.0 --chat" + +sh -c "$CMD" diff --git a/speculative_decoding/requirements.txt b/examples/speculative_decoding/requirements.txt similarity index 80% rename from speculative_decoding/requirements.txt rename to examples/speculative_decoding/requirements.txt index c4f83216b8..245f1c9468 100644 --- a/speculative_decoding/requirements.txt +++ b/examples/speculative_decoding/requirements.txt @@ -1,4 +1,3 @@ -datasets>=2.16.1 flash-attn openai py7zr diff --git a/speculative_decoding/vllm_generate.py b/examples/speculative_decoding/vllm_generate.py similarity index 82% rename from speculative_decoding/vllm_generate.py rename to examples/speculative_decoding/vllm_generate.py index 64259078f4..57c23ec908 100644 --- a/speculative_decoding/vllm_generate.py +++ b/examples/speculative_decoding/vllm_generate.py @@ -13,25 +13,19 @@ # limitations under the License. # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 concurrent.futures @@ -69,7 +63,7 @@ parser = argparse.ArgumentParser() parser.add_argument("--data_path", type=str) parser.add_argument("--output_path", type=str) parser.add_argument("--num_threads", type=int, default=256) -parser.add_argument("--temperature", type=float, default=0.3) +parser.add_argument("--temperature", type=float, default=0.0) parser.add_argument("--max_tokens", type=int, default=2048) parser.add_argument("--chat", action="store_true") args = parser.parse_args() diff --git a/examples/vlm_eval/README.md b/examples/vlm_eval/README.md new file mode 100644 index 0000000000..9e393ea220 --- /dev/null +++ b/examples/vlm_eval/README.md @@ -0,0 +1,36 @@ +# Evaluation scripts for VLM tasks + +This folder includes popular 3rd-party VLM benchmarks for VLM accuracy evaluation. + +The following instructions show how to evaluate the VLM (including Model Optimizer quantized LLM) with the benchmarks, including the TensorRT-LLM deployment. + +## GQA + +[GQA: a dataset for real-world visual reasoning and compositional question answering](https://arxiv.org/abs/1902.09506). Upon completing the benchmark, the model's accuracy (in percentage format) will be displayed, providing a clear metric for performance evaluation. + +First log in to Hugging Face account with your token. + +```bash +huggingface-cli login +``` + +### Baseline + +```bash +bash gqa.sh --hf_model +``` + +### Quantized (simulated) + +```bash +# MODELOPT_QUANT_CFG: Choose from [INT8_SMOOTHQUANT_CFG|FP8_DEFAULT_CFG|INT4_AWQ_CFG|W4A8_AWQ_BETA_CFG] +bash gqa.sh --hf_model --quant_cfg MODELOPT_QUANT_CFG +``` + +### Evaluate the TensorRT-LLM engine + +TensorRT engine could be built following this [guide](../vlm_ptq/README.md) + +```bash +bash gqa.sh --hf_model --visual_engine --llm_engine +``` diff --git a/examples/vlm_eval/convert_gqa_for_eval.py b/examples/vlm_eval/convert_gqa_for_eval.py new file mode 100644 index 0000000000..de4110c83d --- /dev/null +++ b/examples/vlm_eval/convert_gqa_for_eval.py @@ -0,0 +1,56 @@ +# 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. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + +# Adopted from https://github.com/NVlabs/VILA/blob/ec7fb2c264920bf004fd9fa37f1ec36ea0942db5/scripts/convert_gqa_for_eval.py + + +import argparse +import json + +parser = argparse.ArgumentParser() +parser.add_argument("--src", type=str) +parser.add_argument("--dst", type=str) +args = parser.parse_args() + +all_answers = [] +for line_idx, line in enumerate(open(args.src)): + res = json.loads(line) + question_id = res["question_id"] + text = res["text"].rstrip(".").lower() + all_answers.append({"questionId": question_id, "prediction": text}) + +with open(args.dst, "w") as f: + json.dump(all_answers, f) diff --git a/examples/vlm_eval/gqa.sh b/examples/vlm_eval/gqa.sh new file mode 100644 index 0000000000..7b8c9dc48b --- /dev/null +++ b/examples/vlm_eval/gqa.sh @@ -0,0 +1,142 @@ +#!/bin/bash +# 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. + + +set -e + +# Download dataset +script_dir="$(dirname "$(readlink -f "$0")")" + +EVAL_FILE="$script_dir/eval.py" +if [ ! -f $EVAL_FILE ]; then + echo "$EVAL_FILE does not exist. Downloading this file from https://nlp.stanford.edu/data/gqa/eval.zip." + wget https://nlp.stanford.edu/data/gqa/eval.zip + unzip eval.zip "eval.py" -d . + rm eval.zip + + # Changes to eval.py due to the missing assets in GQA v1.2 release + sed -i '77s/{tier}_all_questions.json/{tier}_questions.json/' "$EVAL_FILE" + sed -i '119,120s/^/# /' "$EVAL_FILE" + sed -i '126,128s/^/# /' "$EVAL_FILE" + sed -i '367,373s/^/# /' "$EVAL_FILE" + sed -i '376,379s/^/# /' "$EVAL_FILE" + sed -i '388s/^/# /' "$EVAL_FILE" +fi + +gqa_data=$script_dir/gqa/data +QUESTION=$gqa_data/testdev_balanced_questions.json +if [ ! -f $QUESTION ]; then + echo "$QUESTION does not exist. Downloading this file from https://downloads.cs.stanford.edu/nlp/data/gqa/questions1.2.zip." + wget -P $gqa_data https://downloads.cs.stanford.edu/nlp/data/gqa/questions1.2.zip + unzip $gqa_data/questions1.2.zip "testdev_balanced_questions.json" -d $gqa_data + rm $gqa_data/questions1.2.zip +fi + +# Parse command-line arguments +while [[ $# -gt 0 ]]; do + case "$1" in + --hf_model) + HF_MODEL_DIR="$2" + shift 2 + ;; + --visual_engine) + VISUAL_ENGINE_DIR="$2" + shift 2 + ;; + --llm_engine) + LLM_ENGINE_DIR="$2" + shift 2 + ;; + --batch_size) + BATCH_SIZE="$2" + shift 2 + ;; + --quant_cfg) + QUANT_CFG="$2" + shift 2 + ;; + *) + echo "Unknown option $1" + exit 1 + ;; + esac +done + +# Verify required arguments are set +if [ -z "$HF_MODEL_DIR" ]; then + echo "Error: Missing required argument --hf_model" + exit 1 +fi + +MODEL_NAME=$(basename $HF_MODEL_DIR | sed 's/[^0-9a-zA-Z\-]/_/g' | tr 'A-Z' 'a-z') + +if [[ "$MODEL_NAME" == *"vila"* ]] && [[ -z "$LLM_ENGINE_DIR" ]]; then + # Install required dependency for VILA + pip install -r requirements-vila.txt + # Clone oringinal VILA repo + if [ ! -d "$(dirname "$HF_MODEL_DIR")/VILA" ]; then + echo "VILA repository is needed until it is added to HF model zoo. Cloning the repository parallel to $HF_MODEL_DIR..." + git clone https://github.com/Efficient-Large-Model/VILA.git "$(dirname "$HF_MODEL_DIR")/VILA" && \ + cd "$(dirname "$HF_MODEL_DIR")/VILA" && \ + git checkout ec7fb2c264920bf004fd9fa37f1ec36ea0942db5 && \ + cd - + fi +fi + +# Set batch size defaulted to 20 for VILA and Llava +if [[ -z "$BATCH_SIZE" && ("$MODEL_NAME" == *"vila"* || "$MODEL_NAME" == *"llava"*) ]]; then + BATCH_SIZE=20 +fi + +# Check if TRT engine is provided +if [ -z "$VISUAL_ENGINE_DIR" ] || [ -z "$LLM_ENGINE_DIR" ]; then + echo "Either --visual_engine or --llm_engine not provided, evalluation will be based on Pytorch." + if [ -z "QUANT_CFG" ]; then + ANSWER_DIR="$script_dir/gqa/$MODEL_NAME/llava_gqa_testdev_balanced/answers" + ANSWERS_FILE="$ANSWER_DIR/merge.jsonl" + else + ANSWER_DIR="$script_dir/gqa/$MODEL_NAME_$QUANT_CFG/llava_gqa_testdev_balanced/answers" + ANSWERS_FILE="$ANSWER_DIR/merge.jsonl" + fi +else + echo "Both --visual_engine or --llm_engine are provided, evalluation will be based on TRT engine." + ANSWER_DIR="$script_dir/gqa/$(basename $LLM_ENGINE_DIR)/llava_gqa_testdev_balanced/answers" + ANSWERS_FILE="$ANSWER_DIR/merge.jsonl" +fi + +# Run the Python script with the parsed arguments +if [ ! -f $ANSWERS_FILE ]; then + python model_gqa_loader.py \ + --answers_file "$ANSWERS_FILE" \ + --hf_model_dir "$HF_MODEL_DIR" \ + ${VISUAL_ENGINE_DIR:+--visual_engine_dir "$VISUAL_ENGINE_DIR"} \ + ${LLM_ENGINE_DIR:+--llm_engine_dir "$LLM_ENGINE_DIR"} \ + ${BATCH_SIZE:+--batch_size "$BATCH_SIZE"} \ + ${QUANT_CFG:+--quant_cfg "$QUANT_CFG"} +fi + +# Convert answer to prediction for evaluation +PREDICTION_FILE="$ANSWER_DIR/testdev_balanced_predictions.json" +if [ ! -f $PREDICTION_FILE ]; then + python convert_gqa_for_eval.py \ + --src $ANSWERS_FILE \ + --dst $PREDICTION_FILE +fi + +# Get evaluation result +python eval.py \ + --tier "$gqa_data/testdev_balanced" \ + --predictions $PREDICTION_FILE diff --git a/examples/vlm_eval/model_gqa_loader.py b/examples/vlm_eval/model_gqa_loader.py new file mode 100644 index 0000000000..3a75792d23 --- /dev/null +++ b/examples/vlm_eval/model_gqa_loader.py @@ -0,0 +1,263 @@ +# 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. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + + +import copy +import os +import sys +import time +from pathlib import Path + +from datasets import load_dataset +from tensorrt_llm import logger +from tensorrt_llm.runtime import MultimodalModelRunner +from tqdm import tqdm +from transformers import AutoProcessor + +import modelopt.torch.quantization as mtq + +sys.path.append(str(Path(__file__).resolve().parent / "../llm_ptq")) +from example_utils import get_tokenizer +from vlm_eval_utils import save_jsonl + +sys.path.append(str(Path(__file__).resolve().parent / "../vlm_ptq")) +from utils import parse_arguments + +from modelopt.torch.utils.dataset_utils import ( + create_forward_loop, + get_dataset_dataloader, + get_max_batch_size, +) + + +def quantize_model(model, args, tokenizer): + sample_memory_usage_ratio = ( + 2 if "AWQ" in args.quant_cfg or "SMOOTHQUANT" in args.quant_cfg else 1.1 + ) + batch_size = get_max_batch_size(model, sample_memory_usage_ratio=sample_memory_usage_ratio) + calib_size = args.calib_size + if batch_size > calib_size: + batch_size = calib_size + calib_dataloader = get_dataset_dataloader( + dataset_name="cnn_dailymail", + tokenizer=tokenizer, + batch_size=batch_size, + num_samples=calib_size, + device=model.device, + ) + calibrate_loop = create_forward_loop(dataloader=calib_dataloader) + + quant_cfg = getattr(mtq, args.quant_cfg) + if "AWQ" in args.quant_cfg: + quant_cfg = copy.deepcopy(getattr(mtq, args.quant_cfg)) + weight_quantizer = quant_cfg["quant_cfg"]["*weight_quantizer"] + if isinstance(weight_quantizer, list): + weight_quantizer = weight_quantizer[0] + enable_quant_kv_cache = args.quant_cfg not in ["INT8_SMOOTHQUANT_CFG"] + print(f'{"Enable" if enable_quant_kv_cache else "Disable"} KV cache quantization') + quant_cfg["quant_cfg"]["*output_quantizer"] = { + "num_bits": 8 if args.quant_cfg == "INT8_SMOOTHQUANT_CFG" else (4, 3), + "axis": None, + "enable": enable_quant_kv_cache, + } + + print("Starting quantization...") + start_time = time.time() + model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop) + end_time = time.time() + print(f"Quantization done. Total time used: {end_time - start_time}s") + return model + + +def main(): + args = parse_arguments("eval") + + # Load data + instances = load_dataset( + "lmms-lab/GQA", "testdev_balanced_instructions", split="testdev", token=True + ) + images = load_dataset("lmms-lab/GQA", "testdev_balanced_images", split="testdev", token=True) + id2image = {} + for row in images: + id2image[row["id"]] = row["image"].convert("RGB") + + # Load model + if args.llm_engine_dir is not None: + os.environ["TOKENIZERS_PARALLELISM"] = "false" + logger.set_level(args.log_level) + # Load TensorRT engine + model = MultimodalModelRunner(args) + # Run batch inference + outputs = [] + batch_size = args.batch_size + if model.model_type in ["phi-3-vision"]: + # Phi-3-vision doesn't support batch inference for now + batch_size = 1 + for index in tqdm(range(0, len(instances), batch_size)): + batch = instances[index : index + batch_size] + raw_images = [id2image[imageId] for imageId in batch["imageId"]] + questions = batch["question"] + questions = [ + q + "\nAnswer the question using a single word or phrase." for q in questions + ] + if model.model_type in ["llava"]: + input_text = ["\n" + question for question in questions] + elif model.model_type in ["vila"]: + input_text = ["\n" + question for question in questions] + elif model.model_type in ["phi-3-vision"]: + input_text = questions[0] + _, output_text = model.run(input_text, raw_images, args.max_new_tokens) + outputs.extend( + [ + { + "question_id": id, + "prompt": batch["question"][index], + "text": output_text[index][0], + } + for index, id in enumerate(batch["id"]) + ] + ) + + else: + # Load HF model + if "vila" in args.hf_model_dir.lower(): + sys.path.append(os.path.join(args.hf_model_dir, "..", "VILA")) + import llava + + model = llava.load(args.hf_model_dir) + from llava import conversation as conversation_lib + + if "8b" in args.hf_model_dir.lower(): + conv_mode = "llama_3" + elif "40b" in args.hf_model_dir.lower(): + conv_mode = "hermes-2" + else: + conv_mode = "vicuna_v1" + + conversation_lib.default_conversation = conversation_lib.conv_templates[ + conv_mode + ].copy() + from transformers import GenerationConfig + + generation_config = GenerationConfig.from_pretrained(args.hf_model_dir + "/llm") + generation_config.update(**{"max_new_tokens": args.max_new_tokens}) + else: + processor = AutoProcessor.from_pretrained( + args.hf_model_dir, trust_remote_code=args.trust_remote_code + ) + if "llava" in args.hf_model_dir.lower(): + from transformers import LlavaForConditionalGeneration + + model = LlavaForConditionalGeneration.from_pretrained( + args.hf_model_dir, device_map="auto", torch_dtype="auto" + ) + # To be deprecated for new version transformers + processor.patch_size = model.config.vision_config.patch_size + processor.vision_feature_select_strategy = ( + model.config.vision_feature_select_strategy + ) + elif "phi" in args.hf_model_dir.lower(): + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_pretrained( + args.hf_model_dir, + device_map="auto", + trust_remote_code=args.trust_remote_code, + torch_dtype="auto", + _attn_implementation="flash_attention_2", + ) + # Evaluation for simulated quantization + if args.quant_cfg: + tokenizer = get_tokenizer(args.hf_model_dir, trust_remote_code=args.trust_remote_code) + if "vila" in args.hf_model_dir.lower(): + model.llm = quantize_model(model.llm, args, tokenizer) + elif "llava" in args.hf_model_dir.lower(): + model.language_model = quantize_model(model.language_model, args, tokenizer) + elif "phi" in args.hf_model_dir.lower(): + model = quantize_model(model, args, tokenizer) + + outputs = [] + for instance in tqdm(instances): + image = id2image[instance["imageId"]] + question = instance["question"] + if "vila" in args.hf_model_dir.lower(): + response = model.generate_content( + [image, question], generation_config=generation_config + ) + else: + if "llava" in args.hf_model_dir.lower(): + conversation = [ + { + "role": "user", + "content": [ + {"type": "text", "text": question}, + {"type": "image"}, + ], + }, + ] + prompt = processor.apply_chat_template(conversation, add_generation_prompt=True) + inputs = processor(images=image, text=prompt, return_tensors="pt").to( + "cuda:0", model.dtype + ) + response = model.generate( + **inputs, max_new_tokens=args.max_new_tokens, do_sample=False + ) + elif "phi" in args.hf_model_dir.lower(): + conversation = [ + {"role": "user", "content": f"<|image_1|>\n{question}"}, + ] + prompt = processor.tokenizer.apply_chat_template( + conversation, tokenize=False, add_generation_prompt=True + ) + inputs = processor(images=image, text=prompt, return_tensors="pt").to( + "cuda:0", model.dtype + ) + response = model.generate( + **inputs, + eos_token_id=processor.tokenizer.eos_token_id, + max_new_tokens=args.max_new_tokens, + do_sample=False, + ) + response = processor.decode( + response[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True + ) + + outputs.append({"question_id": instance["id"], "prompt": question, "text": response}) + save_jsonl(args.answers_file, outputs) + + +if __name__ == "__main__": + main() diff --git a/vlm_ptq/requirements-vila.txt b/examples/vlm_eval/requirements-vila.txt similarity index 73% rename from vlm_ptq/requirements-vila.txt rename to examples/vlm_eval/requirements-vila.txt index f114878e0d..d88b118665 100644 --- a/vlm_ptq/requirements-vila.txt +++ b/examples/vlm_eval/requirements-vila.txt @@ -1,2 +1,2 @@ -deepspeed==0.15.4 +deepspeed>=0.16.0 git+https://github.com/bfshi/scaling_on_scales.git diff --git a/examples/vlm_eval/vlm_eval_utils.py b/examples/vlm_eval/vlm_eval_utils.py new file mode 100644 index 0000000000..e06ff5765d --- /dev/null +++ b/examples/vlm_eval/vlm_eval_utils.py @@ -0,0 +1,64 @@ +# 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. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + + +import json +import os +from contextlib import contextmanager +from io import TextIOWrapper +from typing import IO, Any, Iterator, TextIO, Union + + +@contextmanager +def file_descriptor(f: Union[str, IO], mode: str = "r") -> Iterator[IO]: + opened = False + try: + if isinstance(f, str): + f = open(f, mode) + opened = True + yield f + finally: + if opened: + assert isinstance(f, TextIOWrapper), type(f) + f.close() + + +def save_jsonl(f: Union[str, TextIO], obj: Any, **kwargs) -> None: + assert isinstance(f, str), type(f) + os.makedirs(os.path.dirname(f), exist_ok=True) + + with file_descriptor(f, mode="w") as fd: + fd.write("\n".join(json.dumps(datum, **kwargs) for datum in obj)) diff --git a/vlm_ptq/.gitignore b/examples/vlm_ptq/.gitignore similarity index 100% rename from vlm_ptq/.gitignore rename to examples/vlm_ptq/.gitignore diff --git a/vlm_ptq/README.md b/examples/vlm_ptq/README.md similarity index 89% rename from vlm_ptq/README.md rename to examples/vlm_ptq/README.md index 771284e2a9..92e24ab4d9 100644 --- a/vlm_ptq/README.md +++ b/examples/vlm_ptq/README.md @@ -39,6 +39,8 @@ git clone https://huggingface.co/microsoft/Phi-3-vision-128k-instruct scripts/huggingface_example.sh --type phi --model Phi-3-vision-128k-instruct --quant [fp8|int8_sq|int4_awq|w4a8_awq] ``` +The example scripts above also have an additional flag `--tasks gqa`, which will trigger evaluation of the built TensorRT engine using GQA benchmark. Details of the evaluation is explained in this [tutorial](../vlm_eval/README.md). + ### Model Support List Model | type | fp8 | int8_sq | int4_awq | w4a8_awq1 @@ -47,6 +49,6 @@ Llava | llava | Yes | Yes | Yes | Yes VILA | vila | Yes | Yes | Yes | Yes Phi-3-vision | phi | Yes | Yes | Yes | Yes -> *1.The w4a8_awq is an experimental quantization scheme that may result in a higher accuracy penalty. Only available on sm90 GPUs* +> *1.The w4a8_awq is an experimental quantization scheme that may result in a higher accuracy penalty.* > *The accuracy loss after PTQ may vary depending on the actual model and the quantization method. Different models may have different accuracy loss and usually the accuracy loss is more significant when the base model is small. If the accuracy after PTQ is not meeting the requirement, please try either modifying [hf_ptq.py](./hf_ptq.py) and disabling the KV cache quantization or using the [QAT](./../llm_qat/README.md) instead.* diff --git a/vlm_ptq/requirements-mllama.txt b/examples/vlm_ptq/requirements-mllama.txt similarity index 100% rename from vlm_ptq/requirements-mllama.txt rename to examples/vlm_ptq/requirements-mllama.txt diff --git a/examples/vlm_ptq/requirements-vila.txt b/examples/vlm_ptq/requirements-vila.txt new file mode 100644 index 0000000000..d88b118665 --- /dev/null +++ b/examples/vlm_ptq/requirements-vila.txt @@ -0,0 +1,2 @@ +deepspeed>=0.16.0 +git+https://github.com/bfshi/scaling_on_scales.git diff --git a/vlm_ptq/scripts/huggingface_example.sh b/examples/vlm_ptq/scripts/huggingface_example.sh similarity index 71% rename from vlm_ptq/scripts/huggingface_example.sh rename to examples/vlm_ptq/scripts/huggingface_example.sh index c52cbd9250..bf0d699894 100755 --- a/vlm_ptq/scripts/huggingface_example.sh +++ b/examples/vlm_ptq/scripts/huggingface_example.sh @@ -1,4 +1,19 @@ #!/bin/bash +# 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. + set -e script_dir="$(dirname "$(readlink -f "$0")")" @@ -69,12 +84,19 @@ esac GPU_NAME=$(nvidia-smi --id 0 --query-gpu=name --format=csv,noheader,nounits | sed 's/ /_/g') if [ "${MODEL_TYPE}" = "phi" ]; then - BUILD_MAX_OUTPUT_LEN=4608 + BUILD_MAX_INPUT_LEN=4096 else - BUILD_MAX_OUTPUT_LEN=2560 + BUILD_MAX_INPUT_LEN=1024 +fi + +BUILD_MAX_OUTPUT_LEN=512 + +if [ "$MODEL_TYPE" = "llava" ] || [ "$MODEL_TYPE" = "vila" ]; then + BUILD_MAX_BATCH_SIZE=20 +else + BUILD_MAX_BATCH_SIZE=4 fi -BUILD_MAX_BATCH_SIZE=4 echo "Using the following config: max input $BUILD_MAX_INPUT_LEN max output $BUILD_MAX_OUTPUT_LEN max batch $BUILD_MAX_BATCH_SIZE" @@ -107,15 +129,29 @@ if [ $QFORMAT == "fp8_naive" ]; then fi if [ -n "$AUTO_QUANTIZE_BITS" ]; then - PTQ_ARGS+=" --auto_quantize_bits $AUTO_QUANTIZE_BITS" + PTQ_ARGS+=" --auto_quantize_bits $AUTO_QUANTIZE_BITS " +fi + +if $TRUST_REMOTE_CODE; then + PTQ_ARGS+=" --trust_remote_code " +fi + +if [ "${MODEL_TYPE}" = "mllama" ]; then + PTQ_ARGS+=" --dataset scienceqa " fi case "${MODEL_TYPE}" in - "vila" | "phi") - VLM_ARGS=" --max_multimodal_len=$((BUILD_MAX_BATCH_SIZE * 4096)) " + "vila") + VISUAL_FEATURE=196 + VLM_ARGS=" --max_multimodal_len=$((BUILD_MAX_BATCH_SIZE * VISUAL_FEATURE)) " + ;; + "phi") + VISUAL_FEATURE=4096 + VLM_ARGS=" --max_multimodal_len=$((BUILD_MAX_BATCH_SIZE * VISUAL_FEATURE)) " ;; "llava") - VLM_ARGS=" --max_multimodal_len=$((BUILD_MAX_BATCH_SIZE * 576)) " + VISUAL_FEATURE=576 + VLM_ARGS=" --max_multimodal_len=$((BUILD_MAX_BATCH_SIZE * VISUAL_FEATURE)) " ;; "mllama") VLM_ARGS=" --max_encoder_input_len=4100 " @@ -133,9 +169,6 @@ if [ "${MODEL_TYPE}" = "vila" ]; then git checkout ec7fb2c264920bf004fd9fa37f1ec36ea0942db5 && \ cd "$script_dir/.." fi -elif [ "${MODEL_TYPE}" = "llava" ]; then - echo "LLAVA model needs transformers version 4.42.4." - pip install -r ../vlm_ptq/requirements-llava.txt elif [ "${MODEL_TYPE}" = "mllama" ]; then echo "Mllama3.2 model requires transformers version 4.46.2 or try the latest version." pip install -r ../vlm_ptq/requirements-mllama.txt @@ -185,7 +218,6 @@ if [[ $TASKS =~ "build" ]] || [[ ! -d "$ENGINE_DIR" ]] || [[ ! $(ls -A $ENGINE_D $VLM_ARGS fi -echo "Build visual engine" VISUAL_ARGS="" VISION_ENCODER_DIR=${SAVE_PATH}/vision_encoder @@ -199,28 +231,39 @@ case "${MODEL_TYPE}" in ;; esac -python vlm_visual_engine.py \ - --model_path $MODEL_PATH \ - --model_type $VISUAL_MODEL_TYPE \ - --output_dir $VISION_ENCODER_DIR \ - $VISUAL_ARGS + +VISUAL_MAX_BATCH_SIZE=$BUILD_MAX_BATCH_SIZE + +if [[ $TASKS =~ "build" ]] || [[ ! -d "$VISION_ENCODER_DIR" ]] || [[ ! $(ls -A $VISION_ENCODER_DIR) ]]; then + echo "Build visual engine" + python vlm_visual_engine.py \ + --model_path $MODEL_PATH \ + --model_type $VISUAL_MODEL_TYPE \ + --output_dir $VISION_ENCODER_DIR \ + --max_batch_size $VISUAL_MAX_BATCH_SIZE \ + $VISUAL_ARGS +fi echo "Run inference example" -VLM_RUN_ARGS="" - -if [ "$MODEL_TYPE" == "vila" ]; then - VLM_RUN_ARGS+=" --image_path https://github.com/NVlabs/VILA/blob/6b941da19e31ddfdfaa60160908ccf0978d96615/demo_images/av.png?raw=true" -fi - python vlm_run.py \ --hf_model_dir $MODEL_PATH \ --visual_engine_dir $VISION_ENCODER_DIR \ - --llm_engine_dir $ENGINE_DIR \ - $VLM_RUN_ARGS + --llm_engine_dir $ENGINE_DIR -if [ "${MODEL_TYPE}" = "llava" ]; then - echo "For Llava model, current transformers version is 4.42.4, higher version transformers may be needed for other model." +if [[ $TASKS =~ "gqa" ]]; then + echo "Evaluating the TensorRT engine of the quantized model using GQA benchmark." + pushd ../vlm_eval/ + if [[ "$MODEL_PATH" =~ ^/ ]]; then + # If MODEL_PATH is absolute path + source gqa.sh --hf_model $MODEL_PATH --llm_engine $ENGINE_DIR --visual_engine $VISION_ENCODER_DIR + else + # If MODEL_PATH is absolute path + script_parent_dir=$(dirname "$script_dir") + source gqa.sh --hf_model $script_parent_dir/$MODEL_PATH --llm_engine $ENGINE_DIR --visual_engine $VISION_ENCODER_DIR + fi + + popd fi popd diff --git a/examples/vlm_ptq/utils.py b/examples/vlm_ptq/utils.py new file mode 100644 index 0000000000..bf0e458374 --- /dev/null +++ b/examples/vlm_ptq/utils.py @@ -0,0 +1,137 @@ +# 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 argparse + + +def parse_arguments(occasion): + parser = argparse.ArgumentParser() + if occasion == "run": + max_new_tokens = 30 + elif occasion == "eval": + max_new_tokens = 128 + parser.add_argument( + "--max_new_tokens", + type=int, + default=max_new_tokens, + help="Maximum number of new tokens generated by LLM", + ) + parser.add_argument("--batch_size", type=int, default=1) + parser.add_argument("--log_level", type=str, default="info") + parser.add_argument( + "--visual_engine_dir", + type=str, + default=None, + help="Directory containing visual TRT engines", + ) + parser.add_argument( + "--visual_engine_name", type=str, default="model.engine", help="Name of visual TRT engine" + ) + parser.add_argument( + "--llm_engine_dir", type=str, default=None, help="Directory containing TRT-LLM engines" + ) + parser.add_argument( + "--hf_model_dir", type=str, default=None, help="Directory containing tokenizer" + ) + parser.add_argument( + "--input_text", type=str, default=None, help="Text prompt to LLM, separated by symbol" + ) + parser.add_argument("--num_beams", type=int, help="Use beam search if num_beams >1", default=1) + parser.add_argument("--top_k", type=int, default=1) + parser.add_argument("--top_p", type=float, default=0.0) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--repetition_penalty", type=float, default=1.0) + parser.add_argument( + "--run_profiling", action="store_true", help="Profile runtime over several iterations" + ) + parser.add_argument( + "--profiling_iterations", type=int, help="Number of iterations to run profiling", default=20 + ) + parser.add_argument( + "--check_accuracy", action="store_true", help="Check correctness of text output" + ) + parser.add_argument( + "--video_path", type=str, default=None, help="Path to your local video file" + ) + parser.add_argument( + "--image_path", + type=str, + default=None, + help="List of input image paths, separated by symbol", + ) + parser.add_argument("--path_sep", type=str, default=",", help="Path separator symbol") + parser.add_argument("--prompt_sep", type=str, default=",", help="Prompt separator symbol") + parser.add_argument( + "--enable_context_fmha_fp32_acc", + action="store_true", + default=None, + help="Enable FMHA runner FP32 accumulation.", + ) + parser.add_argument( + "--use_py_session", + default=False, + action="store_true", + help="Whether or not to use Python runtime session", + ) + parser.add_argument( + "--enable_chunked_context", + action="store_true", + help="Enables chunked context (only available with cpp session).", + ) + parser.add_argument( + "--kv_cache_free_gpu_memory_fraction", + default=0.2, + type=float, + help="Specify the free gpu memory fraction.", + ) + parser.add_argument( + "--cross_kv_cache_fraction", + default=0.5, + type=float, + help="Specify the kv cache fraction reserved for cross attention. Only " + "applicable for encoder-decoder models. By default 0.5 for self and 0.5 for cross.", + ) + parser.add_argument( + "--multi_block_mode", + type=lambda s: s.lower() + in ("yes", "true", "t", "1"), # custom boolean function to convert input string to boolean + default=True, + help="Distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel.", + ) + parser.add_argument( + "--trust_remote_code", + help="Set trust_remote_code for Huggingface models and tokenizers", + default=False, + action="store_true", + ) + if occasion == "eval": + parser.add_argument("--answers_file", type=str, required=True) + parser.add_argument( + "--quant_cfg", + type=str, + default=None, + help="Specify the modelopt quantization configuration for simulated evaluation", + choices=[ + "INT8_SMOOTHQUANT_CFG", + "FP8_DEFAULT_CFG", + "INT4_AWQ_CFG", + "W4A8_AWQ_BETA_CFG", + ], + ) + parser.add_argument( + "--calib_size", type=int, default=512, help="Number of samples for calibration." + ) + + return parser.parse_args() diff --git a/examples/vlm_ptq/vlm_run.py b/examples/vlm_ptq/vlm_run.py new file mode 100644 index 0000000000..f3a088611d --- /dev/null +++ b/examples/vlm_ptq/vlm_run.py @@ -0,0 +1,93 @@ +# 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 os + +import tensorrt_llm +import tensorrt_llm.profiler as profiler +from tensorrt_llm import logger +from tensorrt_llm.runtime import MultimodalModelRunner +from utils import parse_arguments + + +def print_result(model, input_text, output_text, args): + logger.info("---------------------------------------------------------") + if model.model_type != "nougat": + logger.info(f"\n[Q] {input_text}") + for i in range(len(output_text)): + logger.info(f"\n[A]: {output_text[i]}") + + if args.num_beams == 1: + output_ids = model.tokenizer(output_text[0][0], add_special_tokens=False)["input_ids"] + logger.info(f"Generated {len(output_ids)} tokens") + + if args.check_accuracy: + if model.model_type != "nougat": + if model.model_type == "vila": + if len(args.image_path.split(args.path_sep)) == 1: + assert ( + output_text[0][0].lower() + == "the image captures a bustling city intersection teeming with life. " + "from the perspective of a car's dashboard camera, we see" + ) + elif model.model_type == "fuyu": + assert output_text[0][0].lower() == "4" + elif model.model_type == "pix2struct": + assert ( + "characteristic | cat food, day | cat food, wet | cat treats" + in output_text[0][0].lower() + ) + elif model.model_type in ["blip2", "neva", "phi-3-vision", "llava_next"]: + assert "singapore" in output_text[0][0].lower() + elif model.model_type == "video-neva": + assert "robot" in output_text[0][0].lower() + elif model.model_type == "kosmos-2": + assert "snowman" in output_text[0][0].lower() + else: + assert output_text[0][0].lower() == "singapore" + + if args.run_profiling: + + def msec_per_batch(name): + return 1000 * profiler.elapsed_time_in_sec(name) / args.profiling_iterations + + logger.info("Latencies per batch (msec)") + logger.info("TRT vision encoder: %.1f" % (msec_per_batch("Vision"))) + logger.info("TRTLLM LLM generate: %.1f" % (msec_per_batch("LLM"))) + logger.info("Multimodal generate: %.1f" % (msec_per_batch("Generate"))) + + logger.info("---------------------------------------------------------") + + +if __name__ == "__main__": + os.environ["TOKENIZERS_PARALLELISM"] = "false" + args = parse_arguments("run") + logger.set_level(args.log_level) + + model = MultimodalModelRunner(args) + raw_image = model.load_test_data() + if args.input_text: + input_text = args.input_text.split(args.prompt_sep) + if len(input_text) == 1: + input_text = input_text[0] + else: + input_text = args.input_text + + num_iters = args.profiling_iterations if args.run_profiling else 1 + for _ in range(num_iters): + input_text, output_text = model.run(input_text, raw_image, args.max_new_tokens) + runtime_rank = tensorrt_llm.mpi_rank() + if runtime_rank == 0: + print_result(model, input_text, output_text, args) diff --git a/examples/vlm_ptq/vlm_visual_engine.py b/examples/vlm_ptq/vlm_visual_engine.py new file mode 100644 index 0000000000..f74b8dc197 --- /dev/null +++ b/examples/vlm_ptq/vlm_visual_engine.py @@ -0,0 +1,26 @@ +# 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 argparse + +from tensorrt_llm.tools.multimodal_builder import VisionEngineBuilder, add_multimodal_arguments + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser = add_multimodal_arguments(parser) + args = parser.parse_args() + + builder = VisionEngineBuilder(args) + builder.build() diff --git a/windows/Benchmark.md b/examples/windows/Benchmark.md similarity index 100% rename from windows/Benchmark.md rename to examples/windows/Benchmark.md diff --git a/windows/README.md b/examples/windows/README.md similarity index 95% rename from windows/README.md rename to examples/windows/README.md index e86e4c672c..c80d4cbbe9 100644 --- a/windows/README.md +++ b/examples/windows/README.md @@ -5,8 +5,8 @@ #### A Library to Quantize and Compress Deep Learning Models for Optimized Inference on Native Windows RTX GPUs [![Documentation](https://img.shields.io/badge/Documentation-latest-brightgreen.svg?style=flat)](https://nvidia.github.io/TensorRT-Model-Optimizer/) -[![version](https://img.shields.io/pypi/v/nvidia-modelopt?label=Release)](https://pypi.org/project/nvidia-modelopt/) -[![license](https://img.shields.io/badge/License-MIT-blue)](../LICENSE) +[![version](https://img.shields.io/badge/v0.19.0-orange?label=Release)](https://pypi.org/project/nvidia-modelopt/0.19.0/) +[![license](https://img.shields.io/badge/License-Apache%202.0-blue)](../../LICENSE) [Examples](#examples) | [Benchmark Results](#benchmark-results) @@ -47,7 +47,6 @@ To install ModelOpt-Windows as a standalone toolkit on CUDA 12.x systems, run th ```bash pip install nvidia-modelopt[onnx] --extra-index-url https://pypi.nvidia.com -pip install cupy-cuda12x ``` ### Installation with Olive diff --git a/windows/accuracy_benchmark/README.md b/examples/windows/accuracy_benchmark/README.md similarity index 100% rename from windows/accuracy_benchmark/README.md rename to examples/windows/accuracy_benchmark/README.md diff --git a/windows/accuracy_benchmark/mmlu_benchmark.py b/examples/windows/accuracy_benchmark/mmlu_benchmark.py similarity index 94% rename from windows/accuracy_benchmark/mmlu_benchmark.py rename to examples/windows/accuracy_benchmark/mmlu_benchmark.py index c05ed7c4a6..666024ec3d 100644 --- a/windows/accuracy_benchmark/mmlu_benchmark.py +++ b/examples/windows/accuracy_benchmark/mmlu_benchmark.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 gc import json @@ -430,6 +424,7 @@ def main( subject: Union[str, tuple] = None, ep: str = "pt", max_seq_length: int = 4096, + trust_remote_code: bool = False, **kwargs, ): random.seed(RAND_SEED) @@ -503,7 +498,7 @@ def main( elif ep == "pt": model_ckpt_path = kwargs["model_path"] - tokenizer = get_tokenizer(model_ckpt_path) + tokenizer = get_tokenizer(model_ckpt_path, trust_remote_code=trust_remote_code) model = select_model( max_input_length=MAX_SEQ_LEN, max_output_length=2, dtype=dtype, **kwargs @@ -535,6 +530,7 @@ def main( vocab_file=None, model_name=model_name, model_version=model_version, + trust_remote_code=trust_remote_code, ) runner_cls = ModelRunner if not PYTHON_BINDINGS else ModelRunnerCpp diff --git a/windows/accuracy_benchmark/modeling.py b/examples/windows/accuracy_benchmark/modeling.py similarity index 90% rename from windows/accuracy_benchmark/modeling.py rename to examples/windows/accuracy_benchmark/modeling.py index fe39181793..4d8501c7d3 100644 --- a/windows/accuracy_benchmark/modeling.py +++ b/examples/windows/accuracy_benchmark/modeling.py @@ -1,25 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 typing import Optional import torch import torch.nn as nn @@ -117,7 +111,7 @@ class SeqToSeqModel(EvalModel): assert self.tokenizer is not None, "Tokenizer must be loaded to count text length." return len(self.tokenizer(text).input_ids) - def get_choice(self, text: str, **kwargs) -> Tuple[float, float]: + def get_choice(self, text: str, **kwargs) -> tuple[float, float]: self.load() assert self.model is not None, "Model must be loaded to get choices." assert self.tokenizer is not None, "Tokenizer must be loaded to get choices." @@ -178,7 +172,7 @@ class CausalModel(SeqToSeqModel): batch_size, length = inputs.input_ids.shape return self.tokenizer.decode(outputs[0, length:], skip_special_tokens=True) - def get_choice(self, text: str, **kwargs) -> Tuple[float, float]: + def get_choice(self, text: str, **kwargs) -> tuple[float, float]: self.load() assert self.model is not None, "Model must be loaded to get choices." assert self.tokenizer is not None, "Tokenizer must be loaded to get choices." @@ -235,7 +229,7 @@ class AutoAWQCausalModel(SeqToSeqModel): batch_size, length = inputs.input_ids.shape return self.tokenizer.decode(outputs[0, length:], skip_special_tokens=True) - def get_choice(self, text: str, **kwargs) -> Tuple[float, float]: + def get_choice(self, text: str, **kwargs) -> tuple[float, float]: self.load() assert self.model is not None, "Model must be loaded before running." assert self.tokenizer is not None, "Tokenizer must be loaded before running." diff --git a/windows/accuracy_benchmark/quantization_utils.py b/examples/windows/accuracy_benchmark/quantization_utils.py similarity index 79% rename from windows/accuracy_benchmark/quantization_utils.py rename to examples/windows/accuracy_benchmark/quantization_utils.py index 727326b81a..65bc0e59c5 100644 --- a/windows/accuracy_benchmark/quantization_utils.py +++ b/examples/windows/accuracy_benchmark/quantization_utils.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 tqdm import tqdm from transformers import AutoTokenizer @@ -34,11 +28,14 @@ MAX_SEQ_LEN = 2048 MAX_OUTPUT_LEN = 512 -def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN): +def get_tokenizer(ckpt_path, max_seq_len=MAX_SEQ_LEN, trust_remote_code=False): """Returns the tokenizer from the model ckpt_path.""" print(f"Initializing tokenizer from {ckpt_path}") tokenizer = AutoTokenizer.from_pretrained( - ckpt_path, model_max_length=max_seq_len, padding_side="left", trust_remote_code=True + ckpt_path, + model_max_length=max_seq_len, + padding_side="left", + trust_remote_code=trust_remote_code, ) if type(tokenizer).__name__ == "QWenTokenizer": # qwen use token id 151643 as pad and eos tokens diff --git a/windows/accuracy_benchmark/requirements.txt b/examples/windows/accuracy_benchmark/requirements.txt similarity index 86% rename from windows/accuracy_benchmark/requirements.txt rename to examples/windows/accuracy_benchmark/requirements.txt index 3c3f2edb5e..3c3222a83e 100644 --- a/windows/accuracy_benchmark/requirements.txt +++ b/examples/windows/accuracy_benchmark/requirements.txt @@ -8,3 +8,4 @@ peft>=0.5.0 rwkv>=0.7.3 tiktoken==0.7.0 tqdm==4.66.5 +transformers==4.44.0 diff --git a/windows/accuracy_benchmark/trtllm_utils.py b/examples/windows/accuracy_benchmark/trtllm_utils.py similarity index 92% rename from windows/accuracy_benchmark/trtllm_utils.py rename to examples/windows/accuracy_benchmark/trtllm_utils.py index 5cb5bf97e8..73b3d7cabb 100644 --- a/windows/accuracy_benchmark/trtllm_utils.py +++ b/examples/windows/accuracy_benchmark/trtllm_utils.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 @@ -131,6 +125,7 @@ def load_tokenizer( model_name: str = "GPTForCausalLM", model_version: Optional[str] = None, tokenizer_type: Optional[str] = None, + trust_remote_code: bool = False, ): if vocab_file is None: use_fast = True @@ -142,7 +137,7 @@ def load_tokenizer( legacy=False, padding_side="left", truncation_side="left", - trust_remote_code=True, + trust_remote_code=trust_remote_code, tokenizer_type=tokenizer_type, use_fast=use_fast, ) diff --git a/windows/onnx_ptq/README.md b/examples/windows/onnx_ptq/README.md similarity index 100% rename from windows/onnx_ptq/README.md rename to examples/windows/onnx_ptq/README.md diff --git a/windows/onnx_ptq/quantize.py b/examples/windows/onnx_ptq/quantize.py similarity index 92% rename from windows/onnx_ptq/quantize.py rename to examples/windows/onnx_ptq/quantize.py index a3570d4e01..3d35af75fd 100644 --- a/windows/onnx_ptq/quantize.py +++ b/examples/windows/onnx_ptq/quantize.py @@ -1,23 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT +# SPDX-License-Identifier: Apache-2.0 # -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: +# 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 # -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. +# http://www.apache.org/licenses/LICENSE-2.0 # -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. +# 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 logging @@ -171,14 +165,17 @@ def get_calib_inputs( add_past_kv_inputs, max_calib_rows_to_load, add_position_ids, + trust_remote_code, ): # from transformers import LlamaConfig - # config = LlamaConfig.from_pretrained(model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=True) + # config = LlamaConfig.from_pretrained( + # model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=trust_remote_code + # ) config = AutoConfig.from_pretrained( - model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=True + model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=trust_remote_code ) tokenizer = AutoTokenizer.from_pretrained( - model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=True + model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=trust_remote_code ) tokenizer.add_special_tokens({"pad_token": "[PAD]"}) tokenizer.pad_token = tokenizer.eos_token @@ -349,6 +346,7 @@ def main(args): args.add_past_kv_inputs, 128, args.add_position_ids, + args.trust_remote_code, ) t = time.time() @@ -516,6 +514,12 @@ if __name__ == "__main__": default=["dml", "cpu"], # Default as a list help="Comma-separated list of calibration endpoints. Choose from 'cuda', 'cpu', 'dml'. Default is 'dml,cpu'.", ) + parser.add_argument( + "--trust_remote_code", + help="Set trust_remote_code for Huggingface models and tokenizers", + default=False, + action="store_true", + ) args = parser.parse_args() main(args) diff --git a/windows/onnx_ptq/requirements.txt b/examples/windows/onnx_ptq/requirements.txt similarity index 66% rename from windows/onnx_ptq/requirements.txt rename to examples/windows/onnx_ptq/requirements.txt index 6dee05e6fd..ae1ecc8aa8 100644 --- a/windows/onnx_ptq/requirements.txt +++ b/examples/windows/onnx_ptq/requirements.txt @@ -1,3 +1,3 @@ -datasets>=2.14.4 +datasets>=2.14.5 torch==2.4.0 transformers==4.44.0 diff --git a/llm_distill/requirements.txt b/llm_distill/requirements.txt deleted file mode 100644 index 3262a26dce..0000000000 --- a/llm_distill/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -datasets -pyarrow>=16.0.0 -trl==0.9.4 diff --git a/llm_eval/__init__.py b/llm_eval/__init__.py deleted file mode 100644 index c090f38848..0000000000 --- a/llm_eval/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT -# -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. diff --git a/llm_ptq/benchmarks/benchmark_suite.py b/llm_ptq/benchmarks/benchmark_suite.py deleted file mode 100644 index 46945b4443..0000000000 --- a/llm_ptq/benchmarks/benchmark_suite.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT -# -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. - -# This import overrides the benchmark executor and disable chunked_context -import disable_bench_chunked_context # noqa -import tensorrt_llm.commands.bench as trtllm_bench - -if __name__ == "__main__": - trtllm_bench.main() diff --git a/llm_ptq/benchmarks/disable_bench_chunked_context.py b/llm_ptq/benchmarks/disable_bench_chunked_context.py deleted file mode 100644 index 46ba5af6e1..0000000000 --- a/llm_ptq/benchmarks/disable_bench_chunked_context.py +++ /dev/null @@ -1,42 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT -# -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. - -from pathlib import Path -from typing import Dict, Tuple, Union - -import tensorrt_llm.bench.benchmark.utils as utils -from tensorrt_llm.bench.benchmark.utils import ( - get_settings_from_engine as get_settings_from_engine_original, -) - - -def get_settings_from_engine_override( - engine_path: Path, -) -> Tuple[Dict[str, Union[str, int]], Dict[str, Union[str, int]]]: - exec_settings, build_cfg = get_settings_from_engine_original(engine_path) - # Do not enable chunking if use_paged_context_fmha if not enabled. - exec_settings["settings_config"]["chunking"] = build_cfg["plugin_config"][ - "use_paged_context_fmha" - ] - return exec_settings, build_cfg - - -utils.get_settings_from_engine = get_settings_from_engine_override diff --git a/llm_sparsity/utils.py b/llm_sparsity/utils.py deleted file mode 100644 index 376a4ffcf3..0000000000 --- a/llm_sparsity/utils.py +++ /dev/null @@ -1,37 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT -# -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. - -import io -import json - - -def _make_r_io_base(f, mode: str): - if not isinstance(f, io.IOBase): - f = open(f, mode=mode) - return f - - -def jload(f, mode="r"): - """Load a .json file into a dictionary.""" - f = _make_r_io_base(f, mode) - jdict = json.load(f) - f.close() - return jdict diff --git a/model_hub/download_hf_ckpt.py b/model_hub/download_hf_ckpt.py deleted file mode 100644 index 80d100fc63..0000000000 --- a/model_hub/download_hf_ckpt.py +++ /dev/null @@ -1,37 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT -# -# Permission is hereby granted, free of charge, to any person obtaining a -# copy of this software and associated documentation files (the "Software"), -# to deal in the Software without restriction, including without limitation -# the rights to use, copy, modify, merge, publish, distribute, sublicense, -# and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL -# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -# DEALINGS IN THE SOFTWARE. - -import os - -from huggingface_hub import login, snapshot_download - -login(token="", add_to_git_credential=True) -model_names = [ - "Llama-3.1-8B-Instruct-FP8", - "Llama-3.1-70B-Instruct-FP8", - "Llama-3.1-405B-Instruct-FP8", -] -for model_name in model_names: - hf_repo = "nvidia/" + model_name - local_dir = "quantized_ckpts/" + model_name - os.makedirs(local_dir, exist_ok=True) - print(f"downloading {model_name}...") - snapshot_download(repo_id=hf_repo, local_dir=local_dir) diff --git a/modelopt/__init__.py b/modelopt/__init__.py new file mode 100644 index 0000000000..1244a36f22 --- /dev/null +++ b/modelopt/__init__.py @@ -0,0 +1,42 @@ +# 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. + +"""Nvidia TensorRT Model Optimizer (modelopt).""" + +import sys as _sys +import warnings as _warnings +from importlib.metadata import version as _version + +__version__ = _version("nvidia-modelopt") + + +try: + # Import from local source if available + from . import core + + __core_version__ = __version__ +except ImportError: + # Import from nvidia-modelopt-core wheel installation + import modelopt_core as core # type: ignore[no-redef] + + _sys.modules["modelopt.core"] = core + __core_version__ = _version("nvidia-modelopt-core") + +# Versions need to be the same for compatibility +if __version__.split(".")[:2] != __core_version__.split(".")[:2]: + _warnings.warn( + f"Version mismatch between nvidia-modelopt ({__version__}) and nvidia-modelopt-core" + f" ({__core_version__}). Please ensure both versions are the same for compatibility." + ) diff --git a/modelopt/deploy/__init__.py b/modelopt/deploy/__init__.py new file mode 100644 index 0000000000..b3916619b3 --- /dev/null +++ b/modelopt/deploy/__init__.py @@ -0,0 +1,16 @@ +# 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. + +"""Model Optimizer's deployment package.""" diff --git a/modelopt/deploy/llm/__init__.py b/modelopt/deploy/llm/__init__.py new file mode 100644 index 0000000000..2a97f9c7ee --- /dev/null +++ b/modelopt/deploy/llm/__init__.py @@ -0,0 +1,33 @@ +# 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. + +"""LLM deployment utils with tensorrt_llm.""" + +from mpi4py import MPI + +# Pre load MPI libs to avoid tensorrt_llm importing failures. +print(f"Loaded mpi lib {MPI.__file__} successfully") + +# Pre import tensorrt_llm +try: + import tensorrt_llm +except Exception as e: + print( + "tensorrt_llm package is not installed. Please build or install tensorrt_llm package" + " properly before calling the llm deployment API." + ) + raise (e) + +from .generate import * # noqa diff --git a/modelopt/deploy/llm/generate.py b/modelopt/deploy/llm/generate.py new file mode 100644 index 0000000000..11d6739bf9 --- /dev/null +++ b/modelopt/deploy/llm/generate.py @@ -0,0 +1,255 @@ +# 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. + +"""A wrapper over the TensorRT-LLM high level API runner.""" + +import json +from pathlib import Path +from typing import Any, Iterable, Optional, Union + +import tensorrt_llm +import torch +from packaging.version import parse +from tensorrt_llm.bindings.executor import DecodingConfig +from tensorrt_llm.llmapi import KvCacheConfig as TRT_KvCacheConfig +from tensorrt_llm.llmapi import SamplingParams +from tensorrt_llm.llmapi.llm import LLM as TRT_LLM +from tensorrt_llm.llmapi.tokenizer import TokenizerBase, TransformersTokenizer + + +def _sanitize_temperature_and_top_p(temperature, top_p): + assert temperature >= 0.0, "Temperature must be greater than 0.0." + + # TRT LLM acccepts temperature values only greater than 0.0 + temperature = max(temperature, 0.001) + + kwargs = {"temperature": temperature} + if top_p is not None: + # cpp executor only supports topP.value() > 0.f + top_p = 1e-4 if top_p == 0 else top_p + kwargs["top_p"] = top_p + + return kwargs + + +class LLM(TRT_LLM): + """A wrapper over the ``tensorrt_llm.llmapi.llm.LLM`` for LLM profiling and validation.""" + + def __init__( + self, + engine_dir: Union[str, Path], + tokenizer: Optional[Union[str, Path, TokenizerBase]] = None, + kv_cache_config: dict[str, Union[int, float]] = {}, + medusa_choices: Any = None, + ): + """Initializes the LLM runner class. + + Args: + engine_dir: the directory path of the TensorRT-LLM engine. + tokenizer: the tokenizer. For example, a tokenizer from the Huggingface model. + kv_cache_config: the kv cache config as a dict. Please refer to + https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/performance/perf-best-practices.md + """ + assert parse(tensorrt_llm.__version__) >= parse("0.15.0") + + with open(Path(engine_dir) / "config.json", "r") as engine_config_file: + engine_config = json.load(engine_config_file) + build_config = engine_config["build_config"] + world_size = ( + engine_config.get("pretrained_config", {}).get("mapping", {}).get("world_size", 1) + ) + max_tokens_kv_cache = build_config["max_seq_len"] * build_config["max_batch_size"] + self.gather_context_logits = build_config.get("gather_context_logits", False) + + trt_kv_cache_config = TRT_KvCacheConfig(enable_block_reuse=False) + + # If not specified, free_gpu_memory_fraction is set to the default TRT LLM value 0.9 + trt_kv_cache_config.free_gpu_memory_fraction = kv_cache_config.get( + "free_gpu_memory_fraction", 0.9 + ) + + # If not specified, max_tokens is set to the max value calculated above. + trt_kv_cache_config.max_tokens = kv_cache_config.get("max_tokens", max_tokens_kv_cache) + + kwargs = {} + if medusa_choices is not None: + decoding_config = DecodingConfig() + decoding_config.medusa_choices = medusa_choices + kwargs["decoding_config"] = decoding_config + assert world_size == 1, "decoding_config does not support multi TP in HLAPI." + + if tokenizer is None: + # Assume the tokenizer is stored in the engine_dir if not specified. + tokenizer = engine_dir + + # CustomSentencePieceTokenizer will not be recognized by llmapi, wrapping it around TransformersTokenizer + if type(tokenizer).__name__ in ["CustomSentencePieceTokenizer"]: + tokenizer = TransformersTokenizer(tokenizer) + + super().__init__( + model=engine_dir, + tokenizer=tokenizer, + kv_cache_config=trt_kv_cache_config, + **kwargs, + ) + + @property + def max_input_len(self): + """Get the max input length from the LLM instance.""" + return self.args.build_config.max_input_len + + @property + def max_beam_width(self): + """Get the max beam width from the LLM instance.""" + return self.args.build_config.max_beam_width + + def generate_tokens( + self, + prompts: Union[Iterable[str], Iterable[list[int]]], + max_new_tokens: int, + temperature: float = 1.0, + top_p: float = None, + keep_input_prompt: bool = True, + stop_words: list[str] = None, + ) -> Union[list[list[int]], list[list[list[int]]]]: + """Generates the tokens based on the input prompts. + + Args: + prompts: The input prompts. Could be a list of strings or token lists. + max_new_tokens: The max output token length. + temperature: The sampling temperature. + top_p: The nucleus sampling parameter. + keep_input_prompt: Set to include input prommpts in the outputs. + stop_words: A list of words that the generate stops on. + + Returns: + a list of output token lists if max_beam_width is 1 or a 3D list with shape [batch, beam, sequence_len]. + """ + assert temperature >= 0.0, "Temperature must be greater than 0.0." + + beam_width = self.max_beam_width + kwargs = _sanitize_temperature_and_top_p(temperature, top_p) + sampling_config = SamplingParams( + max_tokens=max_new_tokens, beam_width=beam_width, stop=stop_words, **kwargs + ) + + prompt_ids = [ + self.tokenizer.encode(prompt) if isinstance(prompt, str) else prompt + for prompt in prompts + ] + outputs = self.generate(prompt_ids, sampling_params=sampling_config, use_tqdm=False) + + def _process_output_token_id(output_token_id, prompt_id, with_input, keep_input_prompt): + if with_input == keep_input_prompt: + return output_token_id + + elif with_input: # and not keep_input_prompt + return output_token_id[len(prompt_id) :] + + else: # not with_input and keep_input_prompt: + return prompt_id + output_token_id + + with_input = False + output_tokens = [] + for prompt_id, output in zip(prompt_ids, outputs): + output_token_ids = [out.token_ids for out in output.outputs] + + for output_token_id in output_token_ids: + output_tokens.append( + _process_output_token_id( + output_token_id, prompt_id, with_input, keep_input_prompt + ) + ) + + return ( + output_tokens + if beam_width == 1 + else [ + output_tokens[i : i + beam_width] for i in range(0, len(output_tokens), beam_width) + ] + ) + + def generate_text( + self, + prompts: Union[Iterable[str], Iterable[list[int]]], + max_new_tokens: int, + temperature: float = 1.0, + top_p: float = None, + keep_input_prompt: bool = True, + stop_words: list[str] = None, + ) -> Union[list[str], list[list[str]]]: + """Generates the text based on the input prompts. + + Args: + prompts: The input prompts. Could be a list of strings or token lists. + max_new_tokens: The max output token length. + temperature: The sampling temperature + keep_input_prompt: Set to include input prommpts in the outputs. + stop_words: A list of words the generate will stop on. + + Returns: + a list of output text strings if max_beam_width is 1 or a 2D list with shape [batch, beam]. + """ + beam_width = self.max_beam_width + output_tokens = self.generate_tokens( + prompts, + max_new_tokens, + temperature, + keep_input_prompt=keep_input_prompt, + top_p=top_p, + stop_words=stop_words, + ) + if beam_width == 1: + output_text = [self.tokenizer.decode(batch) for batch in output_tokens] + else: + output_text = [ + [self.tokenizer.decode(beam) for beam in batch] for batch in output_tokens + ] + return output_text + + def generate_context_logits( + self, + prompts: Union[Iterable[str], Iterable[list[int]]], + temperature: float = 1.0, + top_p: float = None, + ) -> list[torch.tensor]: + """Generates the context logits based on the input prompts. + + Args: + prompts: The input prompts. Could be a list of strings or token lists. + temperature: The sampling temperature. + top_p: The nucleus sampling parameter. + keep_input_prompt: Set to include input prommpts in the outputs. + + Returns: + a tensor list of the context_logits. + """ + assert ( + self.gather_context_logits + ), "Please enable gather_context_logits flag when building the engine." + assert temperature >= 0.0, "Temperature must be greater than 0.0." + + kwargs = _sanitize_temperature_and_top_p(temperature, top_p) + kwargs["return_context_logits"] = True + + sampling_config = SamplingParams(max_tokens=1, beam_width=1, **kwargs) + + prompt_ids = [ + self.tokenizer.encode(prompt) if isinstance(prompt, str) else prompt + for prompt in prompts + ] + outputs = self.generate(prompt_ids, sampling_params=sampling_config, use_tqdm=False) + + return [output.context_logits for output in outputs] diff --git a/modelopt/deploy/llm/nemo_utils.py b/modelopt/deploy/llm/nemo_utils.py new file mode 100644 index 0000000000..f5a48df129 --- /dev/null +++ b/modelopt/deploy/llm/nemo_utils.py @@ -0,0 +1,180 @@ +# 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. + +"""The utils to support Nemo models.""" + +import warnings +from pathlib import Path + +import numpy as np +import torch +from transformers import GPT2Tokenizer, PreTrainedTokenizer, T5Tokenizer + +try: + from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer +except Exception: + warnings.warn("Cannot import nemo package, falling back to PreTrainedTokenizer!") + SentencePieceTokenizer = PreTrainedTokenizer + + +class CustomSentencePieceTokenizer(SentencePieceTokenizer): + """Custom tokenizer based on Nemo SentencePieceTokenizer. + + This extension of SentencePieceTokenizer is to make API consistent with HuggingFace tokenizers + in order to run evaluation tools in examples/tensorrt_llm/scripts/nemo_example.sh script. + """ + + def __init__(self, *args, **kwargs): + """Constructor method with extra check for non-legacy SentencePieceTokenizer variant.""" + super().__init__(*args, **kwargs) + assert not self.legacy, "Only non-legacy tokenizer is supported" + + @property + def pad_token(self): + """pad_token.""" + return self.pad_id + + @property + def eos_token(self): + """eos_token.""" + return self.eos_id + + @property + def pad_token_id(self): + """pad_token_id.""" + return self.pad_id + + @property + def eos_token_id(self): + """eos_token_id.""" + return self.eos_id + + def encode(self, text, return_tensors=None, max_length=None, **kwargs): + """Method introduced for HF tokenizers API consistency for evaluation scripts. + + Note: kwargs other than return_tensors and max_length are ignored. + """ + output = self.tokenizer.encode_as_ids(text) + if max_length is not None: + if isinstance(text, str): + output = output[:max_length] + if isinstance(text, list): + output = [x[:max_length] for x in output] + if return_tensors == "pt": + # Only plain text input is supported since for list of strings some padding needs to be introduced + assert isinstance( + text, str + ), "Returning 'pt' tensors is only supported for simple text input" + output = torch.LongTensor(output).reshape((1, -1)) + return output + + def batch_encode_plus(self, texts, **kwargs): + """Method introduced for HF tokenizers API consistency for evaluation scripts. + + Note: kwargs are ignored. + """ + assert isinstance(texts, list), f"Expected list of texts got {texts}" + return {"input_ids": self.tokenizer.encode_as_ids(texts)} + + def decode(self, ids, **kwargs): + """MMethod introduced for HF tokenizers API consistency for evaluation scripts. + + Note: kwargs are ignored. + """ + if isinstance(ids, np.ndarray) or torch.is_tensor(ids): + ids = ids.tolist() + return self.tokenizer.decode(ids) + + def batch_decode(self, ids, **kwargs): + """Method introduced for HF tokenizers API consistency for evaluation scripts.""" + return self.decode(ids, **kwargs) + + +def _build_tokenizer(tokenizer_config: dict): + if tokenizer_config["library"] == "sentencepiece": + # Model Optimizer modification. + # Turn off legacy model by default: See https://github.com/huggingface/transformers/pull/24622 + tokenizer = T5Tokenizer(tokenizer_config["model"], extra_ids=0, legacy=False) + elif "GPT2" in tokenizer_config["type"]: + tokenizer = GPT2Tokenizer(tokenizer_config["vocab_file"], tokenizer_config["merge_file"]) + else: + raise ValueError(f'Tokenizer type {tokenizer_config["library"]} not handled') + + if tokenizer.bos_token_id is None: + tokenizer.add_special_tokens({"bos_token": ""}) + if tokenizer.eos_token_id is None: + tokenizer.add_special_tokens({"eos_token": ""}) + + return tokenizer + + +def get_tokenzier(tokenizer_dir_or_path: Path) -> PreTrainedTokenizer: + """Loads the tokenizer from the decoded NEMO weights dir.""" + # TODO: Remove this function sometime in the future + warnings.warn( + ( + "Function get_tokenzier is deprecated and may be removed soon. " + "Please consider using get_nemo_tokenizer instead." + ), + DeprecationWarning, + ) + model_path = ( + tokenizer_dir_or_path / "tokenizer.model" + if tokenizer_dir_or_path.is_dir() + else tokenizer_dir_or_path + ) + tokenizer_config = {"library": "sentencepiece", "model": str(model_path)} + return _build_tokenizer(tokenizer_config) + + +def get_nemo_tokenizer(tokenizer_cfg_path: str): + """Build tokenizer from Nemo tokenizer config. + + Refer to the logic of get_nmt_tokenizer function on how to instantiate tokenizers in Nemo, see + https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/tokenizer_utils.py. + """ + from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer + from omegaconf import OmegaConf + + print(f"Initializing tokenizer from tokenizer config {tokenizer_cfg_path}") + tokenizer_cfg = OmegaConf.load(tokenizer_cfg_path) + + library = tokenizer_cfg.library + legacy = tokenizer_cfg.get("sentencepiece_legacy", library == "sentencepiece") + special_tokens_dict = tokenizer_cfg.get("special_tokens", {}) + + if library == "huggingface": + print(f"Getting HuggingFace AutoTokenizer with pretrained_model_name: {tokenizer_cfg.type}") + tokenizer = AutoTokenizer( + pretrained_model_name=tokenizer_cfg.type, + vocab_file=tokenizer_cfg.get("vocab_file", None), + merges_file=tokenizer_cfg.get("merge_file", None), + use_fast=tokenizer_cfg.get("use_fast", False), + **special_tokens_dict, + ) + print("Unwrapping HuggingFace tokenizer from Nemo AutoTokenizer") + # TODO: Either Nemo tokenizer API alignment is needed https://github.com/NVIDIA/NeMo/pull/8818 + # or alternatively we can use only HuggingFace tokenizers in future. Current workaround + # proposed here is to unwrap the HuggingFace tokenizer from Nemo AutoTokenizer class: + tokenizer = tokenizer.tokenizer + elif library == "sentencepiece": + print(f"Getting SentencePiece with model: {tokenizer_cfg.model}") + tokenizer = CustomSentencePieceTokenizer(model_path=tokenizer_cfg.model, legacy=legacy) + else: + raise NotImplementedError( + "Currently we only support 'huggingface' and 'sentencepiece' tokenizer libraries." + ) + + return tokenizer diff --git a/modelopt/onnx/__init__.py b/modelopt/onnx/__init__.py new file mode 100644 index 0000000000..fbeb8b4015 --- /dev/null +++ b/modelopt/onnx/__init__.py @@ -0,0 +1,35 @@ +# 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. + +"""Model optimization subpackage for onnx.""" + +import sys +import warnings + +MIN_PYTHON_VERSION = (3, 10) + +try: + from . import quantization # noqa: F401 +except ImportError as e: + raise ImportError(f"{e}\nPlease install optional ``[onnx]`` dependencies.") + + +# Check the current Python version +if sys.version_info < MIN_PYTHON_VERSION: + warnings.warn( + f"This package requires Python {MIN_PYTHON_VERSION[0]}.{MIN_PYTHON_VERSION[1]} or higher. " + f"You are using Python {sys.version_info[0]}.{sys.version_info[1]}.", + UserWarning, + ) diff --git a/modelopt/onnx/op_types.py b/modelopt/onnx/op_types.py new file mode 100644 index 0000000000..c4439f8ed6 --- /dev/null +++ b/modelopt/onnx/op_types.py @@ -0,0 +1,287 @@ +# 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. + +"""Utility functions to categorize onnx ops.""" + + +def is_unary_op(op_type: str): + """Returns whether the given op is a unary operator or not.""" + return op_type in [ + "Neg", + "Sqrt", + "Abs", + "Log", + "Exp", + "Not", + "Cast", + "Floor", + "Ceil", + "Round", + "Erf", + "Gelu", + "Sin", + "Cos", + "Atan", + "Sign", + "IsNaN", + "IsInf", + "Log", + "LeakyRelu", + # "Relu", + "Elu", + "Tanh", + "Sigmoid", + # "BatchNormalization", + "Softmax", + "Softplus", + "InstanceNormalization", + "CumSum", + ] + + +def is_binary_op(op_type: str): + """Returns whether the given op is a binary operator or not.""" + return op_type in [ + "Add", + "Sub", + "Mul", + "Pow", + "Div", + "Min", + "Max", + "Greater", + "GreaterOrEqual", + "Less", + "LessOrEqual", + "Equal", + "BitwiseOr", + "BitwiseAnd", + "BitwiseXor", + "BitShift", + ] + + +def is_fusible_reduction_op(op_type: str): + """Returns whether the given op type is of reduction category and fusible by the compiler.""" + return op_type in [ + "ReduceMax", + "ReduceMin", + "ReduceMean", + "ReduceProd", + "ReduceSum", + "TopK", # Transformed to BottomK based on `largest` param + ] + + +def is_fusible_scaling_op(op_type: str): + """Returns whether the given op type is of scaling category and fusible with input.""" + return op_type in [ + "SkipSimplifiedLayerNormalization", + "SimplifiedLayerNormalization", + "MatMul", + "Gemm", + "Mul", + ] + + +def is_copy_op(op_type: str): + """Returns whether the given op is a copy operator or not.""" + return op_type in [ + "Flatten", + "Transpose", + "Concat", + "Split", + "Squeeze", + "Expand", + "ReverseSequence", + "Reshape", + "Tile", + "Gather", + "Slice", + "GatherElements", + "GatherND", + "ScatterElements", + "ScatterND", + "OneHot", + ] + + +def is_linear_op(op_type: str): + """Returns whether the given op type is of Linear category or not.""" + return op_type in ["Conv", "Gemm", "MatMul"] + + +def is_pointwise_or_elementwise_op(op_type: str): + """Returns whether the given op type is of Pointwise or Elementwise category or not. + + This considers only the fusible types. + """ + return is_unary_op(op_type) or is_binary_op(op_type) + + +def is_pooling_or_window_op(op_type: str): + """Returns whether the given op type is of Pooling/Window category or not.""" + return op_type in [ + "AveragePool", + "GlobalAveragePool", + "MaxPool", + "GlobalMaxPool", + "GlobalLpPool", + "LpPool", + "MaxPoolGridSample", + "HammingWindow", + "BlackmanWindow", + "HannWindow", + ] + + +def is_normalization_op(op_type: str): + """Returns whether the given op type is of Normalization category or not.""" + return op_type in [ + "BatchNormalization", + "InstanceNormalization", + "LRN", + "LpNormalization", + "GroupNormalization", + "LayerNormalization", + ] + + +def is_conversion_op(op_type: str): + """Returns whether the given op type is of Conversion category or not.""" + return op_type in ["Cast", "QuantizeLinear", "DequantizeLinear"] + + +def is_non_reshape_copy_op(op_type: str): + """Returns whether the given op is a non-reshape copy op or not.""" + return is_copy_op(op_type) and (op_type != "Reshape") + + +def is_irregular_mem_access_op(op_type: str): + """Returns whether the given op type is of Irreggular mem access category or not.""" + return op_type in [ + "Gather", + "GatherElements", + "GatherND", + "MaxRoiPool", + "RoiAlign", + "ScatterND", + "ScatterElements", + "NonMaxSuppression", + ] + + +def is_generator_op(op_type: str): + """Returns whether the given op type is of Generator category or not.""" + return op_type in [ + "Const", + "ConstOfShape", + "EyeLike", + "OneHot", + "Multinomial", + "RandomNormal", + "RandomUniform", + "Bernoulli", + ] + + +def is_modifier_op(op_type: str): + """Returns whether the given op type is of Modifier category or not.""" + return op_type in [ + "Identity", + "Trilu", + "Expand", + "Pad", + "Dropout", + "TileDropout", + "Col2Im", + "MaxUnpool", + ] + + +def is_sequence_op(op_type: str): + """Returns whether the given op type is of Sequence category or not.""" + return op_type in [ + "SequenceAt", + "SequenceConstruct", + "SequenceEmpty", + "SequenceErase", + "SequenceInsert", + "SequenceLength", + ] + + +def is_selection_op(op_type: str): + """Returns whether the given op type is of Selection category or not.""" + return op_type in ["Where", "Compress"] + + +def is_control_flow_op(op_type: str): + """Returns whether the given op type is of Control Flow category or not.""" + return op_type in ["If", "Loop"] + + +def is_multiclass_op(op_type: str): + """Returns whether the given op type is of Multiclass category or not.""" + return op_type in ["Einsum"] + + +def is_recurrent_op(op_type: str): + """Returns whether the given op type is of Recurrent category or not.""" + return op_type in ["LSTM", "RNN", "GRU"] + + +def is_shape_op(op_type: str): + """Returns whether the given op type is of Shape category or not.""" + return op_type in ["Shape", "Size"] + + +def is_default_quantizable_op_by_ort(op_type: str): + """Returns if ORT quantizes the op type by default. + + Note. Subject to change with different ORT versions. + Note. Users can use nodes_to_quantize and/or op_types_to_quantize arguments to quantize + non-default operations. + Reference: https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/quantization/registry.py + """ + return op_type in [ + "Conv", + "Gemm", + "ArgMax", + "Relu", + "Split", + "MaxPool", + "InstanceNormalization", + "Softmax", + "Where", + "Squeeze", + "GlobalAveragePool", + "Pad", + "Resize", + "ConvTranspose", + "Gather", + "Sigmoid", + "EmbedLayerNormalization", + "Reshape", + "Unsqueeze", + "Transpose", + "MatMul", + "Concat", + "Mul", + "Clip", + "Add", + "LeakyRelu", + "AveragePool", + ] diff --git a/modelopt/onnx/quantization/__init__.py b/modelopt/onnx/quantization/__init__.py new file mode 100644 index 0000000000..47b3ae22d6 --- /dev/null +++ b/modelopt/onnx/quantization/__init__.py @@ -0,0 +1,20 @@ +# 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. + +"""Model optimization subpackage for onnx quantization.""" + +# TODO: Use high-level quantize instead of quantize_int4 +from .int4 import quantize as quantize_int4 +from .quantize import quantize diff --git a/modelopt/onnx/quantization/__main__.py b/modelopt/onnx/quantization/__main__.py new file mode 100644 index 0000000000..b4e06c77af --- /dev/null +++ b/modelopt/onnx/quantization/__main__.py @@ -0,0 +1,233 @@ +# 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. + +"""Command-line entrypoint for ONNX PTQ.""" + +import argparse + +import numpy as np + +from modelopt.onnx.quantization.quantize import quantize + +__all__ = ["main"] + + +def parse_args(): + argparser = argparse.ArgumentParser("python -m modelopt.onnx.quantization") + group = argparser.add_mutually_exclusive_group(required=False) + argparser.add_argument( + "--onnx_path", required=True, type=str, help="Input onnx model without Q/DQ nodes." + ) + argparser.add_argument( + "--quantize_mode", + type=str, + choices=["fp8", "int8", "int4"], + default="int8", + help=("Quantization mode for the given ONNX model."), + ) + argparser.add_argument( + "--calibration_method", + type=str, + choices=["max", "entropy", "awq_clip", "rtn_dq"], + help=( + "Calibration method choices for fp8: {max (default)}, " + "int8: {entropy (default), max}, int4: {awq_clip (default), rtn_dq}." + ), + ) + group.add_argument( + "--calibration_data_path", + type=str, + help="Calibration data in npz/npy format. If None, random data for calibration will be used.", + ) + group.add_argument( + "--calibration_cache_path", + type=str, + help="Pre-calculated activation tensor scaling factors aka calibration cache path.", + ) + argparser.add_argument( + "--calibration_shapes", + type=str, + required=False, + help=( + "Optional model input shapes for calibration." + "Users should provide the shapes specifically if the model has non-batch dynamic dimensions." + "Example input shapes spec: input0:1x3x256x256,input1:1x3x128x128" + ), + ) + argparser.add_argument( + "--calibration_eps", + type=str, + default=["cuda:0", "cpu", "trt"], + nargs="+", + help=( + "Priority order for the execution providers (EP) to calibrate the model. " + "Any subset of ['cuda:x', dml:x, 'cpu', 'trt'], where 'x' is the device id." + ), + ) + argparser.add_argument( + "--op_types_to_quantize", + type=str, + default=[], + nargs="+", + help="A space-separated list of node types to quantize.", + ) + argparser.add_argument( + "--op_types_to_exclude", + type=str, + default=[], + nargs="+", + help="A space-separated list of node types to exclude from quantization.", + ) + argparser.add_argument( + "--nodes_to_quantize", + type=str, + default=[], + nargs="+", + help="A space-separated list of node names to quantize. Regular expressions are supported.", + ) + argparser.add_argument( + "--nodes_to_exclude", + type=str, + default=[], + nargs="+", + help="A space-separated list of node names to exclude from quantization. Regular expressions are supported.", + ) + argparser.add_argument( + "--use_external_data_format", + action="store_true", + help="If True, .onnx_data will be used to load and/or write weights and constants.", + ) + argparser.add_argument( + "--keep_intermediate_files", + action="store_true", + help=( + "If True, keep the files generated during the ONNX models' conversion/calibration. " + "Otherwise, only the converted ONNX file is kept for the user." + ), + ) + argparser.add_argument( + "--output_path", + type=str, + help=( + "Output filename to save the converted ONNX model. If None, save it in the same dir as " + "the original ONNX model with an appropriate suffix." + ), + ) + argparser.add_argument( + "--verbose", + action="store_true", + help="If verbose, print all the debug info.", + ) + argparser.add_argument( + "--trt_plugins", + type=str, + default=None, + help=( + "Specifies custom TensorRT plugin library paths in .so format (compiled shared library). " + 'For multiple paths, separate them with a semicolon, i.e.: "lib_1.so;lib_2.so". ' + "If this is not None, the TensorrtExecutionProvider is invoked, so make sure that the TensorRT libraries " + "are in the PATH or LD_LIBRARY_PATH variables." + ), + ) + argparser.add_argument( + "--trt_plugins_precision", + type=str, + default=None, + nargs="+", + help=( + "A space-separated list indicating the precision for each custom op. " + "Each item should have the format :, where precision can be fp32 (default) or fp16. " + "For example: op_type_1:fp16 op_type_2:fp32." + ), + ) + argparser.add_argument( + "--high_precision_dtype", + type=str, + default=None, + help=( + "High precision data type, one of ['fp32', 'fp16']. For int8 quantization, the default value is 'fp32' and " + "'fp16' for other quantization modes." + ), + ) + argparser.add_argument( + "--mha_accumulation_dtype", + type=str, + default="fp32", + help=( + "Accumulation dtype of MHA. This flag will only take effect when mha_accumulation_dtype == 'fp32' " + "and quantize_mode == 'fp8'. One of ['fp32', 'fp16']" + ), + ) + argparser.add_argument( + "--disable_mha_qdq", + action="store_true", + help="If True, Q/DQ will not be added to MatMuls in MHA pattern.", + ) + argparser.add_argument( + "--dq_only", + action="store_true", + help=( + "If True, FP32/FP16 weights will be converted to INT8/FP8 weights. Q nodes will get removed from the " + "weights and have only DQ nodes with those converted INT8/FP8 weights in the output model." + ), + ) + argparser.add_argument( + "--use_zero_point", + type=bool, + default=False, + help=( + "If True, zero-point based quantization will be used - currently, applicable for awq_lite algorithm." + ), + ) + return argparser.parse_args() + + +def main(): + """Command-line entrypoint for ONNX PTQ.""" + args = parse_args() + calibration_data = None + if args.calibration_data_path: + calibration_data = np.load(args.calibration_data_path, allow_pickle=True) + + default_high_precision_dtype = "fp32" if args.quantize_mode == "int8" else "fp16" + + quantize( + args.onnx_path, + quantize_mode=args.quantize_mode, + calibration_data=calibration_data, + calibration_method=args.calibration_method, + calibration_cache_path=args.calibration_cache_path, + calibration_shapes=args.calibration_shapes, + calibration_eps=args.calibration_eps, + op_types_to_quantize=args.op_types_to_quantize, + op_types_to_exclude=args.op_types_to_exclude, + nodes_to_quantize=args.nodes_to_quantize, + nodes_to_exclude=args.nodes_to_exclude, + use_external_data_format=args.use_external_data_format, + keep_intermediate_files=args.keep_intermediate_files, + output_path=args.output_path, + verbose=args.verbose, + trt_plugins=args.trt_plugins, + trt_plugins_precision=args.trt_plugins_precision, + high_precision_dtype=args.high_precision_dtype or default_high_precision_dtype, + mha_accumulation_dtype=args.mha_accumulation_dtype, + disable_mha_qdq=args.disable_mha_qdq, + dq_only=args.dq_only, + use_zero_point=args.use_zero_point, + ) + + +if __name__ == "__main__": + main() diff --git a/modelopt/onnx/quantization/calib_utils.py b/modelopt/onnx/quantization/calib_utils.py new file mode 100644 index 0000000000..1bb8b4b7be --- /dev/null +++ b/modelopt/onnx/quantization/calib_utils.py @@ -0,0 +1,151 @@ +# 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. + +"""Provides basic calibration utils.""" + +import struct +from typing import Union + +import numpy as np +import onnx +from onnxruntime.quantization.calibrate import CalibrationDataReader + +from modelopt.onnx.utils import ( + gen_random_inputs, + get_input_names, + get_input_shapes, + parse_shapes_spec, +) + +CalibrationDataType = Union[np.ndarray, dict[str, np.ndarray]] + + +class CalibrationDataProvider(CalibrationDataReader): + """Calibration data provider class.""" + + def __init__( + self, onnx_path: str, calibration_data: CalibrationDataType, calibration_shapes: str = None + ): + """Intializes the data provider class with the calibration data iterator. + + Args: + onnx_path: Path to the ONNX model. + calibration_data: Numpy data to calibrate the model. + Ex. If a model has input shapes like {"sample": (2, 4, 64, 64), "timestep": (1,), + "encoder_hidden_states": (2, 16, 768)}, the calibration data should have dictionary + of tensors with shapes like {"sample": (1024, 4, 64, 64), "timestep": (512,), + "encoder_hidden_states": (1024, 16, 768)} to calibrate with 512 samples. + calibration_shapes: A string representing the shape of each input tensors for one calibration step. + If the shape is not provided for an input tensor, the shape is inferred from the onnx model directly, + with all the unknown dims filled with 1. + """ + onnx_model = onnx.load(onnx_path) + input_names = get_input_names(onnx_model) + input_shapes = {} if calibration_shapes is None else parse_shapes_spec(calibration_shapes) + inferred_input_shapes = get_input_shapes(onnx_model) + for name in input_names: + if name not in input_shapes: + input_shapes[name] = inferred_input_shapes[name] + + # Validate calibration data against expected inputs by the model + if isinstance(calibration_data, np.ndarray): + assert len(input_names) == 1, "Calibration data has only one tensor." + calibration_data = {input_names[0]: calibration_data} + elif isinstance(calibration_data, dict): + assert len(input_names) == len( + calibration_data + ), "Model input count and calibration data doesn't match." + for input_name in input_names: + assert input_name in calibration_data + else: + raise ValueError( + f"calibration data must be numpy array or dict, got {type(calibration_data)}" + ) + + # Create list of model inputs with appropriate batch size + # So that we can create an input iterator + n_itr = int(calibration_data[input_names[0]].shape[0] / input_shapes[input_names[0]][0]) + self.calibration_data_list = [{}] * n_itr + for input_name in input_names: + for idx, calib_data in enumerate( + np.array_split(calibration_data[input_name], n_itr, axis=0) + ): + self.calibration_data_list[idx][input_name] = calib_data + + self.calibration_data_reader = iter(self.calibration_data_list) + + def get_next(self): + """Returns the next available calibration input from the reader.""" + return next(self.calibration_data_reader, None) + + def get_first(self): + """Returns the first calibration input from the reader without incrementing the iterator. + + This is useful when doing a test run for the session. + """ + assert len(self.calibration_data_list) > 0, "Calibration data list is empty!" + return self.calibration_data_list[0] + + +class RandomDataProvider(CalibrationDataReader): + """Calibration data reader class with random data provider.""" + + def __init__( + self, onnx_model: Union[str, onnx.onnx_pb.ModelProto], calibration_shapes: str = None + ): + """Initializes the data reader class with random calibration data.""" + if isinstance(onnx_model, str): + onnx_path = onnx_model + onnx_model = onnx.load(onnx_path) + self.calibration_data_list: list[dict[str, np.ndarray]] = [ + gen_random_inputs(onnx_model, calibration_shapes) + ] + self.calibration_data_reader = iter(self.calibration_data_list) + + def get_next(self): + """Returns the next available calibration input from the reader.""" + return next(self.calibration_data_reader, None) + + def get_first(self): + """Returns the first calibration input from the reader without incrementing the iterator. + + This is useful when doing a test run for the session. + """ + assert len(self.calibration_data_list) > 0, "Calibration data list is empty!" + return self.calibration_data_list[0] + + +def import_scales_from_calib_cache(cache_path: str) -> dict[str, float]: + """Reads TensorRT calibration cache and returns as dictionary. + + Args: + cache_path: Calibration cache path. + + Returns: + Dictionary with scales in the format {tensor_name: float_scale}. + """ + with open(cache_path, "r") as f: + scales_dict = {} + lines = f.readlines() + for i, line in enumerate(lines): + if i > 0: # Skips the first line (i.e., TRT-8501-EntropyCalibration2) + layer_name, hex_value = line.replace("\n", "").split(": ") + try: + scale = struct.unpack("!f", bytes.fromhex(hex_value))[0] + scales_dict[layer_name + "_scale"] = scale + except Exception: + raise ValueError(f"Scale value for tensor {layer_name} was not found!") + + return scales_dict diff --git a/modelopt/onnx/quantization/extensions.py b/modelopt/onnx/quantization/extensions.py new file mode 100644 index 0000000000..8dd2b9bf61 --- /dev/null +++ b/modelopt/onnx/quantization/extensions.py @@ -0,0 +1,34 @@ +# 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. + +"""Module to load C++ extensions.""" + +import os +import sys +import warnings + +import cppimport + +try: + print("Loading extension modelopt_round_and_pack_ext...\n") + path = os.path.join(os.path.dirname(__file__), "src") + sys.path.append(path) + round_and_pack_ext = cppimport.imp("modelopt_round_and_pack_ext") + sys.path.remove(path) +except Exception as e: + warnings.warn( + f"{e}\nUnable to load `modelopt_round_and_pack_ext', falling back to python based optimized version." + ) + round_and_pack_ext = None diff --git a/modelopt/onnx/quantization/fp8.py b/modelopt/onnx/quantization/fp8.py new file mode 100755 index 0000000000..3e883bc19b --- /dev/null +++ b/modelopt/onnx/quantization/fp8.py @@ -0,0 +1,307 @@ +# 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. + +"""Performs FP8 GEMM only quantization of an ONNX model, and returns the ONNX ModelProto.""" + +import logging +import os +import tempfile +from functools import reduce + +import numpy as np +import onnx +import onnx_graphsurgeon as gs +from onnx import numpy_helper +from onnx_graphsurgeon.ir.graph import Graph +from onnxmltools.utils.float16_converter import convert_float_to_float16 +from onnxruntime.quantization import CalibrationMethod +from onnxruntime.quantization.calibrate import CalibrationDataReader + +from modelopt.onnx.quantization.graph_utils import ( + build_non_residual_input_map, + convert_fp16_io, + expand_node_names_from_patterns, + find_nodes_to_exclude, + get_tensor_producer_nodes, + insert_fp8_mha_casts, + remove_partial_input_qdq, +) +from modelopt.onnx.quantization.ort_patching import _quantize_static as quantize_static +from modelopt.onnx.quantization.ort_utils import configure_ort + + +def _find_unsupported_fp8_convs_to_exclude(graph: Graph): + """Find unsupported FP8 Conv nodes to exclude. + + The input and output channel alignment requirement for FP8 + conv kernels for input and output type FP8E4M3 should be both 16. + The filter size for FP8 conv kernels should be less than 32. + + Args: + graph: Onnx model graph. + + Returns: + List of Conv nodes. + """ + unsupported_conv_nodes = [] + for node in graph.nodes: + if node.op == "Conv": + weight = node.inputs[1] + assert ( + 3 <= len(weight.shape) <= 5 + ), f"Invalid weight shape {weight.shape}. Only 1D, 2D, and 3D convolutions are supported." + output_channel = weight.shape[0] + input_channel = weight.shape[1] + if output_channel % 16 != input_channel % 16: + logging.info(f"Found unpaddable conv for FP8: {node.name}") + unsupported_conv_nodes.append(node.name) + continue + + filter_size = reduce(lambda x, y: x * y, weight.shape[2:]) + if filter_size > 32: + logging.info(f"Found large filter conv for FP8: {node.name}") + unsupported_conv_nodes.append(node.name) + + return unsupported_conv_nodes + + +def int8_to_fp8(onnx_path: str, verbose: bool = False) -> onnx.onnx_pb.ModelProto: + """Converts the INT8 quantized model to FP8 quantized model. + + Note. This conversion works only for max calibrated INT8 models. + + Args: + onnx_path: Path to the INT8 quantized ONNX model. + verbose: Whether to print verbose logs or not. + + Returns: + FP8 quantized ONNX model. + """ + onnx_model = onnx.load(onnx_path) + graph = onnx_model.graph + initializers = graph.initializer + tensor_producers = get_tensor_producer_nodes(graph) + processed_tensor = set() + initializer_indices = { + initializer.name: idx for idx, initializer in enumerate(graph.initializer) + } + + def _int8_scale_to_fp8_scale(scale: np.ndarray, scale_name: str): + np_scale = onnx.numpy_helper.to_array(scale) + np_fp8_scale = (np_scale * 448.0) / 127.0 + dtype = onnx.helper.tensor_dtype_to_np_dtype(scale.data_type) + return numpy_helper.from_array(np_fp8_scale.astype(dtype), scale_name) + + def _convert(node: onnx.onnx_ml_pb2.NodeProto): + if verbose: + logging.info(f"Processing {node.name}") + + scale_name = node.input[1] + zero_point_name = node.input[2] + + if scale_name not in processed_tensor: + scale_idx = initializer_indices.get(scale_name, None) + if scale_idx is not None: + scale = initializers[scale_idx] + fp8_scale = _int8_scale_to_fp8_scale(scale, scale_name) + initializers[scale_idx].CopyFrom(fp8_scale) + else: + producer_node = tensor_producers[scale_name] + scale = producer_node.attribute[0].t + fp8_scale = _int8_scale_to_fp8_scale(scale, scale_name) + producer_node.attribute[0].t.CopyFrom(fp8_scale) + processed_tensor.add(scale_name) + + if zero_point_name not in processed_tensor: + zero_point_idx = initializer_indices.get(zero_point_name, None) + assert ( + zero_point_idx is not None + ), f"Expected '{zero_point_name}' to be found in 'graph.initializer', but it was not present." + zero_point = initializers[zero_point_idx] + dtype = onnx.helper.tensor_dtype_to_np_dtype(zero_point.data_type) + vals = np.array(zero_point.int32_data, dtype=dtype).tolist() + + np_zero_point = onnx.helper.make_tensor( + zero_point_name, onnx.TensorProto.FLOAT8E4M3FN, zero_point.dims, vals + ) + initializers[zero_point_idx].CopyFrom(np_zero_point) + processed_tensor.add(zero_point_name) + + # Iterate through the nodes and convert the scales and zero points + for node in graph.node: + if node.op_type in ["DequantizeLinear", "QuantizeLinear"]: + _convert(node) + + return onnx_model + + +def upgrade_opset_21(onnx_model: onnx.onnx_pb.ModelProto) -> onnx.onnx_pb.ModelProto: + """Modifies the ONNX graph such that it follows the opset 21 requirements. + + This is necessary for FP8+FP16 quantization since FP8 QuantizeLinear/DequantizeLinear ops do not support FP16 + scaling factors until opset 21. + """ + graph = gs.import_onnx(onnx_model) + + for node in graph.nodes: + # QuantizeLinear/DequantizeLinear op with FP16 scales are only supported with empty domain + # and opset_import version=21. + if node.op == "QuantizeLinear" or node.op == "DequantizeLinear": + node.domain = "" + + # ReduceMean op no longer has "axes" attribute in opset 21. Instead, it should be the second input tensor. + if node.op == "ReduceMean" and "axes" in node.attrs: + axes = gs.Constant( + name=node.name + "_axes", values=np.array(node.attrs["axes"], dtype=np.int64) + ) + del node.attrs["axes"] + node.inputs.append(axes) + + onnx_model = gs.export_onnx(graph) + + # Set opset_import version to 21. + for opset_import in onnx_model.opset_import: + if opset_import.domain == "": + opset_import.version = 21 + + return onnx_model + + +def quantize( + onnx_path: str, + calibration_method: str = "max", + calibration_data_reader: CalibrationDataReader = None, + calibration_cache_path: str = None, + calibration_shapes: str = None, + calibration_eps: list[str] = ["cuda:0", "cpu", "trt"], + op_types_to_quantize: list[str] = None, + op_types_to_exclude: list[str] = None, + nodes_to_quantize: list[str] = None, + nodes_to_exclude: list[str] = None, + use_external_data_format: bool = True, + intermediate_generated_files: list[str] = [], + verbose: bool = False, + trt_extra_plugin_lib_paths: str = None, + high_precision_dtype: str = "fp16", + mha_accumulation_dtype: str = "fp32", +) -> onnx.onnx_pb.ModelProto: + """Applies FP8 GEMM only quantization to an ONNX file. + + Currently, ['Conv', 'Gemm', 'MatMul'] quantization is supported. + """ + logging.info("Quantization Mode: fp8") + if calibration_method != "max": + raise RuntimeError("Only the max calibration method is supported for FP8 quantization.") + + # Load the onnx graph + onnx_model = onnx.load(onnx_path, load_external_data=use_external_data_format) + graph = gs.import_onnx(onnx_model) + graph.toposort() + + # The quantizable op types for FP8 are limited to Conv, Gemm, and Matmul + fp8_supported_op_types = ["Gemm", "MatMul", "Conv"] + op_types_to_quantize = op_types_to_quantize or fp8_supported_op_types + if not set(op_types_to_quantize) <= set(fp8_supported_op_types): + raise RuntimeError( + f"Unsupported op types in fp8 mode: '{set(op_types_to_quantize) - set(fp8_supported_op_types)}'" + ) + + # Change the default configuration of ORT quantization + op_types = set([node.op for node in graph.nodes]) + trt_guided_options, _ = configure_ort( + list(op_types), op_types_to_quantize, trt_extra_plugin_lib_paths, calibration_eps + ) + logging.info( + "Quantizable op types in the model:" + f" {[t for t in op_types_to_quantize if t in op_types]}" + ) + + # Collect node names to include in quantization + no_quantize_inputs = [] + nodes_to_quantize = expand_node_names_from_patterns(graph, nodes_to_quantize) + if not nodes_to_quantize: + nodes_to_quantize = [node.name for node in graph.nodes if node.op in op_types_to_quantize] + _, no_quantize_inputs = build_non_residual_input_map(graph) + if no_quantize_inputs: + op_types_to_quantize.append("Add") + add_nodes = [dst.name for _, dst, _ in no_quantize_inputs] + nodes_to_quantize.extend(add_nodes) + + # Collect node names to exclude from quantization + nodes_to_exclude = find_nodes_to_exclude(graph, nodes_to_exclude, op_types_to_exclude) + nodes_to_exclude.extend(_find_unsupported_fp8_convs_to_exclude(graph)) + + # Update the list of nodes to quantize + nodes_to_quantize = [ + node_name for node_name in nodes_to_quantize if node_name not in nodes_to_exclude + ] + + logging.info(f"Total number of nodes: {len(graph.nodes)}") + if not nodes_to_quantize: + logging.info( + "No node or node type is selected for quantization or model does not have them!" + ) + return + elif verbose: + logging.info(f"Selected nodes to quantize: {nodes_to_quantize}") + + # Create a temp file for intermediate model + tmp_onnx_file, tmp_onnx_path = tempfile.mkstemp(suffix=".onnx") + os.close(tmp_onnx_file) + + # Quantize in INT8 mode using ORT's MinMax calibration method, with + # ActivationSymmetric as True, which is equivalent to max calibration + quantize_static( + onnx_path, + tmp_onnx_path, + calibration_data_reader, + op_types_to_quantize=op_types_to_quantize, + nodes_to_quantize=nodes_to_quantize, + per_channel=True, + extra_options=trt_guided_options, + use_external_data_format=use_external_data_format, + calibrate_method=CalibrationMethod.MinMax, + ) + intermediate_generated_files.append(tmp_onnx_path) + if use_external_data_format: + intermediate_generated_files.append(tmp_onnx_path + ".data") + + # Post-processing of the onnx model after ORT quantization + onnx_model = int8_to_fp8(tmp_onnx_path, verbose) + graph = gs.import_onnx(onnx_model) + remove_partial_input_qdq(graph, no_quantize_inputs) + onnx_model = gs.export_onnx(graph) + + if high_precision_dtype == "fp16": + # We need to convert float to float16 so as to speed up layers like LayerNorm or GroupNorm. + logging.info("Converting float tensors to float16") + graph = gs.import_onnx(onnx_model) + convert_fp16_io(graph) + onnx_model = gs.export_onnx(graph) + onnx_model = convert_float_to_float16( + onnx_model, keep_io_types=True, disable_shape_infer=True + ) + + # We need to convert the ONNX model to opset 21 since FP8 QuantizeLinear/DequantizeLinear ops do not support + # FP16 scaling factors until opset 21. + onnx_model = upgrade_opset_21(onnx_model) + + if mha_accumulation_dtype == "fp32": + # Insert Cast nodes in MHA's BMM1 and BMM2's input and output tensors because + # The compiler only has FP32 accumulation kernels for FP8 MHAs. + logging.info("Inserting Cast nodes to enable FP8+FP16 MHA") + onnx_model = insert_fp8_mha_casts(onnx_model) + + return onnx_model diff --git a/modelopt/onnx/quantization/graph_utils.py b/modelopt/onnx/quantization/graph_utils.py new file mode 100755 index 0000000000..6827882a89 --- /dev/null +++ b/modelopt/onnx/quantization/graph_utils.py @@ -0,0 +1,1098 @@ +# 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. + +"""Provides ONNX graph related utils for QDQ placement.""" + +import logging +import os +import re +from collections import defaultdict +from typing import Optional, Union + +import numpy as np +import onnx +import onnx_graphsurgeon as gs +from onnx_graphsurgeon.ir.graph import Graph +from onnx_graphsurgeon.ir.node import Node +from onnx_graphsurgeon.ir.tensor import Constant, Tensor, Variable +from onnxruntime.quantization.calibrate import CalibrationDataReader + +from modelopt.onnx.op_types import is_copy_op, is_linear_op +from modelopt.onnx.quantization.ort_utils import create_inference_session +from modelopt.onnx.utils import find_lowest_common_ancestor, get_child_nodes, get_parent_nodes + + +def is_const_input(tensor: Tensor) -> bool: + """Returns whether the given tensor is an initializer or produced by const-foldable nodes.""" + if isinstance(tensor, Constant): + return True + + # Tensor is a graph input variable + if len(tensor.inputs) == 0: + return False + + producer_node = tensor.inputs[0] # Generally tensors has single producer + if producer_node.op in ["Constant", "Identity"]: + return True + + # Second axes input to Squeeze/Unsqueeze is a constant, we need to check the first input + if producer_node.op in ["Squeeze", "Unsqueeze"] and is_const_input(producer_node.inputs[0]): + return True + + # Const -> Clip -> Exp -> Mul pattern matching for swin_v2 + if producer_node.op == "Exp": + clip_node = producer_node.i() + if clip_node.op == "Clip" and has_const_input(clip_node): + return True + + return False + + +def has_const_input(node: Node) -> bool: + """Returns whether the given node has any constant input.""" + for tensor in node.inputs: + if is_const_input(tensor): + return True + + return False + + +def has_path_type( + node: Node, + graph: Graph, + path_type: list[str], + is_forward: bool, + wild_card_types: list[str] = [], + path_nodes: list[Node] = [], +) -> bool: + """Checks if the given node is start/end of a given forward/backward path type. + + Note, Path can be forward or backward wrt a node depending on the next level nodes. + Additionally, this method can work with optional nodes and collect the traversed path. + + Args: + node: Start node of the path. + graph: ONNX model graph. + path_type: Path types to match from the given node. + is_forward: Whether to match forward or backward path. + wild_card_types: Wild card types, these type of nodes are skipped and not matched with the path_type. + path_nodes: Accumulated nodes in the matched path. + + Returns: + Bool, whether the given node is start/end of the given forward/backward path type. + """ + optional_path_types = ["BiasAdd", "ConstMul"] + if not path_type: + # All types matched + return True + + # Current node type and special type conversion for optional BiasAdd and ConstMul + # Note, matching path with Add/Mul type nodes with const input will fail + node_type = node.op + if node_type == "Add" and has_const_input(node): + node_type = "BiasAdd" + elif node_type == "Mul" and has_const_input(node): + node_type = "ConstMul" + + # Check if current non-wild node type does not match the expected path type + # And if path type is not optional (ex. BiasAdd) + is_match = (node_type == path_type[0]) or (node.op == path_type[0]) + is_wild_match = node_type in wild_card_types + if not is_match and not is_wild_match and (path_type[0] not in optional_path_types): + return False + + # Add current node name in the path + if is_match: + path_nodes.append(node) + + # If current node type matches the expected path type or path type is optional (ex. BiasAdd), we have a type match + # Update the remaining path types to match + next_path_type = path_type[:] + + # Non-repeatable optional types should be consumed + if is_match or (path_type[0] in ["BiasAdd", "ConstMul"]): + next_path_type = path_type[1:] + + # If current node is not wild card and didn't match, go ahead and match with the + # remaining path types starting with the current node + if not is_match and not is_wild_match: + assert path_type[0] in optional_path_types + return has_path_type( + node, + graph, + next_path_type, + is_forward, + wild_card_types, + path_nodes, + ) + + if is_forward: + next_level_nodes = get_child_nodes(node) + else: + next_level_nodes = get_parent_nodes(node) + + # Check if any child (forward path) or parent (backward path) can match the remaining path types + for next_node in next_level_nodes: + sub_path = [] + if has_path_type(next_node, graph, next_path_type, is_forward, wild_card_types, sub_path): + path_nodes.extend(sub_path) + return True + + # Path type matches if there is no remaining types to match + return not next_path_type + + +def get_fusible_backbone(node: Node, graph: Graph) -> Optional[Node]: + """Returns the linear backbone node for a given node if it matches the pattern. + + TensorRT fuses convolution with BN, Relu etc. when in some specific pattern. + This rule tries to match some of those patterns. + Note. BiasAdd and ConstMul are optional in path types. + + Args: + node: Start node of the pattern. + graph: ONNX model graph. + + Returns: + Backbone node of the given node, None if not found. + """ + + def _get_backbone(root: Node): + if root.op == "Conv": + return root + + for tensor in root.inputs: + if not isinstance(tensor, Constant): + parent_node = tensor.inputs[0] + bb = _get_backbone(parent_node) + if bb: + return bb + + fusible_linear_path_types = [ + # ["Sigmoid", "Conv"], # With following Mul + # ["Resize", "Relu", "Conv"], # Note. this causes regression in MTL_v1 + ["BiasAdd", "ConstMul", "Conv"], + ["Relu", "BiasAdd", "ConstMul", "Conv"], + ["BatchNormalization", "BiasAdd", "Conv"], + ["Relu", "BatchNormalization", "BiasAdd", "Conv"], + ] + for idx, path_type in enumerate(fusible_linear_path_types): + if has_path_type(node, graph, path_type, is_forward=False, wild_card_types=[]): + return _get_backbone(node) + + return None + + +def get_tensor_producer_nodes( + graph: onnx.onnx_ml_pb2.GraphProto, +) -> dict[str, onnx.onnx_ml_pb2.NodeProto]: + """Returns a dictionary of tensor name and their producer node object mapping. + + Note. we create a special Root type node as external inputs producer for ease of implementation. + + Args: + graph: ONNX model graph. + + Returns: + Dictionary, key is tensor name and value is their producer node object + """ + # Create a dictionary to store tensor producer nodes + tensor_producers = defaultdict(None) + + # Special Root type producer node + root_node = onnx.helper.make_node( + op_type="Root", + inputs=[], + outputs=[i.name for i in graph.input], + name="root_0", + ) + + input_names = [graph_input.name for graph_input in graph.input] + initializer_names = [initializer.name for initializer in graph.initializer] + external_input_names = list(np.setdiff1d(input_names, initializer_names)) + + # Note. We are marking external inputs as non-constant by adding a parent, + # so that we can quantize the first node of the graph if appropriate + for graph_input in external_input_names: + tensor_producers[graph_input] = root_node + + # Traverse the graph to find producer nodes for each tensor + for node in graph.node: + for output_name in node.output: + tensor_producers[output_name] = node + + return tensor_producers + + +def get_tensor_consumer_nodes( + graph: onnx.onnx_ml_pb2.GraphProto, +) -> dict[str, list[onnx.onnx_ml_pb2.NodeProto]]: + """Returns a dictionary of tensor name and their consumer node object mapping. + + Args: + graph: ONNX model graph. + + Returns: + Dictionary, key is tensor name and value is their consumer node object + """ + # Create a dictionary to store tensor consumer nodes + tensor_consumers = defaultdict(list) + + # Traverse the graph to find consumer nodes for each tensor + for node in graph.node: + for input_name in node.input: + tensor_consumers[input_name].append(node) + + return tensor_consumers + + +def filter_quantizable_kgen_heads( + cask_fusible_partitions: list[list[Node]], + kgen_partitions: list[list[Node]], + quantizable_op_types: list[str], +) -> tuple[list[Node], list[tuple[Node, Node, str]]]: + """Returns the list of kgen head names if it follows a CASK partition.""" + cask_partition_nodes = set() + for partition in cask_fusible_partitions: + cask_partition_nodes.update([node.name for node in partition]) + + cask_partition_heads = [partition[0] for partition in cask_fusible_partitions] + + def _is_following_cask_partition(node: Node): + # Checking if cask fusible partition can be reached backward + # ignoring the copy ops + if node.name in cask_partition_nodes: + return True + + if not is_copy_op(node.op): + return False + + for parent in get_parent_nodes(node): + if _is_following_cask_partition(parent): + return True + + return False + + def _has_other_quantizable_consumer( + tensor: Tensor, quantizable_kgen_heads: list[Node], head_name: str + ): + # Note. this is kinda approximate analysis, + # all quantizable kgen heads may haven't got discovered yet + quantizable_ops = [node.name for node in cask_partition_heads + quantizable_kgen_heads] + + # Look for other quantizable consumer than the current kgen head + if head_name in quantizable_ops: + quantizable_ops.remove(head_name) + + for consumer in tensor.outputs: + if consumer.name in quantizable_ops: + return True + + return False + + quantizable_kgen_heads = [] + no_quantize_inputs = [] # list of tuple [(src_node_name, dst_node_name, input_name), ...] + output_quantization_candidates = [ + "AveragePool", + "BatchNormalization", + "GlobalAveragePool", + "MaxPool", + ] + + for partition in kgen_partitions: + head_node = partition[0] + # Check if partition head is of default quantizable type + if head_node.op not in quantizable_op_types: + continue + + # If the node has cost input, do not quantize + if has_const_input(head_node): + continue + + head_parents = get_parent_nodes(head_node) + no_quantize_inputs_of_head = [] + has_quantizable_input = False + + # Check each of the parent (input producer for partition head) + # or predecessor nodes and see if output quantization is needed for them + # and decide which input of kgen head needs quantization + for parent in head_parents: + # If the head is consuming output of any quantizable op, then it is quantizable + if _is_following_cask_partition(parent) or parent.op in output_quantization_candidates: + quantizable_kgen_heads.append(partition[0]) + has_quantizable_input = True + # If the input from the current parent has no other quantizable consumer, do not quantize that input + elif not _has_other_quantizable_consumer( + parent.outputs[0], quantizable_kgen_heads, head_node.name + ): + no_quantize_inputs_of_head.append((parent, partition[0], parent.outputs[0].name)) + + # If at least one input of Add is quantizable, collect if there is any non-quantizable inputs + if head_node.op == "Add" and has_quantizable_input: + no_quantize_inputs.extend(no_quantize_inputs_of_head) + + return quantizable_kgen_heads, no_quantize_inputs + + +def classify_partition_nodes( + partitions: list[list[Node]], +) -> tuple[list[Node], list[Node], list[tuple[Node, Node, str]]]: + """We should partially quantize the partition nodes with inputs outside of the partition. + + Args: + partitions: Partitions created by modelopt ptq algo. + + Returns: + List of non-quantizable nodes. + List of quantizable nodes. + List of partially-quantizable inputs with non-quantizable input info as (src, dst, input_name) + """ + non_quantizable_partition_nodes = [] # list of Node [node1, ...] + quantizable_partition_nodes = [] # list of Node [node1, ...] + no_quantize_inputs = [] # list of tuple [(src_node, dst_node, input_name), ...] + + for partition in partitions: + partition_root_type = partition[0].op + assert is_linear_op(partition_root_type) + + # Collect tensor names produced by partition nodes + partition_node_outputs = [] + for node in partition: + for node_output in node.outputs: + partition_node_outputs.append(node_output.name) + + for node in partition: + has_external_inputs = False + internal_inputs = [] # Keeps (producer, consumer, tensor) + for tensor in node.inputs: + if is_const_input(tensor): + continue + + # If a KGEN op has external non-constant input, it is considered partially quantizable + if tensor.name not in partition_node_outputs: + # partition heads will be fully quantizable and added + has_external_inputs = True + else: + producer_node = tensor.inputs[0] + # format: source, target, input + # Note. it might happen that this node was not quantized + # We just ignore it from no_quantize_inputs list in post-processing + internal_inputs.append((producer_node, node, tensor.name)) + + if not has_external_inputs: + non_quantizable_partition_nodes.append(node) + elif has_external_inputs and internal_inputs: + no_quantize_inputs.extend(internal_inputs) + else: + # partition head is quantizable + quantizable_partition_nodes.append(node) + + return non_quantizable_partition_nodes, quantizable_partition_nodes, no_quantize_inputs + + +def build_non_residual_input_map( + graph: Graph, +) -> tuple[dict[str, str], list[tuple[Node, Node, str]]]: + """Builds a map of non-residual Add input name to the Add node name from the given graph. + + This assumes that the Add layer only has 2 inputs. + + We will refer to a subgraph which has a Convolution node with a single output that is summed (element-wise) + with another non-constant input-tensor as a “residual-add” subgraph, because it occurs in modern + convnets that use residual connections. + + Args: + graph: Onnx model graph. + + Returns: + Dictionary of Add node names vs their non-residual input name. + List of partially-quantizable inputs with non-quantizable input info as (src, dst, input_name) + """ + non_residual_inputs = {} + no_quantize_inputs = [] + for node in graph.nodes: + if node.op in ["Add"]: + # Add nodes with constant or graph input does not have non-residual input + # Here, A = node.inputs[0], B = node.inputs[1] and A.inputs means producer nodes of A + # TODO: make this check a util? + if ( + has_const_input(node) + or len(node.inputs[0].inputs) == 0 + or len(node.inputs[1].inputs) == 0 + ): + non_residual_inputs[node.name] = None + continue + + input1_producer = node.i(0, 0) + input2_producer = node.i(1, 0) + + backbone1 = get_fusible_backbone(input1_producer, graph) + backbone2 = get_fusible_backbone(input2_producer, graph) + + # Input in the longest path to LCA is the non-residual input + lca, d1, d2 = find_lowest_common_ancestor(input1_producer, input2_producer) + + # Generally if both the inputs have a backbone then both backbones are of the same type + if backbone1 and backbone2: + if backbone1 == backbone2: + non_residual_inputs[node.name] = None + continue + + assert backbone1.op == backbone2.op, ( + f"{backbone1.name} and {backbone2.name} are different types of backbone for" + f" {node.name}!" + ) + if d1 > d2: + non_residual_inputs[node.name] = node.inputs[0].name + no_quantize_inputs.append((input1_producer, node, node.inputs[0].name)) + else: + non_residual_inputs[node.name] = node.inputs[1].name + no_quantize_inputs.append((input2_producer, node, node.inputs[1].name)) + elif backbone1: + # ConvNext pattern + # Conv ---------------------- add + # \---- non backbone---/ + # This case LCA being backbone itself is not residual Add case. + if lca and lca == backbone1.name: + # Not a residual Add node + non_residual_inputs[node.name] = None + else: + non_residual_inputs[node.name] = node.inputs[0].name + no_quantize_inputs.append((input1_producer, node, node.inputs[0].name)) + elif backbone2: + if lca and lca == backbone2.name: + # Not a residual Add node + non_residual_inputs[node.name] = None + else: + non_residual_inputs[node.name] = node.inputs[1].name + no_quantize_inputs.append((input2_producer, node, node.inputs[1].name)) + else: + # Not a residual Add node + non_residual_inputs[node.name] = None + + return non_residual_inputs, no_quantize_inputs + + +def remove_partial_input_qdq( + graph: Graph, + no_quantize_inputs: list[tuple[Node, Node, str]], +) -> None: + """Modifies the onnx model by removing QDQ nodes from the marked inputs, ex. non-residual inputs etc. + + Args: + graph: Onnx model graph. + no_quantize_inputs: List non-quantizable input info as (src, dst, input_name) + """ + logging.info("Deleting QDQ nodes from marked inputs to make certain operations fusible ...") + graph_nodes = {node.name: node for node in graph.nodes} + for source, target, non_qdq_input_name in no_quantize_inputs: + # Note. no_quantize_inputs objects are from non-quantized input graph + # we are deleting some QDQ from the new quantized output graph + source_node = graph_nodes[source.name] + try: + dq_node = source_node.o().o() + except Exception: + # Reached end of the graph + continue + if dq_node.op == "DequantizeLinear": + dq_node = dq_node.outputs[0] # source_node->Q->DQ->target_node + while len(dq_node.outputs): + # Find the input index in the target connecting with source_node + target_input_idx_arr = [ + idx + for idx, inp in enumerate(dq_node.outputs[0].inputs) + if inp.name == dq_node.name + ] + target_input_idx = target_input_idx_arr[0] if target_input_idx_arr else 0 + + # Connect the output of source_node with the outputs of DQ until DQ is not connected to any other + # layers. Note that when a connection is removed, this is also deleted from dq_node.outputs, thus + # why we keep iterating over the same idx=0 in dq_node.outputs[0]. + dq_node.outputs[0].inputs[target_input_idx] = source_node.outputs[0] + + graph.cleanup() + graph.toposort() + + +def _find_nodes_from_op_types_to_exclude(graph: Graph, op_types_to_exclude=None) -> list[str]: + nodes_to_exclude = [] + if op_types_to_exclude: + nodes_to_exclude = [node.name for node in graph.nodes if node.op in op_types_to_exclude] + return nodes_to_exclude + + +def expand_node_names_from_patterns( + graph: Union[onnx.onnx_pb.GraphProto, Graph], name_patterns: list[str] +) -> list[str]: + """Expand the node names from the given patterns.""" + node_list = getattr(graph, "nodes", None) or getattr(graph, "node", None) or [] + matched_node_names = [] + for pattern in name_patterns: + for node in node_list: + if re.match(pattern, node.name): + matched_node_names.append(node.name) + return matched_node_names + + +def find_nodes_to_exclude( + graph: Graph, nodes_to_exclude: list[str], op_types_to_exclude: list[str] +): + """Find the node names from the ONNX graph which matches user's exclusion patterns.""" + nodes_to_exclude = nodes_to_exclude or [] + nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude) + nodes_to_exclude.extend(_find_nodes_from_op_types_to_exclude(graph, op_types_to_exclude)) + + # Remove duplicates from the exclusion list + return [*set(nodes_to_exclude)] + + +def get_extended_model_outputs( + onnx_path: str, + extended_model: onnx.onnx_pb.ModelProto, + use_external_data_format: bool, + intermediate_generated_files: list[str], + calibration_data_reader: CalibrationDataReader, + calibration_eps: list[str], +) -> dict[str, np.ndarray]: + """Run one inference step on an onnx model which has some intermediate tensor marked as model outputs. + + The first calibration data is used for the dummy inference. This is useful when we want to know the shape of an + intermediate tensor given the calibration data. + + Args: + onnx_path: + Path to the original onnx model. + extended_model: + The onnx model with some intermediate tensors marked as model outputs. + use_external_data_format: + If not None, this path will be used to store the weights of the quantized model. + intermediate_generated_files: + List of intermediate generated files that will be deleted after quantization. + calibration_data_reader: + Calibration data reader for running inference. + calibration_eps: + Priority order for the execution providers (EP) to calibrate the model. + Any subset of ['cuda:x', 'cpu', 'trt'], where 'x' is the device id. + + Returns: a map with each output name pointed to the corresponding output numpy ndarray. + """ + # Get the first calibration input data. + inputs = calibration_data_reader.get_first() + + # Initialize ORT session. + if use_external_data_format: + extended_onnx_path = f"{onnx_path[:-5]}.extended.onnx" + extended_model_external_data_path = f"{onnx_path[:-5]}.extended.onnx_data" + onnx.save_model( + extended_model, + extended_onnx_path, + save_as_external_data=True, + location=os.path.basename(extended_model_external_data_path), + ) + intermediate_generated_files.append(extended_onnx_path) + intermediate_generated_files.append(extended_model_external_data_path) + session = create_inference_session(extended_onnx_path, calibration_eps) + else: + session = create_inference_session(extended_model.SerializeToString(), calibration_eps) + + # Run extended model's inference. + extended_model_output_names = [output.name for output in session.get_outputs()] + outputs = session.run(extended_model_output_names, inputs) + output_map = {name: output for name, output in zip(extended_model_output_names, outputs)} + + return output_map + + +def find_nodes_from_matmul_to_exclude( + onnx_path: str, + use_external_data_format: bool = False, + intermediate_generated_files: list[str] = None, + calibration_data_reader: CalibrationDataReader = None, + calibration_eps: list[str] = ["cuda:0", "cpu", "trt"], + verbose: bool = False, +) -> list[str]: + """Find MatMul nodes that meets gemv condition to exclude. + + Either of m or n in matmul is 1, this matmul cannot utilize + TensorCores. The perf of adding Q/DQ layers is not good in + TRT. Thus, in this case, do not add Q/DQ layers to this matmul. + + Args: + onnx_path: + Path to the onnx model. + use_external_data_format: + If not None, this path will be used to store the weights of the quantized model. + intermediate_generated_files: + List of intermediate generated files that will be deleted after quantization. + calibration_data_reader: + Calibration data reader for running inference. + calibration_eps: + Priority order for the execution providers (EP) to calibrate the model. + Any subset of ['cuda:x', 'cpu', 'trt'], where 'x' is the device id. + verbose: + If True, print the matmul nodes to exclude. + + Returns: + List of Nodes to exclude from quantization. + """ + model = onnx.load(onnx_path, load_external_data=use_external_data_format) + graph = gs.import_onnx(model) + + matmul_nodes = [] + for node in graph.nodes: + if node.op == "MatMul" or node.op == "Gemm": + matmul_nodes.append(node) + + if len(matmul_nodes) == 0: + return [] + + nodes_to_exclude = [] + + # Then, add each matmul output as model's extended outputs. + for matmul_node in matmul_nodes: + model.graph.output.extend([onnx.ValueInfoProto(name=matmul_node.outputs[0].name)]) + + # To get the shape info of the model, we run model inference once to get the shape info. + output_map = get_extended_model_outputs( + onnx_path, + model, + use_external_data_format, + intermediate_generated_files, + calibration_data_reader, + calibration_eps, + ) + + for matmul_node in matmul_nodes: + matmul_output = output_map[matmul_node.outputs[0].name] + if ( + len(matmul_output.shape) < 2 + or matmul_output.shape[-1] == 1 + or matmul_output.shape[-2] == 1 + ): + nodes_to_exclude.append(matmul_node.name) + + if verbose: + logging.info(f"Matmul nodes to exclude: {nodes_to_exclude}") + + return [*set(nodes_to_exclude)] + + +def find_nodes_from_mha_to_exclude( + onnx_path: str, + use_external_data_format: bool = False, + nodes_to_exclude: list[str] = None, + disable_mha_qdq: bool = False, + quantize_mode: str = "int8", + intermediate_generated_files: list[str] = None, + calibration_data_reader: CalibrationDataReader = None, + calibration_eps: list[str] = ["cuda:0", "cpu", "trt"], + verbose: bool = False, +) -> list[str]: + """Find MatMul nodes in MHA pattern to exclude. + + If disable_mha_qdq is set, don't add Q/DQ layers to MatMuls in MHA pattern. + else when quantize_mode == "fp8", if head_size > 256 or head_size % 16 != 0 or + mha doesn't meet fp8 fMHA v2 pattern, don't add Q/DQ layers to MatMuls in MHA pattern. + else when quantize_mode == "int8", if seq_len > 512, don't add Q/DQ layers + to MatMuls in MHA pattern. + + Args: + onnx_path: + Path to the onnx model. + use_external_data_format: + If not None, this path will be used to store the weights of the quantized model. + nodes_to_exclude: + List of Nodes to exclude from quantization. + disable_mha_qdq: + If True, all MHA's BMM1 and BMM2 will be added to nodes_to_exclude. + Else, each MHA will be checked whether to enable QDQ or not when is_fp8fp16 is True. + quantize_mode: + Quantization mode. One of 'int8' (default), 'int4' and 'fp8'. + intermediate_generated_files: + List of intermediate generated files that will be deleted after quantization. + calibration_data_reader: + Calibration data reader for running inference. + calibration_eps: + Priority list of execution providers (EP) for calibration. + verbose: + If True, print the matmul nodes to exclude. + + Returns: + List of Nodes to exclude from quantization. + """ + model = onnx.load(onnx_path, load_external_data=use_external_data_format) + graph = gs.import_onnx(model) + + mha_partitions = find_mha_partitions(graph) + if len(mha_partitions) == 0: + return nodes_to_exclude + + matmul_nodes_to_exclude = [] + if disable_mha_qdq: + for mha_partition in mha_partitions: + matmul_nodes_to_exclude.append(mha_partition[0].name) + matmul_nodes_to_exclude.append(mha_partition[2].name) + elif quantize_mode == "fp8" or quantize_mode == "int8": + # Add each BMM1's second input as BS1 model's extended outputs. + for mha_partition in mha_partitions: + bmm1_node = mha_partition[0] + model.graph.output.extend([onnx.ValueInfoProto(name=bmm1_node.inputs[1].name)]) + + # To get head_size and seq_len of MHA of the model, we run extended model inference once to get the shape info. + output_map = get_extended_model_outputs( + onnx_path, + model, + use_external_data_format, + intermediate_generated_files, + calibration_data_reader, + calibration_eps, + ) + + # For each MHA block, + # In quantize_mode == int8, if seq_len > 512, add bmm to nodes_to_exclude. + # In quantize_mode == fp8, if head_size > 256 or head_size % 16 != 0, add its bmm to nodes_to_exclude. + for mha_partition in mha_partitions: + bmm1_node = mha_partition[0] + softmax_node = mha_partition[1] + bmm1_input_name = bmm1_node.inputs[1].name + bmm1_input = output_map[bmm1_input_name] + seq_len = bmm1_input.shape[-1] + head_size = bmm1_input.shape[-2] + enable_mha_qdq = True + if quantize_mode == "int8": + if seq_len > 512: + enable_mha_qdq = False + elif quantize_mode == "fp8": + if head_size > 256 or head_size % 16 != 0: + enable_mha_qdq = False + else: + fp8_fmha_v2_pattern = match_fp8_mha_pattern(graph, softmax_node, False) + if len(fp8_fmha_v2_pattern) == 0: + enable_mha_qdq = False + if not enable_mha_qdq: + matmul_nodes_to_exclude.append(mha_partition[0].name) + matmul_nodes_to_exclude.append(mha_partition[2].name) + + if verbose: + logging.info(f"Matmul nodes From MHA to exclude: {matmul_nodes_to_exclude}") + + nodes_to_exclude.extend(matmul_nodes_to_exclude) + # Remove duplicates from the exclusion list + return [*set(nodes_to_exclude)] + + +def add_fp16_fp32_cast(onnx_path, custom_ops_to_cast_to_fp16): + """Adds cast_to_fp16 nodes to the inputs of a layer and cast_to_fp32 to the outputs.""" + name_dict = {} + + def _get_unique_name(old_name): + if old_name not in name_dict: + name_dict[old_name] = 0 + return old_name + name_dict[old_name] = name_dict[old_name] + 1 + return old_name + "_" + str(name_dict[old_name]) + + def _add_cast_node_inp(tensor, precision="fp16", suffix=""): + if precision == "fp16": + onnx_precision = int(onnx.TensorProto.FLOAT16) + np_precision = "float16" + else: + onnx_precision = int(onnx.TensorProto.FLOAT) + np_precision = "float32" + + cast_out = Variable( + name=_get_unique_name(tensor.name + f"_{precision}{suffix}"), + dtype=np_precision, + shape=tensor.shape, + ) + cast_node = Node( + op="Cast", + name=_get_unique_name(tensor.name + f"_cast_to_{precision}{suffix}"), + attrs={"to": onnx_precision}, + inputs=[tensor], + outputs=[cast_out], + ) + graph.nodes.append(cast_node) + return cast_out + + def _add_cast_node_out(tensor, inp_precision="fp16", out_precision="fp32", suffix=""): + cast_precision = ( + int(onnx.TensorProto.FLOAT16) + if out_precision == "fp16" + else int(onnx.TensorProto.FLOAT) + ) + np_precision = "float16" if inp_precision == "fp16" else "float32" + + cast_inp = gs.Variable( + name=_get_unique_name(tensor.name + f"_{inp_precision}{suffix}"), + dtype=np_precision, + shape=tensor.shape, + ) + cast_node = gs.Node( + op="Cast", + name=_get_unique_name(tensor.name + f"_cast_to_{out_precision}{suffix}"), + attrs={"to": cast_precision}, + inputs=[cast_inp], + outputs=[tensor], + ) + graph.nodes.append(cast_node) + return cast_inp + + graph = gs.import_onnx(onnx.load(onnx_path)) + castable_nodes = [n for n in graph.nodes if n.op in custom_ops_to_cast_to_fp16] + + for node in castable_nodes: + # Cast all inputs to FP16 + for inp_idx, inp in enumerate(node.inputs): + cast_out = _add_cast_node_inp(inp) + node.inputs[inp_idx] = cast_out + + # Cast all outputs from FP16 back to FP32 + for out_idx, out in enumerate(node.outputs): + cast_inp = _add_cast_node_out(out) + node.outputs[out_idx] = cast_inp + + graph.cleanup().toposort() + + new_onnx_path = onnx_path.replace(".onnx", "_castFP16.onnx") + onnx.save(gs.export_onnx(graph), new_onnx_path) + return new_onnx_path + + +def print_stat(graph: Graph, verbose: bool) -> None: + """Collect and print stats of the quantized model.""" + count = 0 + quantized_type_counts = {} + quantized_nodes = [] + output_names = [output_node.name for output_node in graph.outputs] + for node in graph.nodes: + for tensor in node.inputs: + if len(tensor.inputs) == 0: + continue + + producer_node = tensor.inputs[0] + if producer_node.op == "DequantizeLinear": + quantized_type_counts[node.op] = quantized_type_counts.get(node.op, 0) + 1 + quantized_nodes.append(node.name) + count += 1 + break + else: + # Sometimes "_DequantizeLinear_Output" is not suffix of the "DequantizeLinear" typed node, + # if that node is also in final model output. Ex. CLIP-ViT-L-14-opset16.onnx + assert tensor.name in output_names or producer_node.op != "DequantizeLinear" + + if verbose: + logging.info(f"Quantized nodes: {quantized_nodes}") + logging.info(f"Total number of quantized nodes: {count}") + logging.info(f"Quantized type counts: {quantized_type_counts}") + + +def find_mha_partitions(graph): + """Match MHA: BMM1 -> ... -> Softmax -> ... -> BMM2.""" + mha_chain_type = ["MatMul", "Softmax", "MatMul"] + wild_card_types = [ + "Div", + "Mul", + "ConstMul", + "Add", + "BiasAdd", + "Reshape", + "Transpose", + "Flatten", + "Cast", + ] + mha_partitions = [] + for node in graph.nodes: + if node.op == "MatMul": + mha_partition = [] + if has_path_type(node, graph, mha_chain_type, True, wild_card_types, mha_partition): + if ( + len(mha_partition) == 3 + and mha_partition[0].op == "MatMul" + and mha_partition[2].op == "MatMul" + ): + mha_partitions.append(mha_partition) + + return mha_partitions + + +def match_fp8_mha_pattern(graph: Graph, softmax_op: Node, has_fp8_qdq: bool) -> list[Node]: + """Match FP8 fMHA v2 with the given softmax_op. + + If has_fp8_qdq == True, we match this FP8 fMHA v2 pattern: + Q -> DQ -> BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> Q -> DQ -> BMM2 -> Q -> DQ. + If has_fp8_qdq == False, we match this FP8 fMHA v2 pattern: + BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> BMM2. + + Args: + graph: + The graph to match FP8 MHA pattern. + softmax_op: + The softmax op of FP8 MHA we want to match. + nodes_to_exclude: + List of Nodes to exclude from quantization. + has_fp8_qdq: + If True, match the FP8 MHA with Q/DQs. + Else, match the FP8 MHA without Q/DQs. + + Returns: + List of BMM1 node, Softmax node and BMM2 node. + """ + if has_fp8_qdq: + softmax_bmm1_chain_types = [ + ["Softmax", "MatMul", "DequantizeLinear", "QuantizeLinear"], + ["Softmax", "Add", "MatMul", "DequantizeLinear", "QuantizeLinear"], + ["Softmax", "Div", "MatMul", "DequantizeLinear", "QuantizeLinear"], + ["Softmax", "Mul", "MatMul", "DequantizeLinear", "QuantizeLinear"], + ["Softmax", "Add", "Div", "MatMul", "DequantizeLinear", "QuantizeLinear"], + ["Softmax", "Add", "Mul", "MatMul", "DequantizeLinear", "QuantizeLinear"], + ] + softmax_bmm2_chain_type = [ + "Softmax", + "QuantizeLinear", + "DequantizeLinear", + "MatMul", + "QuantizeLinear", + "DequantizeLinear", + ] + else: + softmax_bmm1_chain_types = [ + ["Softmax", "MatMul"], + ["Softmax", "Add", "MatMul"], + ["Softmax", "Div", "MatMul"], + ["Softmax", "Mul", "MatMul"], + ["Softmax", "Add", "Div", "MatMul"], + ["Softmax", "Add", "Mul", "MatMul"], + ] + softmax_bmm2_chain_type = [ + "Softmax", + "MatMul", + ] + wild_card_types = [ + "Reshape", + "Transpose", + "Flatten", + "Cast", + ] + + bmm1_index = 1 + bmm2_index = 7 if has_fp8_qdq else 3 + # Maps chain_idx to bmm position offsets for indexing the fp8_mha_partition + # chain_idx=0 -> offset=0 + # chain_idx=1,2,3 -> offset=1 + # chain_idx=4,5 -> offset=2 + bmm_offset_map = {0: 0, 1: 1, 2: 1, 3: 1, 4: 2, 5: 2} + for chain_idx, softmax_bmm1_chain_type in enumerate(softmax_bmm1_chain_types): + fp8_mha_partition = [] + if has_path_type( + softmax_op, graph, softmax_bmm1_chain_type, False, wild_card_types, fp8_mha_partition + ) and has_path_type( + softmax_op, graph, softmax_bmm2_chain_type, True, wild_card_types, fp8_mha_partition + ): + offset = bmm_offset_map.get(chain_idx) + bmm1_node = fp8_mha_partition[bmm1_index + offset] + bmm2_node = fp8_mha_partition[bmm2_index + offset] + assert bmm1_node.op == "MatMul" and bmm2_node.op == "MatMul" + return [bmm1_node, softmax_op, bmm2_node] + return [] + + +def insert_matmul_casts(graph, matmul_node): + """Insert three cast nodes for MatMul's two inputs and output.""" + matmul_input0 = matmul_node.inputs[0] + matmul_input0_cast_output = gs.Variable( + name=f"{matmul_input0.name}/Cast_output", dtype=np.float32 + ) + graph.layer( + op="Cast", + name=f"{matmul_input0.name}/Cast", + inputs=[matmul_input0], + outputs=[matmul_input0_cast_output], + attrs={"to": np.float32}, + ) + matmul_node.inputs[0] = matmul_input0_cast_output + + matmul_input1 = matmul_node.inputs[1] + matmul_input1_cast_output = gs.Variable( + name=f"{matmul_input1.name}/Cast_output", dtype=np.float32 + ) + graph.layer( + op="Cast", + name=f"{matmul_input1.name}/Cast", + inputs=[matmul_input1], + outputs=[matmul_input1_cast_output], + attrs={"to": np.float32}, + ) + matmul_node.inputs[1] = matmul_input1_cast_output + + matmul_output = matmul_node.outputs[0] + matmul_output_cast_input = gs.Variable( + name=f"{matmul_output.name}/Cast_output", dtype=np.float32 + ) + graph.layer( + op="Cast", + name=f"{matmul_output.name}/Cast", + inputs=[matmul_output_cast_input], + outputs=[matmul_output], + attrs={"to": np.float16}, + ) + matmul_node.outputs[0] = matmul_output_cast_input + + +def insert_fp8_mha_casts(onnx_model): + r"""Insert three cast ops. + + The first cast will be added before the input0 of MatMul to cast fp16 to fp32. + The second cast will be added before the input1 of MatMul to cast fp16 to fp32. + The third cast will be added after the output of MatMul to cast fp32 back to fp16. + The insertion of Cast ops in the FP8 MHA part actually forbids the MHAs to run + with FP16 accumulation because the compiler only has FP32 accumulation kernels for FP8 MHAs. + """ + graph = gs.import_onnx(onnx_model) + graph.cleanup().toposort() + + # Match FP8 MHA: Q -> DQ -> BMM1 -> (Mul/Div) -> (Add) -> Softmax -> (Cast) -> Q -> DQ -> BMM2 -> Q -> DQ + for node in graph.nodes: + if node.op == "Softmax": + fp8_fmha_v2_pattern = match_fp8_mha_pattern(graph, node, True) + # Insert cast nodes on BMM2's input and output tensors. + if len(fp8_fmha_v2_pattern) == 3: + insert_matmul_casts(graph, fp8_fmha_v2_pattern[0]) + insert_matmul_casts(graph, fp8_fmha_v2_pattern[2]) + + graph.cleanup().toposort() + + return gs.export_onnx(graph) + + +def convert_fp16_io(graph): + """Convert graph I/O to FP16.""" + convertible_dtypes = [ + onnx.TensorProto.FLOAT, + onnx.TensorProto.DOUBLE, + onnx.TensorProto.BFLOAT16, + ] + for input_tensor in graph.inputs: + input_tensor.dtype = ( + onnx.TensorProto.FLOAT16 + if input_tensor.dtype in convertible_dtypes + else input_tensor.dtype + ) + for output_tensor in graph.outputs: + output_tensor.dtype = ( + onnx.TensorProto.FLOAT16 + if output_tensor.dtype in convertible_dtypes + else output_tensor.dtype + ) diff --git a/modelopt/onnx/quantization/gs_patching.py b/modelopt/onnx/quantization/gs_patching.py new file mode 100644 index 0000000000..cd0f905aa5 --- /dev/null +++ b/modelopt/onnx/quantization/gs_patching.py @@ -0,0 +1,112 @@ +# 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. + +"""Patches onnx_graphsurgeon to support explicitly setting a dtype.""" + +from typing import Sequence, Union + +import numpy as np +import onnx +import onnx_graphsurgeon as gs +from onnx.mapping import TENSOR_TYPE_MAP +from onnx_graphsurgeon.ir.tensor import LazyValues + +from modelopt.onnx.quantization.quant_utils import pack_float32_to_4bit_cpp_based + +GS_LOGGER = gs.logger.logger.G_LOGGER +ONNX_MAJOR, ONNX_MINOR = tuple(map(int, onnx.__version__.split(".")[:2])) + + +def _onnx_supports_int4(): + return ONNX_MAJOR > 1 or ONNX_MAJOR == 1 and ONNX_MINOR >= 16 + + +def _make_constant( + name: str, values: Union[np.ndarray, LazyValues], dtype: onnx.TensorProto.DataType +) -> gs.Constant: + """Creates a constant with a specified dtype.""" + converted_dtype = ( + dtype if isinstance(values, LazyValues) else TENSOR_TYPE_MAP[int(dtype)].np_dtype + ) + if values.dtype != converted_dtype: + GS_LOGGER.critical( + f"Trying to create tensor with incompatible types: `{values.dtype}`, `{dtype}`" + ) + + t = gs.Constant(name, values) + setattr(t, "explicit_dtype", dtype) + return t + + +def _make_variable( + name: str, dtype: onnx.TensorProto.DataType, shape: Sequence[Union[int, str]] +) -> gs.Constant: + """Creates a variable with a specified dtype.""" + x = gs.Variable(name, TENSOR_TYPE_MAP[int(dtype)].np_dtype, shape) + setattr(x, "explicit_dtype", dtype) + return x + + +def _export_tensor_proto(tensor: gs.Constant) -> onnx.TensorProto: + if isinstance(tensor._values, LazyValues): + onnx_tensor = tensor._values.tensor + else: + # is numpy array. + dtype = getattr( + tensor, "explicit_dtype", onnx.helper.np_dtype_to_tensor_dtype(tensor.values.dtype) + ) + + vals = tensor.values + if _onnx_supports_int4() and dtype in [onnx.TensorProto.INT4, onnx.TensorProto.UINT4]: + signed = dtype == onnx.TensorProto.INT4 + np_dtype = onnx.helper.tensor_dtype_to_np_dtype(dtype) + vals = pack_float32_to_4bit_cpp_based(tensor.values, signed=signed).astype(np_dtype) + + onnx_tensor = onnx.helper.make_tensor( + tensor.name, + dtype, + dims=tensor.values.shape, + vals=vals.tobytes(), + raw=True, + ) + if tensor.data_location is not None: + onnx_tensor.data_location = tensor.data_location + onnx_tensor.name = tensor.name + return onnx_tensor + + +def _export_value_info_proto(tensor: gs.Variable, do_type_check: bool) -> onnx.ValueInfoProto: + if do_type_check and tensor.dtype is None: + GS_LOGGER.critical( + "Graph input and output tensors must include dtype information. Please set the dtype" + " attribute for: {:}".format(tensor) + ) + + if tensor.dtype is not None: + dtype = getattr( + tensor, "explicit_dtype", onnx.helper.np_dtype_to_tensor_dtype(np.dtype(tensor.dtype)) + ) + onnx_tensor = onnx.helper.make_tensor_value_info(tensor.name, dtype, tensor.shape) + else: + onnx_tensor = onnx.helper.make_empty_tensor_value_info(tensor.name) + return onnx_tensor + + +def patch_gs_modules(): + """Dynamically patch graphsurgeon modules.""" + gs.make_constant = _make_constant + gs.make_variable = _make_variable + gs.exporters.onnx_exporter.OnnxExporter.export_tensor_proto = _export_tensor_proto + gs.exporters.onnx_exporter.OnnxExporter.export_value_info_proto = _export_value_info_proto diff --git a/modelopt/onnx/quantization/int4.py b/modelopt/onnx/quantization/int4.py new file mode 100644 index 0000000000..9f2e35745a --- /dev/null +++ b/modelopt/onnx/quantization/int4.py @@ -0,0 +1,1326 @@ +# 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. + +"""Performs INT4 WoQ on an ONNX model, and returns the ONNX ModelProto.""" + +import copy +import gc +import logging +import math +import os +import tempfile +import time +from typing import Any, Optional, Union, cast + +import numpy +import onnx +import onnx.numpy_helper as numpy_helper +import onnx_graphsurgeon as gs +from onnxruntime.quantization.calibrate import CalibrationDataReader +from tqdm import tqdm + +import modelopt.onnx.quantization.qdq_utils as qdq +from modelopt.onnx.op_types import is_fusible_scaling_op +from modelopt.onnx.quantization.calib_utils import RandomDataProvider +from modelopt.onnx.quantization.graph_utils import ( + expand_node_names_from_patterns, + get_tensor_consumer_nodes, + get_tensor_producer_nodes, +) +from modelopt.onnx.quantization.gs_patching import patch_gs_modules +from modelopt.onnx.quantization.ort_utils import create_inference_session +from modelopt.onnx.utils import save_onnx + +__all__ = ["quantize"] + +has_cupy = False +cupy_warning_msg = "" +try: + import cupy as np + + if not np.cuda.is_available(): + raise ImportError("Unable to use cupy as CUDA is not available.") + + has_cupy = True +except Exception as e: + # ImportError if cupy is not installed + # cupy_backends.cuda.api.runtime.CUDARuntimeError if CUDA is not available + import numpy as np + + cupy_warning_msg = f"Using slower INT4 ONNX quantization using numpy: {e}" + + +# Set logging level to info +logging.getLogger().setLevel(logging.INFO) +NUM_BITS = 4 +INT4_SCALE = 7.0 +INT4_MIN = -(2 ** (NUM_BITS - 1)) # -8 +INT4_MAX = 2 ** (NUM_BITS - 1) - 1 # 7 +UINT4_MIN = 0 +UINT4_MAX = 15 + +# following min-value for clip is taken from AutoAWQ where zero-point based quantization is +# supported and working +CLIP_MIN = 1e-5 + + +def _next_block_size_multiple(x: float, block_size: int) -> float: + return math.ceil(x / block_size) * block_size + + +def _pad(w: np.ndarray, block_size: int) -> np.ndarray: + """Pads `w` to next largest multiple of block_size, on axis 0.""" + if w.shape[0] % block_size == 0: + return w + + pad_width = _next_block_size_multiple(w.shape[0], block_size) - w.shape[0] + pads = [(0, 0) for _ in range(len(w.shape))] + pads[0] = (0, pad_width) + return np.pad(w, pads, mode="constant", constant_values=0) + + +def _depad(w: np.ndarray, orig_shape: tuple) -> np.ndarray: + """Depad axis 0 to original shape.""" + if w.shape == orig_shape: + return w + return w[0 : orig_shape[0], ...] + + +def find_scales(w: np.ndarray, block_size: int, alpha: float = 1.0, use_zero_point: bool = False): + """Find scale factors for `w` via `s = max(w.block(block_size)) / 7`.""" + w = _pad(w, block_size) + w = w.T + s_last_dim = w.shape[-1] // block_size + s_shape = list(w.shape) + s_shape[-1] = s_last_dim + if not use_zero_point: + w_amax = np.abs(w.reshape(-1, block_size)).max(axis=-1) + s = (w_amax * alpha) / INT4_SCALE + s = s.reshape(s_shape).T + z = None + else: + max_val = w.reshape(-1, block_size).max(axis=-1) + min_val = w.reshape(-1, block_size).min(axis=-1) + max_int = UINT4_MAX + min_int = UINT4_MIN + s = (max_val - min_val).clip(min=CLIP_MIN) / max_int + # z = -np.round(temp).clip(min=min_int, max=max_int) # gives 0 - need to check + temp = min_val / s + temp = np.round(temp) + temp = -temp + temp = temp.clip(min=min_int, max=max_int) + z = temp + assert s.shape == z.shape, "s and z shape mismatch" + s = s.reshape(s_shape).T + z = z.reshape(s_shape).T + return s, z + + +def rtn(w: np.ndarray, s: np.ndarray, block_size: int, zp: np.ndarray = None) -> np.ndarray: + """Quantizes `w` with scale factors `s` via Round-to-Nearest. + + Ties are broken by rounding to the nearest even number. + """ + w_padded = _pad(w, block_size) + num_blocks = w_padded.shape[0] // s.shape[0] + if zp is None: + w_padded = ( + np.rint(w_padded / s.repeat(num_blocks, axis=0)) + .clip(INT4_MIN, INT4_MAX) + .astype(np.int8) + ) + else: + w_padded = ( + (np.rint(w_padded / s.repeat(num_blocks, axis=0)) + zp.repeat(num_blocks, axis=0)) + .clip(UINT4_MIN, UINT4_MAX) + .astype(np.int8) + ) + return _depad(w_padded, w.shape) + + +def dq_tensor(w: np.ndarray, s: np.ndarray, block_size: int, zp: np.ndarray = None) -> np.ndarray: + """Dequantizes `w` with scale factors `s`.""" + w_padded = _pad(w, block_size) + num_blocks = w_padded.shape[0] // s.shape[0] + if zp is None: + w_padded = w_padded * s.repeat(num_blocks, axis=0) + else: + w_padded = (w_padded - zp.repeat(num_blocks, axis=0)) * s.repeat(num_blocks, axis=0) + return _depad(w_padded, w.shape) + + +def quantize_rtn( + onnx_model: onnx.onnx_pb.ModelProto, + gemm_io_type: onnx.TensorProto.DataType, + block_size: int, + dq_only: bool = False, +) -> onnx.onnx_pb.ModelProto: + """Quantizes `onnx_model` using the RTN (Round-to-Nearest) algorithm. + + This algorithm computes scale factors by computing s = max(abs(block)) / 8, for each block. The + quantized weights are computed via Q(w) = round_to_even(w / s), where `round_to_even` denotes + rounding ties to the nearest even integer (i.e. 1.5, 2.5 both round to 2). + + Always selects the first dimension (0) to block over. This is because we must batch over the Cin + dimension, and in ONNX, weights are always plugged into the RHS (i.e. y = x @ W). + """ + graph = gs.import_onnx(onnx_model) + gemm_nodes = [node for node in graph.nodes if node.op in ["Gemm", "MatMul"]] + gemm_tensors = {} + act_tensors = [] + for gemm in gemm_nodes: + for in_tensor in gemm.inputs: + if not isinstance(in_tensor, gs.Constant): + continue + if len(in_tensor.values.shape) == 1: + # 1D blocked quantization not supported. + continue + gemm_tensors[in_tensor.name] = in_tensor + act_tensors.append(gemm.inputs[0]) + + gemm_weights = {name: tensor.values for name, tensor in gemm_tensors.items()} + + scales = {} + for name, w in gemm_weights.items(): + s, zp = find_scales(np.asarray(w), block_size) + assert zp is None, "zero-point is not enabled but zp is found non-None" + scales[name] = s + + logging.info("Computed scales.") + + # Change the scale type to the expected type, fp16 by default + for name, _ in scales.items(): + s = scales[name] + scales[name] = s.astype(onnx.mapping.TENSOR_TYPE_MAP[gemm_io_type].np_dtype) + + # Change the input activation type to the expected type, fp16 by default + for act_tensor in act_tensors: + _change_input_type(onnx_model.graph, act_tensor.name, gemm_io_type) + + # Import the update graph + graph = gs.import_onnx(onnx_model) + + if dq_only: + # Calculate actual quantized weights. + gemm_weights_quantized = {} + for name, w in gemm_weights.items(): + qw = rtn(np.asarray(w), scales[name], block_size) + if has_cupy: + qw = np.asnumpy(qw) + scales[name] = np.asnumpy(scales[name]) + gemm_weights_quantized[name] = numpy.asarray(qw) + + qdq.insert_dq_nodes(graph, scales, quantized_weights=gemm_weights_quantized) + else: + if has_cupy: + for name, _ in scales.items(): + scales[name] = np.asnumpy(scales[name]) + qdq.insert_qdq_nodes(graph, scales, weight_map=gemm_tensors) + + logging.info(f"Inserted {'DQ' if dq_only else 'Q/DQ'} nodes.") + return gs.export_onnx(graph) + + +def quant_tensor(w: np.ndarray, block_size: int, alpha: float = 1.0, use_zero_point: bool = False): + """Quantize a tensor using alpha etc. and return the quantized tensor.""" + scale, zp = find_scales(w, block_size, alpha, use_zero_point) + wq = rtn(w, scale, block_size, zp) + return wq, scale, zp + + +class AWQClipHelper: + """AWQ calibration helper class.""" + + min_alpha = 0.5 + alpha_step = 0.05 + + def __init__(self, w, block_size: int, **kwargs): + """Initializes AWQClipHelper with a module weight.""" + self.alpha_step = kwargs.get("awqclip_alpha_step", self.alpha_step) + self.min_alpha = kwargs.get("awqclip_alpha_min", self.min_alpha) + self.alphas = [ + round(float(k), 2) for k in np.arange(self.min_alpha, 1.0, self.alpha_step) + ] + [1.0] + + ci, co = w.shape + self.block_size = block_size if block_size != -1 else w.shape[0] + w = _pad(w, block_size).T + self.w_amax = np.abs(w.reshape(-1, block_size)).max(axis=-1) + + self.loss = { + k: np.zeros((co, math.ceil(ci / self.block_size)), dtype=np.float32) + for k in self.alphas + } + self.best_loss = np.full_like(self.w_amax, float("inf")) + self.best_alpha = np.ones_like(self.w_amax) + + def update_best_params(self): + """Updates the loss dictionary.""" + for alpha, loss in self.loss.items(): + loss = loss.reshape(self.w_amax.shape) + indices = loss < self.best_loss + self.best_loss = np.where(indices, loss, self.best_loss) + self.best_alpha = np.where(indices, alpha, self.best_alpha) + + +def _clip_search( + x: np.ndarray, w: np.ndarray, awq_clip: AWQClipHelper, max_tokens: int = 64, **kwargs +): + """Apply AWQ algorithm on a weight and return optimum alpha. + + This algorithm defines a simple search space for the optimal scales: S = Sx ^ α. + S is only related to the magnitude of activation Sx, and a single hyper-parameter α is used to balance + between the protection of salient and non-salient channels. The algorithm finds the best α by a fast grid search + over the interval of [0, 1] (0 means do not scale; 1 corresponds to the most aggressive scaling). + Further weight clipping is also applied by minimizing the MSE error. + """ + # Select max_tokens from input + co_bsz = kwargs.get("awqclip_bsz_col", 1024) + x = np.reshape(x, (-1, x.shape[-1])) # _, ci + x = x[0 :: max(1, x.shape[0] // max_tokens)] # max_tokens, ci + + ci, co = w.shape + block_size = awq_clip.block_size + + # Pad weight and input if necessary + if ci % block_size != 0: + w = _pad(w, block_size) + x = _pad(x.T, block_size).T + + # Make a copy of the original padded weight to quantize with generated scales + w_copy = copy.deepcopy(w) + + # Reshape weight and input for batch processing over co dimension + w = w.T # co, ci + w = w.reshape(co, 1, -1, block_size) # co, 1, n_block, block_size + x = x.reshape(1, x.shape[0], -1, block_size) # 1, max_tokens, n_block, block_size + + # Loop over co dimension of the weight and generate scales + for co_batch in range(math.ceil(co / co_bsz)): + slice_s, slice_e = co_batch * co_bsz, min((co_batch + 1) * co_bsz, co) + weight = w[slice_s:slice_e] + org_out = np.sum(x * weight, axis=-1) # co_bsz, max_tokens, n_block + + # Compute loss for each alpha value + for alpha in awq_clip.loss.keys(): + # Perform QDQ on the whole original weight tensor + qw, scales, _ = quant_tensor(w_copy, block_size, alpha) + cur_w = dq_tensor(qw, scales, block_size) + + # Reshape before getting the batch of size co_bsz to multiply with input + cur_w = cur_w.T # ci, co -> co, ci + cur_w = cur_w.reshape(co, 1, -1, block_size) # co, 1, n_block, block_size + cur_w = cur_w[slice_s:slice_e] + + # Compute loss for each batch + cur_out = np.sum(x * cur_w, axis=-1) # co_bsz, max_tokens, n_block + loss = np.mean(np.power((org_out - cur_out), 2), axis=1) # co_bsz, n_block + awq_clip.loss[alpha][slice_s:slice_e] += loss + + if has_cupy: + np.get_default_memory_pool().free_all_blocks() + + # Update the best alpha value for the weight blocks + awq_clip.update_best_params() + if has_cupy: + np.get_default_memory_pool().free_all_blocks() + + +def _find_quantizable_weights( + graph: onnx.onnx_pb.GraphProto, + nodes_to_exclude: list[str], +) -> list[tuple[onnx.onnx_pb.ValueInfoProto, onnx.onnx_pb.ValueInfoProto, bool, int]]: + """Finds the quantizable weights from the graph.""" + wa_pack = [] + gemm_nodes = [node for node in graph.node if node.op_type in ["Gemm", "MatMul"]] + initializer_idxs = {initializer.name: idx for idx, initializer in enumerate(graph.initializer)} + for gemm in gemm_nodes: + exclude_this_node = False + + for i in range(len(nodes_to_exclude)): + if nodes_to_exclude[i] in gemm.name: + exclude_this_node = True + break + + if exclude_this_node: + continue + + if gemm.input[0] in initializer_idxs: + # Ex. two const input to MatMul_115 in fastvit0.onnx + # Note. RTN algorithm will quantize these weights though + continue + + if gemm.input[1] not in initializer_idxs: + continue + + weight_tensor = graph.initializer[initializer_idxs[gemm.input[1]]] + if len(weight_tensor.dims) == 1: # 1D blocked quantization not supported + continue + + gemm_io_type = cast(int, weight_tensor.data_type) + + act_tensor = onnx.helper.ValueInfoProto() + act_tensor.name = gemm.input[0] + + # TODO: support transA by transposing activation tensors in _clip_search + do_transpose = gemm.op_type == "Gemm" and any( + [attr.name == "transB" and attr.i > 0 for attr in gemm.attribute] + ) + + wa_pack.append((act_tensor, weight_tensor, do_transpose, gemm_io_type)) + + return wa_pack + + +def _augment_graph( + graph: onnx.onnx_pb.GraphProto, + wa_pack: list[tuple[gs.Tensor, gs.Tensor, bool, int]], +): + """Extend graph outputs with MatMuls activation input.""" + augmented_outputs = set([tensor.name for tensor in graph.output]) + for act_tensor, _, _, _ in wa_pack: + if act_tensor.name not in augmented_outputs: + graph.output.append(act_tensor) + augmented_outputs.add(act_tensor.name) + + +def _change_input_type( + graph: onnx.onnx_pb.GraphProto, input_name: str, gemm_io_type: onnx.TensorProto.DataType +): + # Find the corresponding value info in the graph + done = False + for value_info in graph.value_info: + if value_info.name == input_name: + value_info.type.tensor_type.elem_type = gemm_io_type + done = True + break + + if not done: + # If input not in value_info, it must be a graph input + for input_info in graph.input: + if input_info.name == input_name: + input_info.type.tensor_type.elem_type = gemm_io_type + break + + +def _quantize_awq_clip( + onnx_model: onnx.onnx_pb.ModelProto, + data_reader: CalibrationDataReader, + use_external_data_format: bool, + calibration_eps: list[str], + block_size: int, + force_fp16: bool = False, + nodes_to_exclude: list[str] = [], + **kwargs: Any, +) -> onnx.onnx_pb.ModelProto: + """Quantizes `onnx_model` using the Activation aware quantization a.k.a AWQ algorithm.""" + logging.info("Finding quantizable weights and augmenting graph output with input activations") + t = time.time() + augmented_model = copy.deepcopy(onnx_model) + graph = augmented_model.graph + + nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude) + # Collect quantizable weight tensors + wa_pack = _find_quantizable_weights(graph, nodes_to_exclude) + + # Add input activations to graph output + _augment_graph(augmented_model.graph, wa_pack) + logging.info(f"Augmenting took {time.time() - t} seconds") + + scales = {} + gemm_weights_quantized = {} + + t = time.time() + + # Create a temp file for augmented model + augmented_onnx_file, augmented_onnx_path = tempfile.mkstemp(suffix=".onnx") + os.close(augmented_onnx_file) + + # TODO: ONNX version issue, onnx_export uses current ONNX IR version. + augmented_model.ir_version = 9 + save_onnx(augmented_model, augmented_onnx_path, use_external_data_format) + logging.info(f"Saving the model took {time.time() - t} seconds") + + # Creating inference session and preparing inputs for calibration + session = create_inference_session(augmented_onnx_path, calibration_eps) + inputs = [] + for inp_d in data_reader: + inputs.append(inp_d) + assert isinstance(inp_d, dict) + + # Apply AWQ clip on selected weights + t = time.time() + alphas = {} + for i in tqdm(range(len(wa_pack)), desc="Running clip search..."): + act_tensor, weight_tensor, do_transpose, gemm_io_type = wa_pack[i] + + # First capture all the activation values after calibration data sweep + output_dicts = {} + for inp_d in inputs: + np_inp_d = {name: numpy.asarray(tensor) for name, tensor in inp_d.items()} + output = session.run([act_tensor.name], np_inp_d) + out = np.asarray(output[0]) + output_dicts.setdefault(act_tensor.name, []).append(out) + + # Concatenating the activation tensors over all calib data + x = np.concatenate(output_dicts[act_tensor.name], axis=0) # n_token, ci + w = numpy_helper.to_array( + weight_tensor, base_dir=os.path.dirname(augmented_onnx_path) + ).copy() + if do_transpose: + w = w.T + w = np.asarray(w) + + awq_clip = AWQClipHelper(w, block_size, **kwargs) + _clip_search(x, w, awq_clip, **kwargs) + alphas[weight_tensor.name] = awq_clip.best_alpha + + logging.info(f"Clip search for all weights took {time.time() - t} seconds") + + del session + + # Compute quantized weights and scales which are needed for DQ nodes + t = time.time() + for i in tqdm(range(len(wa_pack)), desc="Quantizing the weights..."): + act_tensor, weight_tensor, do_transpose, gemm_io_type = wa_pack[i] + gemm_io_type = cast(onnx.TensorProto.DataType, gemm_io_type) + + if force_fp16: + gemm_io_type = onnx.TensorProto.FLOAT16 + + w = numpy_helper.to_array( + weight_tensor, base_dir=os.path.dirname(augmented_onnx_path) + ).copy() + if do_transpose: + w = w.T + w = np.asarray(w) + + alpha = alphas.get(weight_tensor.name, 1) + qw, scale, _ = quant_tensor(w, block_size, alpha) + if has_cupy: + qw = np.asnumpy(qw) + scale = np.asnumpy(scale) + if do_transpose: + qw = qw.T + scale = scale.T + scales[weight_tensor.name] = scale.astype( + onnx.mapping.TENSOR_TYPE_MAP[gemm_io_type].np_dtype + ) + gemm_weights_quantized[weight_tensor.name] = numpy.asarray(qw).astype(numpy.int8) + + # Change the input activation type to the expected type, fp16 by default + # TODO: cast input C for Gemm + _change_input_type(onnx_model.graph, act_tensor.name, gemm_io_type) + + logging.info(f"Quantizing actual weights took {time.time() - t} seconds") + + logging.info("Inserting DQ nodes using quantized weights and scales ...") + t = time.time() + graph_gs = gs.import_onnx(onnx_model) + dq_node_attributes = {"axis": 0, "block_size": block_size} + qdq.insert_dq_nodes( + graph_gs, scales, quantized_weights=gemm_weights_quantized, attributes=dq_node_attributes + ) + logging.info(f"Inserting DQ nodes took {time.time() - t} seconds") + + logging.info("Exporting the quantized graph ...") + t = time.time() + model = gs.export_onnx(graph_gs) + model.ir_version = 9 + logging.info(f"Exporting took {time.time() - t} seconds") + + try: + os.remove(augmented_onnx_path) + if use_external_data_format: + os.remove(augmented_onnx_path + "_data") + except OSError: + logging.warn("Augmented ONNX model or external data file was not found!") + + return model + + +class AWQLiteHelper: + """AWQ Lite calibration helper class.""" + + alpha_step = 0.1 + + def __init__(self, x, w, block_size: int, **kwargs): + """Initializes AWQLiteHelper with a module weight.""" + self.alpha_step = kwargs.get("awqlite_alpha_step", self.alpha_step) + self.block_size = block_size if block_size != -1 else w.shape[0] + # w = _pad(w, block_size).T + # self.w_amax = np.abs(w.reshape(-1, block_size)).max(axis=-1) + self.weight_scale = None if w is None else get_weight_scale(w, self.block_size) + self.act_scale = None if x is None else get_act_scale(x) + self.loss = { + k.item(): 0.0 + for k in np.arange(0, 1.0 + AWQLiteHelper.alpha_step, AWQLiteHelper.alpha_step) + } + self.best_scale = None + self.best_alpha = None + + def update_best_params(self): + """Updates best-alpha and best-scale.""" + self.best_alpha = min(self.loss, key=self.loss.__getitem__) + self.best_scale = get_scale( + self.act_scale, + self.weight_scale, + self.best_alpha, + ) + + +def get_act_scale(x): + """Get scale tensors for inputs.""" + return x.__abs__().mean(0) + + +def get_weight_scale(weight, block_size=None): + """Get scale tensors for weights.""" + org_shape = weight.shape + slice_after_padding = None + if block_size: + if org_shape[0] % block_size != 0: + slice_after_padding = slice(org_shape[0]) + weight = _pad(weight, block_size) + org_shape = weight.shape + weight = weight.reshape(block_size, -1) + weight_abs_amax = weight.__abs__().max(axis=0, keepdims=True) + scale = weight.__abs__() / (weight_abs_amax + np.finfo(weight_abs_amax.dtype).tiny) + scale = scale.reshape(org_shape) + if slice_after_padding is not None: + scale = scale[slice_after_padding, ...] + scale = scale.mean(1) + return scale + + +def get_scale(x_max, w_max, alpha): + """Get AWQ lite scales as described by 's' in the paper.""" + scales = (x_max.__pow__(alpha) / (w_max.__pow__(1 - alpha) + np.finfo(w_max.dtype).tiny)).clip( + min=1e-4, max=1e4 + ) + scales = scales / np.sqrt(scales.max() * scales.min()) + return scales + + +def run_awq_scale_search_per_node( + wa_pack: list[tuple[gs.Tensor, gs.Tensor, bool, int]], + augmented_onnx_path, + block_size, + use_zero_point, + session, + awq_lite, + inputs, + tqdm_msg_append_str, + enable_weight_clipping, + enable_fast_path_using_high_sysram, + output_data, + clip_alphas, + **kwargs: Any, +): + """Method that iterates over each quantizable node for scale search.""" + assert len(awq_lite) == len(wa_pack) + + for i in tqdm( + range(len(wa_pack)), + desc="Running AWQ scale search per node" + tqdm_msg_append_str, + ): + act_tensor, weight_tensor, do_transpose, gemm_io_type = wa_pack[i] + + output_dicts = {} + + if enable_fast_path_using_high_sysram: + assert ( + len(output_data) > 0 + ), "fast-path is enabled but node-inputs are not pre-determined before grid search" + node_inputs = [] + for j in range(len(output_data)): + node_inputs.append(np.asarray(output_data[j][i])) + # want to free system RAM asap since that data is here copied to GPU for this node + output_data[j][i] = None + output_dicts[act_tensor.name] = node_inputs + else: + # First capture all the activation values after calibration data sweep + for inp_d in inputs: + np_inp_d = {name: numpy.asarray(tensor) for name, tensor in inp_d.items()} + output = session.run([act_tensor.name], np_inp_d) + out = np.asarray(output[0]) + output_dicts.setdefault(act_tensor.name, []).append(out) + + # Concatenating the activation tensors over all calib data + w = numpy_helper.to_array( + weight_tensor, base_dir=os.path.dirname(augmented_onnx_path) + ).copy() + if do_transpose: + w = w.T + w = np.asarray(w) + + x = np.concatenate(output_dicts[act_tensor.name], axis=0).reshape( + (-1, w.shape[0]) + ) # n_token, ci + + awq_lite[i] = AWQLiteHelper(x, w, block_size, **kwargs) + + out_actual = x.__matmul__(w) + + for alpha in awq_lite[i].loss.keys(): + awq_scale = get_scale( + awq_lite[i].act_scale, + awq_lite[i].weight_scale, + alpha, + ) + x_scaled = x * 1.0 / awq_scale + w_scaled = w * awq_scale[:, np.newaxis] + + qw, scale, zp = quant_tensor(w_scaled, block_size, use_zero_point=use_zero_point) + dqw = dq_tensor(qw, scale, block_size, zp) + out_curr = x_scaled.__matmul__(dqw) + loss = np.mean(np.power((out_actual - out_curr), 2)) + del out_curr + awq_lite[i].loss[alpha] = loss + + awq_lite[i].update_best_params() + if enable_weight_clipping: + w = w * (awq_lite[i].best_scale[:, np.newaxis]) + awq_clip = AWQClipHelper(w, block_size, **kwargs) + _clip_search(x, w, awq_clip, **kwargs) + clip_alphas[weight_tensor.name] = awq_clip.best_alpha + del x, w, out_actual, output_dicts + if has_cupy: + np.get_default_memory_pool().free_all_blocks() + + return awq_lite, clip_alphas + + +def get_act_to_weight_map_and_act_to_wa_pack_map( + wa_pack: list[tuple[gs.Tensor, gs.Tensor, bool, int]], +): + """Method to return subgraph related maps based on activation-name as key. + + This method returns 2 maps: + (a) map of act-name to input-node's weights dimensions + (b) map of act-name to wa_pack indices with same act-name + """ + act_to_wa_pack_map = {} + act_to_quant_nodes_weight_shape_map = {} + for i in tqdm(range(len(wa_pack)), desc="Getting activation names maps..."): + act_tensor, weight_tensor, do_transpose, gemm_io_type = wa_pack[i] + # wa_pack index is stored in map to represent quant nodes + act_to_wa_pack_map.setdefault(act_tensor.name, []).append(i) + act_to_quant_nodes_weight_shape_map.setdefault(act_tensor.name, []).append( + weight_tensor.dims + ) + if len(act_to_quant_nodes_weight_shape_map[act_tensor.name]) > 1: + assert ( + weight_tensor.dims[0] == act_to_quant_nodes_weight_shape_map[act_tensor.name][0][0] + ) + + return act_to_wa_pack_map, act_to_quant_nodes_weight_shape_map + + +def get_x_w_mean_for_subgraph( + wa_pack: list[tuple[gs.Tensor, gs.Tensor, bool, int]], + wa_pack_idx_list, + augmented_onnx_path, + x, + block_size, +): + """This method returns x-mean and w-mean.""" + x_sum = np.sum(np.abs(x), axis=0, dtype=np.float32) + x_m = x_sum / x.shape[0] + x_m = x_m.astype(x.dtype) + del x_sum + + w_concatenated = None + for wa_pack_idx in wa_pack_idx_list: + act_tensor, weight_tensor, do_transpose, gemm_io_type = wa_pack[wa_pack_idx] + w = numpy_helper.to_array( + weight_tensor, base_dir=os.path.dirname(augmented_onnx_path) + ).copy() + if do_transpose: + w = w.T + w = np.asarray(w) + if w_concatenated is None: + w_concatenated = w + else: + w_concatenated = np.concatenate([w_concatenated, w], axis=1) + del w + + if has_cupy: + np.get_default_memory_pool().free_all_blocks() + + assert w_concatenated is not None + + org_shape = w_concatenated.shape + w_concatenated = w_concatenated.reshape(block_size, -1) + div_by = np.amax(np.abs(w_concatenated), axis=0) + div_by = div_by + 1e-6 # damping factor 1e-6 is taken from AutoAWQ + w_concatenated = np.abs(w_concatenated) / div_by + w_concatenated = w_concatenated.reshape(org_shape) + w_m = np.mean(w_concatenated, axis=1) + + del w_concatenated + gc.collect() + if has_cupy: + np.get_default_memory_pool().free_all_blocks() + + x_m = x_m.reshape(-1) + w_m = w_m.reshape(-1) + assert x_m.shape == w_m.shape + + return x_m, w_m + + +def run_awq_scale_search_per_subgraph( + wa_pack: list[tuple[gs.Tensor, gs.Tensor, bool, int]], + act_to_wa_pack_map, + act_to_quant_nodes_weight_shape_map, + augmented_onnx_path, + block_size, + use_zero_point, + session, + awq_lite, + inputs, + tqdm_msg_append_str, + **kwargs: Any, +): + """Method that iterates over each quantizable subgraph/siblings for scale search.""" + for act, wa_pack_idx_list in tqdm( + act_to_wa_pack_map.items(), + desc="Running AWQ scale search per subgraph" + tqdm_msg_append_str, + ): + output_dicts = {} + + for inp_d in inputs: + np_inp_d = {name: numpy.asarray(tensor) for name, tensor in inp_d.items()} + output = session.run([act], np_inp_d) + out = np.asarray(output[0]) + output_dicts.setdefault(act, []).append(out) + assert len(act_to_quant_nodes_weight_shape_map[act]) > 0 + common_dim = act_to_quant_nodes_weight_shape_map[act][0][0] + x = np.concatenate(output_dicts[act], axis=0).reshape((-1, common_dim)) # n_token, ci + + del output_dicts + + x_m, w_m = get_x_w_mean_for_subgraph( + wa_pack, wa_pack_idx_list, augmented_onnx_path, x, block_size + ) + + best_error = float("inf") + best_alpha = None + best_scale = None + + out_actual = {} + + alpha_step = kwargs.get("awqlite_alpha_step", AWQLiteHelper.alpha_step) + alpha_values = np.arange(0, 1.0 + alpha_step, alpha_step) + + for alpha in alpha_values: + loss = 0.0 + # offset and clip value in the following formula is taken from AutoAWQ + awq_scale = (x_m.__pow__(alpha) / (w_m.__pow__(1 - alpha) + 1e-4)).clip(min=1e-4) + awq_scale = awq_scale / np.sqrt(awq_scale.max() * awq_scale.min()) + awq_scale[np.isinf(awq_scale)] = 1 + awq_scale[np.isnan(awq_scale)] = 1 + for wa_pack_idx in wa_pack_idx_list: + _, weight_tensor, do_transpose, _ = wa_pack[wa_pack_idx] + w = numpy_helper.to_array( + weight_tensor, base_dir=os.path.dirname(augmented_onnx_path) + ).copy() + if do_transpose: + w = w.T + w = np.asarray(w) + out_act = out_actual.get(wa_pack_idx, None) + if out_act is None: + out_act = x.__matmul__(w) + out_actual[wa_pack_idx] = out_act + assert out_act is not None + x_scaled = x * 1.0 / awq_scale + w_scaled = w * awq_scale[:, np.newaxis] + qw, scale, zp = quant_tensor(w_scaled, block_size, use_zero_point=use_zero_point) + dqw = dq_tensor(qw, scale, block_size, zp) + out_curr = x_scaled.__matmul__(dqw) + loss += np.mean(np.power((out_act - out_curr), 2)) + del out_curr, out_act + del w, x_scaled, w_scaled, qw, scale, dqw + if has_cupy: + np.get_default_memory_pool().free_all_blocks() + is_best = loss < best_error + if is_best: + best_error = loss + best_alpha = alpha + best_scale = awq_scale + + for wa_pack_idx in wa_pack_idx_list: + assert np.isnan(best_scale).sum() == 0, best_scale + assert awq_lite[wa_pack_idx] is None + awq_lite[wa_pack_idx] = AWQLiteHelper(None, None, block_size) + awq_lite[wa_pack_idx].best_alpha = best_alpha + awq_lite[wa_pack_idx].best_scale = best_scale + assert awq_lite[wa_pack_idx] is not None + + del x, out_actual, x_m, w_m + if has_cupy: + np.get_default_memory_pool().free_all_blocks() + + return awq_lite + + +def get_parent_child_nodes_map( + graph: onnx.onnx_pb.GraphProto, + wa_pack: list[tuple[gs.Tensor, gs.Tensor, bool, int]], +): + """Get mapping of parent nodes to their MatMul/Gemm nodes with quantizable weights.""" + parent_child_nodes_map = {} + output_name_to_node = get_tensor_producer_nodes(graph) + input_name_to_nodes = get_tensor_consumer_nodes(graph) + + for act_tensor, _, _, _ in wa_pack: + parent_name = output_name_to_node[act_tensor.name].name + parent_child_nodes_map[parent_name] = [] + for node in input_name_to_nodes[act_tensor.name]: + if node.op_type in ["Gemm", "MatMul"]: + parent_child_nodes_map[parent_name].append(node) + + return parent_child_nodes_map, input_name_to_nodes + + +def _quantize_awq_lite( + onnx_model: onnx.onnx_pb.ModelProto, + data_reader: CalibrationDataReader, + use_external_data_format: bool, + calibration_eps: list[str], + block_size: int, + force_fp16: bool = False, + enable_fast_path_using_high_sysram: bool = False, + enable_weight_clipping: bool = False, + use_zero_point: bool = False, + nodes_to_exclude: list[str] = [], + **kwargs: Any, +) -> onnx.onnx_pb.ModelProto: + """Quantizes `onnx_model` using the Activation aware quantization a.k.a AWQ algorithm.""" + logging.info("Finding quantizable weights and augmenting graph output with input activations") + t = time.time() + + run_per_subgraph = kwargs.get("awqlite_run_per_subgraph", False) + fuse_nodes = kwargs.get("awqlite_fuse_nodes", True) + + # TODO - add weight-clipping support in per-subgraph implementation + assert not run_per_subgraph or not enable_weight_clipping + + # TODO - evaluate/add sysram based fast-path support in per-subgraph implementation + assert not run_per_subgraph or not enable_fast_path_using_high_sysram + + augmented_model = copy.deepcopy(onnx_model) + graph = augmented_model.graph + + nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude) + # Collect quantizable weight tensors + wa_pack = _find_quantizable_weights(graph, nodes_to_exclude) + if fuse_nodes: + parent_child_nodes_map, input_name_to_nodes = get_parent_child_nodes_map( + onnx_model.graph, wa_pack + ) + + # Add input activations to graph output + _augment_graph(augmented_model.graph, wa_pack) + logging.info(f"Augmenting took {time.time() - t} seconds") + + zero_points = {} + scales = {} + gemm_weights_quantized = {} + input_tensors = {} + pre_quant_scale = {} + + t = time.time() + + # Create a temp file for augmented model + augmented_onnx_file, augmented_onnx_path = tempfile.mkstemp(suffix=".onnx") + os.close(augmented_onnx_file) + + # TODO: ONNX version issue, onnx_export uses current ONNX IR version. + augmented_model.ir_version = 9 + save_onnx(augmented_model, augmented_onnx_path, use_external_data_format) + logging.info(f"Saving the model took {time.time() - t} seconds") + + # Creating inference session and preparing inputs for calibration + session = create_inference_session(augmented_onnx_path, calibration_eps) + inputs = [] + for inp_d in data_reader: + inputs.append(inp_d) + assert isinstance(inp_d, dict) + + gc.collect() + + output_data = [] + + if enable_fast_path_using_high_sysram: + print("Fast-path-using-high-sysram is enabled.\n") + + tensor_names_list = [] + for i in tqdm(range(len(wa_pack)), desc="Getting tensor names..."): + act_tensor, weight_tensor, do_transpose, gemm_io_type = wa_pack[i] + tensor_names_list.append(act_tensor.name) + + for i in tqdm(range(len(inputs)), desc="Caching activations..."): + inp_d = inputs[i] + np_inp_d = {name: numpy.asarray(tensor) for name, tensor in inp_d.items()} + output = session.run(tensor_names_list, np_inp_d) + output_data.append(output) + + del session + session = None + gc.collect() + + # Apply AWQ lite on selected weights + t = time.time() + awq_lite = [None] * len(wa_pack) + clip_alphas = {} + + msg = "..." + if enable_weight_clipping: + msg = " and clip-range search..." + + act_to_wa_pack_map, act_to_quant_nodes_weight_shape_map = ( + get_act_to_weight_map_and_act_to_wa_pack_map(wa_pack) + ) + + if run_per_subgraph: + awq_lite = run_awq_scale_search_per_subgraph( + wa_pack, + act_to_wa_pack_map, + act_to_quant_nodes_weight_shape_map, + augmented_onnx_path, + block_size, + use_zero_point, + session, + awq_lite, + inputs, + msg, + **kwargs, + ) + else: + awq_lite, clip_alphas = run_awq_scale_search_per_node( + wa_pack, + augmented_onnx_path, + block_size, + use_zero_point, + session, + awq_lite, + inputs, + msg, + enable_weight_clipping, + enable_fast_path_using_high_sysram, + output_data, + clip_alphas, + **kwargs, + ) + assert len(awq_lite) == len(wa_pack) + for i in range(len(awq_lite)): + assert awq_lite[i] is not None + + if enable_weight_clipping: + assert len(clip_alphas.keys()) == len(wa_pack) + + logging.info("AWQ scale search" + msg.strip(".") + f" took {time.time() - t} seconds") + + if session is not None: + del session + session = None + if has_cupy: + np.get_default_memory_pool().free_all_blocks() + del output_data + gc.collect() + + # Compute quantized weights and scales which are needed for DQ nodes + t = time.time() + # Use a common mean scale for weights within a sub-graph + if fuse_nodes and not run_per_subgraph: + for _, wa_pack_idx_list in act_to_wa_pack_map.items(): + group_awq_scale = [] + for wa_pack_idx in wa_pack_idx_list: + group_awq_scale.append(awq_lite[wa_pack_idx].best_scale[:, np.newaxis]) + mean_awq_scale = np.concatenate(group_awq_scale, axis=1) + mean_awq_scale = mean_awq_scale.mean(axis=1) + for wa_pack_idx in wa_pack_idx_list: + awq_lite[wa_pack_idx].best_scale = mean_awq_scale + + for i in tqdm(range(len(wa_pack)), desc="Quantizing the weights..."): + act_tensor, weight_tensor, do_transpose, gemm_io_type = wa_pack[i] + gemm_io_type = cast(onnx.TensorProto.DataType, gemm_io_type) + + if force_fp16: + gemm_io_type = onnx.TensorProto.FLOAT16 + + w = numpy_helper.to_array( + weight_tensor, base_dir=os.path.dirname(augmented_onnx_path) + ).copy() + if do_transpose: + w = w.T + w = np.asarray(w) + + w_scaled = w * awq_lite[i].best_scale[:, np.newaxis] + alpha = clip_alphas.get(weight_tensor.name, 1) + assert enable_weight_clipping or ( + alpha == 1 + ), "clip range enabled without enabling weight-clipping param" + qw, scale, zp = quant_tensor(w_scaled, block_size, alpha, use_zero_point=use_zero_point) + assert use_zero_point is True or zp is None, "zp is not according to use-zero-point setting" + if do_transpose: + qw = qw.T + scale = scale.T + if zp is not None: + zp = zp.T + if has_cupy: + qw = np.asnumpy(qw) + scale = np.asnumpy(scale) + if zp is not None: + zp = np.asnumpy(zp) + scales[weight_tensor.name] = scale.astype( + onnx.mapping.TENSOR_TYPE_MAP[gemm_io_type].np_dtype + ) + weight_dtype = numpy.int8 + if zp is not None: + zero_points[weight_tensor.name] = numpy.asarray(zp).astype(numpy.uint8) + weight_dtype = numpy.uint8 + gemm_weights_quantized[weight_tensor.name] = numpy.asarray(qw).astype(weight_dtype) + input_tensors[weight_tensor.name] = act_tensor.name + pqs_value = ( + awq_lite[i] + .best_scale[:, np.newaxis] + .astype(onnx.mapping.TENSOR_TYPE_MAP[gemm_io_type].np_dtype) + ).T + if has_cupy: + pqs_value = np.asnumpy(pqs_value) + pre_quant_scale[weight_tensor.name] = pqs_value + + # Change the input activation type to the expected type, fp16 by default + # TODO: cast input C for Gemm + _change_input_type(onnx_model.graph, act_tensor.name, gemm_io_type) + + logging.info(f"Quantizing actual weights took {time.time() - t} seconds") + + # Fuse Mul nodes with parent node if possible + if fuse_nodes: + logging.info("Fusing pre-quant scale Mul nodes with parent node ...") + t = time.time() + updated_nodes = set() + name_to_node_map = {node.name: node for node in onnx_model.graph.node} + initializer_map = { + initializer.name: initializer for initializer in onnx_model.graph.initializer + } + for parent, child_nodes in parent_child_nodes_map.items(): + if parent == "root_0": + continue + parent = name_to_node_map[parent] + if parent.name in updated_nodes: + continue + weight_tensor_names = [] + # When fuse_nodes or run_per_subgraph is True, + # scales computed for each child_nodes will be same. + # Hence, picking pre_quant_scale corresponding to any child_nodes is acceptable + input_scale = np.asarray(pre_quant_scale[child_nodes[0].input[1]]) + for node in child_nodes: + weight_tensor_names.append(node.input[1]) + if ( + is_fusible_scaling_op(parent.op_type) + and not all([initializer_map.get(inp) is None for inp in parent.input]) + and len(input_name_to_nodes[child_nodes[0].input[0]]) == len(child_nodes) + ): + for inp in parent.input: + if initializer_map.get(inp) is not None: + tensor = initializer_map[inp] + tensor_array = numpy_helper.to_array( + tensor, + base_dir=os.path.dirname(augmented_onnx_path), + ) + new_tensor = np.asarray(tensor_array) / input_scale + new_tensor = numpy_helper.from_array(new_tensor.get(), tensor.name) + # replace initializer with new scaled array + tensor.CopyFrom(new_tensor) + for w_name in weight_tensor_names: + del pre_quant_scale[w_name] + updated_nodes.add(parent.name) + else: + scale_tensor = onnx.helper.make_tensor( + name=parent.output[0] + "_pre_quant_scale", + data_type=onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[input_scale.dtype], + dims=input_scale.shape, + vals=(1.0 / input_scale).flatten().tolist(), + ) + mul_op_name = parent.output[0] + "_pre_quant_scale_out" + mul_node = onnx.helper.make_node( + "Mul", + inputs=[child_nodes[0].input[0], scale_tensor.name], + outputs=[mul_op_name], + name=child_nodes[0].input[0] + "_pre_quant_scale_mul", + ) + for node in child_nodes: + node.input[0] = mul_node.output[0] + for w_name in weight_tensor_names: + del pre_quant_scale[w_name] + onnx_model.graph.initializer.append(scale_tensor) + onnx_model.graph.node.append(mul_node) + + logging.info(f"Fusing pre-quant scale Mul nodes took {time.time() - t} seconds") + + logging.info( + "Inserting DQ nodes and input_pre_quant_scale node using quantized weights and scales ..." + ) + t = time.time() + graph_gs = gs.import_onnx(onnx_model) + dq_node_attributes = {"axis": 0, "block_size": block_size} + qdq.insert_dq_nodes( + graph_gs, + scales, + quantized_weights=gemm_weights_quantized, + attributes=dq_node_attributes, + zero_points=zero_points if use_zero_point else None, + ) + if pre_quant_scale: + qdq.insert_pre_quant_scale_nodes(graph_gs, input_tensors, pre_quant_scale) + + logging.info(f"Inserting nodes took {time.time() - t} seconds") + + logging.info("Exporting the quantized graph ...") + t = time.time() + model = gs.export_onnx(graph_gs) + model.ir_version = 9 + logging.info(f"Exporting took {time.time() - t} seconds") + + try: + os.remove(augmented_onnx_path) + if use_external_data_format: + os.remove(augmented_onnx_path + "_data") + except OSError: + logging.warn("Augmented ONNX model or external data file was not found!") + + return model + + +def quantize( + onnx_path: Union[str, onnx.onnx_pb.ModelProto], + calibration_method: str = "awq_lite", + calibration_data_reader: CalibrationDataReader = None, + calibration_eps: list[str] = ["cuda:0", "dml:0", "cpu"], + use_external_data_format: bool = True, + use_zero_point: bool = False, + block_size: Optional[int] = None, + nodes_to_exclude: Optional[list[str]] = [r"/lm_head"], + **kwargs: Any, +) -> onnx.onnx_pb.ModelProto: + """Applies INT4 Weight-Only-Quantization (WoQ) to an ONNX model. + + Currently, only ``MatMul`` nodes quantization is supported. + + Args: + onnx_path: Input ONNX model (base model) + calibration_method: It determines the quantization algorithm. Few important algorithms are: + + - *awq_lite*: Applies AWQ scaling (Alpha search) followed by INT4 quantization. + - *awq_clip*: Executes weight clipping and INT4 quantization. + calibration_data_reader: It can be assigned a list of model inputs. If it is ``None``, then + a randomly generated model input will be used for calibration in AWQ implementation. + calibration_eps: It denotes ONNX Execution Providers (EPs) to use for base model calibration. + This list of EPs is then passed to create-session API of the onnxruntime (ORT) to + perform base model calibration. + + .. note:: + + Make sure that ORT package for chosen calibration-EPs is setup properly along + with their dependencies. + use_external_data_format: If True, save tensors to external file(s) for quantized model. + use_zero_point: If True, enables zero-point based quantization. + block_size: Block size parameter for int4 quantization. Default value of 128 is used for + ``block_size`` parameter. + nodes_to_exclude: List of node-names (or substrings of node-names) denoting the nodes to + exclude from quantization. + + .. note:: + + By default, ``lm-head`` node is NOT quantized. + kwargs: It denotes additional keyword arguments for int4 quantization. It includes: + + - **awqlite_alpha_step** (float): Step size to find best Alpha in awq-lite.Range: [0, 1]. + Default: 0.1. + - **awqclip_alpha_step** (float): Step size to find best Alpha in awq-clip. + Default: 0.05 + - **awqclip_alpha_min** (float): Minimum threshold for weight-clipping in awq-clip. + Default: 0.5. + - **awqclip_bsz_col** (int): Batch size for processing the column dimension in awq-clip. + Default: 1024. + **Returns**: A quantized ONNX model in ONNX ModelProto format. + """ + logging.info("Quantization Mode: int4") + + if cupy_warning_msg: + logging.warning(cupy_warning_msg) + + # Check if block_size is None and set default to 128 + if block_size is None: + block_size = 128 + + gemm_io_type: onnx.TensorProto.DataType = onnx.TensorProto.FLOAT + + # set config params + nodes_to_exclude = nodes_to_exclude or [] + + # Patch GS modules to support INT4. + patch_gs_modules() + + if isinstance(onnx_path, str): + onnx_model = onnx.load(onnx_path, load_external_data=use_external_data_format) + else: + onnx_model = onnx_path + + # Initialize calibration_data_reader if not provided + if calibration_data_reader is None: + calibration_data_reader = RandomDataProvider(onnx_model) + + if "trt" in calibration_method: + qdq.use_trt_qdq_ops() + + if calibration_method in ["rtn", "rtn_dq", "rtn_trt", "rtn_trt_dq"]: + onnx_model = quantize_rtn( + onnx_model, gemm_io_type, block_size, dq_only="dq" in calibration_method + ) + elif calibration_method in ["awq_lite", "awq_full"]: + do_weight_clipping = False + if calibration_method == "awq_full": + do_weight_clipping = True + onnx_model = _quantize_awq_lite( + onnx_model, + calibration_data_reader, + use_external_data_format, + calibration_eps, + block_size, + nodes_to_exclude=nodes_to_exclude, + use_zero_point=use_zero_point, + enable_weight_clipping=do_weight_clipping, + **kwargs, + ) + elif calibration_method in ["awq_clip", "awq_clip_trt"]: + onnx_model = _quantize_awq_clip( + onnx_model, + calibration_data_reader, + use_external_data_format, + calibration_eps, + block_size, + nodes_to_exclude=nodes_to_exclude, + **kwargs, + ) + else: + raise RuntimeError(f"Unsupported calibration method: '{calibration_method}'") + + return onnx_model diff --git a/modelopt/onnx/quantization/int8.py b/modelopt/onnx/quantization/int8.py new file mode 100755 index 0000000000..87d9062772 --- /dev/null +++ b/modelopt/onnx/quantization/int8.py @@ -0,0 +1,260 @@ +# 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. + +"""Performs INT8 quantization of an ONNX model, and returns the ONNX ModelProto.""" + +import logging +import os +import tempfile + +import onnx +import onnx_graphsurgeon as gs +from onnx_graphsurgeon.ir.graph import Graph +from onnx_graphsurgeon.ir.node import Node +from onnxmltools.utils.float16_converter import convert_float_to_float16 +from onnxruntime.quantization import CalibrationMethod +from onnxruntime.quantization.calibrate import CalibrationDataReader + +from modelopt.onnx.quantization.calib_utils import import_scales_from_calib_cache +from modelopt.onnx.quantization.graph_utils import ( + build_non_residual_input_map, + classify_partition_nodes, + expand_node_names_from_patterns, + filter_quantizable_kgen_heads, + find_nodes_from_matmul_to_exclude, + find_nodes_to_exclude, + remove_partial_input_qdq, +) +from modelopt.onnx.quantization.ort_patching import _quantize_static as quantize_static +from modelopt.onnx.quantization.ort_utils import configure_ort +from modelopt.onnx.quantization.partitioning import ( + find_fusible_partitions, + find_non_quantizable_partitions_from_patterns, + find_quantizable_nodes, + get_skiped_output_layers, +) +from modelopt.onnx.quantization.qdq_utils import replace_scale_values + +# Set logging level to info +logging.getLogger().setLevel(logging.INFO) + + +def _find_nodes_to_quantize( + graph: Graph, + quantizable_op_types: list[str], + verbose: bool, +) -> tuple[list[Node], list[tuple[Node, Node, str]]]: + # Build a map of add nodes to their non-residual inputs, i.e. fusible with Conv group + logging.info("Building non-residual Add input map ...") + non_residual_inputs, _ = build_non_residual_input_map(graph) + + logging.info( + "Searching for hard-coded patterns like MHA, LayerNorm, etc. to avoid quantization." + ) + non_quantizable_hard_coded_partitions = find_non_quantizable_partitions_from_patterns(graph) + + logging.info("Building KGEN/CASK targeted partitions ...") + # partitioned_nodes keeps track of nodes that are already part of some partition. + # Certain nodes of those partitions are quantizable. For example, heads. + partitioned_nodes = set(sum(non_quantizable_hard_coded_partitions, [])) + cask_fusible_partitions, kgen_partitions = find_fusible_partitions( + graph, + partitioned_nodes, + non_residual_inputs, + ) + if verbose: + logging.info( + "CASK fusible partitions:" + f" {[[node.name for node in partition] for partition in cask_fusible_partitions]}" + ) + logging.info( + "KGEN partitions:" + f" {[[node.name for node in partition] for partition in kgen_partitions]}" + ) + + logging.info("Classifying the partition nodes ...") + _, quantizable_partition_nodes, no_quantize_inputs = classify_partition_nodes( + cask_fusible_partitions, + ) + quantizable_kgen_heads, no_quantize_kgen_inputs = filter_quantizable_kgen_heads( + cask_fusible_partitions, + kgen_partitions, + quantizable_op_types, + ) + + quantizable_nodes = quantizable_kgen_heads + quantizable_partition_nodes + paritially_quantizable_nodes = [dst for _, dst, _ in no_quantize_inputs] + + # Quantize all inputs of partially quantizable nodes by ORT + # but remove QDQ from non-quantizable inputs in the post-processing step + quantizable_nodes.extend(paritially_quantizable_nodes) + + quantizable_nodes.extend( + find_quantizable_nodes(graph, quantizable_nodes, partitioned_nodes, quantizable_op_types) + ) + + skip_list = get_skiped_output_layers(graph, paritially_quantizable_nodes) + quantizable_nodes = [node for node in quantizable_nodes if node.name not in skip_list] + + return quantizable_nodes, no_quantize_inputs + no_quantize_kgen_inputs + + +def quantize( + onnx_path: str, + calibration_method: str = "entropy", + calibration_data_reader: CalibrationDataReader = None, + calibration_cache_path: str = None, + calibration_shapes: str = None, + calibration_eps: list[str] = ["cuda:0", "cpu", "trt"], + op_types_to_quantize: list[str] = None, + op_types_to_exclude: list[str] = None, + nodes_to_quantize: list[str] = None, + nodes_to_exclude: list[str] = None, + use_external_data_format: bool = True, + intermediate_generated_files: list[str] = [], + verbose: bool = False, + trt_extra_plugin_lib_paths: str = None, + high_precision_dtype: str = "fp32", + **kwargs, +) -> onnx.onnx_pb.ModelProto: + """Applies INT8 quantization to an ONNX file using the compiler friendly heuristics. + + Quantization of ['Add', 'AveragePool', 'BatchNormalization', 'Clip', 'Conv', 'ConvTranspose', + 'Gemm', 'GlobalAveragePool', 'MatMul', 'MaxPool', 'Mul'] op types are supported. + """ + logging.info("Quantization Mode: int8") + + # Take the onnx graph + onnx_model = onnx.load(onnx_path, load_external_data=use_external_data_format) + graph = gs.import_onnx(onnx_model) + graph.toposort() + + # Change the default configuration of ORT quantization + op_types_to_quantize = op_types_to_quantize or [] + op_types = set([node.op for node in graph.nodes]) + trt_guided_options, quantizable_op_types = configure_ort( + list(op_types), op_types_to_quantize, trt_extra_plugin_lib_paths, calibration_eps + ) + logging.info( + "Quantizable op types in the model:" + f" {[t for t in quantizable_op_types if t in op_types]}" + ) + + # Collect node names to include in quantization + no_quantize_inputs = [] + nodes_to_quantize = expand_node_names_from_patterns(graph, nodes_to_quantize) + if not nodes_to_quantize: + # If nodes_to_quantize is not passed, use user supplied op_types_to_quantize list + nodes_to_quantize = [node.name for node in graph.nodes if node.op in op_types_to_quantize] + + # If op_types_to_quantize is not provided, use default QDQ placement algorithm + if not nodes_to_quantize: + quantizable_nodes, no_quantize_inputs = _find_nodes_to_quantize( + graph, + quantizable_op_types, + verbose, + ) + nodes_to_quantize = [node.name for node in quantizable_nodes] + + # Read the calibration cache and quantize nodes for which activation scale values are cached + if calibration_cache_path: + act_scales_dict = import_scales_from_calib_cache(calibration_cache_path) + logging.info(f"Using calibration cache from {calibration_cache_path}") + iq_quantized_nodes = [] + quantized_tensors = [ + tensor_name.replace("_scale", "") for tensor_name in act_scales_dict.keys() + ] + for node in graph.nodes: + for node_input in node.inputs: + if node_input.name in quantized_tensors: + iq_quantized_nodes.append(node.name) + + logging.info( + f"Skipping quantization of nodes: {set(nodes_to_quantize) - set(iq_quantized_nodes)}" + ) + nodes_to_quantize = list(set(nodes_to_quantize).intersection(iq_quantized_nodes)) + + # Either of m or n in matmul is 1, this matmul cannot utilize TensorCores. + # The perf of adding Q/DQ layers is not good in TRT. Thus, in this case, + # do not add Q/DQ layers to this matmul. + matmul_nodes_to_exclude = find_nodes_from_matmul_to_exclude( + onnx_path, + use_external_data_format, + intermediate_generated_files, + calibration_data_reader, + calibration_eps, + verbose, + ) + nodes_to_exclude.extend(matmul_nodes_to_exclude) + + # Collect node names to exclude from quantization + nodes_to_exclude = find_nodes_to_exclude(graph, nodes_to_exclude, op_types_to_exclude) + + # Update the list of nodes to quantize + nodes_to_quantize = [ + node_name for node_name in nodes_to_quantize if node_name not in nodes_to_exclude + ] + + logging.info(f"Total number of nodes: {len(graph.nodes)}") + if not nodes_to_quantize: + logging.info( + "No node or node type is selected for quantization or model does not have them!" + ) + return + elif verbose: + logging.info(f"Selected nodes to quantize: {nodes_to_quantize}") + + # Create a temp file for intermediate model + tmp_onnx_file, tmp_onnx_path = tempfile.mkstemp(suffix=".onnx") + os.close(tmp_onnx_file) + + # Use ORT api to quantize the onnx model + quantize_static( + onnx_path, + tmp_onnx_path, + calibration_data_reader, + op_types_to_quantize=op_types_to_quantize, + nodes_to_quantize=nodes_to_quantize, + per_channel=True, + extra_options=trt_guided_options, + use_external_data_format=use_external_data_format, + calibrate_method=( + CalibrationMethod.Entropy + if calibration_method == "entropy" + # With ActivationSymmetric as True, MinMax calibration is equivalent to max calibration + else CalibrationMethod.MinMax + ), + ) + + intermediate_generated_files.append(tmp_onnx_path) + if use_external_data_format: + intermediate_generated_files.append(tmp_onnx_path + ".data") + + # Post-processing of the onnx model after ORT quantization + onnx_model = onnx.load(tmp_onnx_path) + graph = gs.import_onnx(onnx_model) + remove_partial_input_qdq(graph, no_quantize_inputs) + onnx_model = gs.export_onnx(graph) + if calibration_cache_path: + replace_scale_values(onnx_model.graph, act_scales_dict) + + if high_precision_dtype == "fp16": + # We need to convert float to float16 so as to speed up layers like LayerNorm or GroupNorm. + logging.info("Converting float tensors to float16") + onnx_model = convert_float_to_float16( + onnx_model, keep_io_types=True, disable_shape_infer=True + ) + + return onnx_model diff --git a/modelopt/onnx/quantization/operators.py b/modelopt/onnx/quantization/operators.py new file mode 100644 index 0000000000..783737f0ed --- /dev/null +++ b/modelopt/onnx/quantization/operators.py @@ -0,0 +1,96 @@ +# Adapted from https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/operators/conv.py +# and https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/operators/gather.py +# +# MIT License +# +# Copyright (c) Microsoft Corporation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# 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. + +"""Additional or modified QDQ operators on top of ORT quantized operators.""" + +from onnxruntime.quantization.operators.qdq_base_operator import QDQOperatorBase + +from ..op_types import is_normalization_op + + +class QDQNormalization(QDQOperatorBase): + """By default, ORT does not quantize Normalization ops. This module is intended to help with that. + + Note. QDQOperatorBase is not sufficient for dynamic input only quantization. + """ + + def __init__(self, onnx_quantizer, onnx_node): + """Normalization quantizer init.""" + super().__init__(onnx_quantizer, onnx_node) + + def quantize(self): + """Main function to quantize the Normalization ops.""" + node = self.node + assert is_normalization_op(node.op_type) + + # Quantize only the dynamic input (first input to the op) + self.quantizer.quantize_activation_tensor(node.input[0]) + if not self.disable_qdq_for_node_output: + self.quantizer.quantize_activation_tensor(node.output[0]) + + +class QDQConvTranspose(QDQOperatorBase): + """QDQ for ConvTranspose operator.""" + + def __init__(self, onnx_quantizer, onnx_node): + """ConvTranspose quantizer init.""" + super().__init__(onnx_quantizer, onnx_node) + + def quantize(self): + """Main function to quantize the ConvTranspose ops.""" + node = self.node + assert node.op_type == "ConvTranspose" + + self.quantizer.quantize_activation_tensor(node.input[0]) + if not self.disable_qdq_for_node_output: + self.quantizer.quantize_activation_tensor(node.output[0]) + + is_weight_per_channel, weight_axis = self.quantizer.is_tensor_per_channel( + node.input[1], default_axis=1 + ) + if is_weight_per_channel: + self.quantizer.quantize_weight_tensor_per_channel(node.input[1], weight_axis) + else: + self.quantizer.quantize_weight_tensor(node.input[1]) + + if len(node.input) == 3: + self.quantizer.quantize_bias_tensor( + node.name, node.input[2], node.input[0], node.input[1] + ) diff --git a/modelopt/onnx/quantization/ort_patching.py b/modelopt/onnx/quantization/ort_patching.py new file mode 100755 index 0000000000..213fd0c569 --- /dev/null +++ b/modelopt/onnx/quantization/ort_patching.py @@ -0,0 +1,610 @@ +# Adapted from https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/quant_utils.py +# and https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/calibrate.py +# and https://github.com/microsoft/onnxruntime/blob/2ac381c55397dffff327cc6efecf6f95a70f90a1/onnxruntime/python/tools/quantization/onnx_quantizer.py +# +# MIT License +# +# Copyright (c) Microsoft Corporation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# 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. + +"""This module contains all the patched functions from ORT.""" + +import logging +import tempfile +from pathlib import Path +from typing import Optional, Sequence, Union + +import numpy as np +import onnx +import onnxruntime as ort +from onnx import onnx_pb +from onnxruntime.quantization import calibrate +from onnxruntime.quantization.calibrate import ( + CalibraterBase, + CalibrationDataReader, + CalibrationMethod, + DistributionCalibrater, + EntropyCalibrater, + HistogramCalibrater, + HistogramCollector, + MinMaxCalibrater, + PercentileCalibrater, + TensorData, + TensorsData, +) +from onnxruntime.quantization.onnx_quantizer import ONNXQuantizer as QDQQuantizer +from onnxruntime.quantization.quant_utils import ( + QuantFormat, + QuantizationMode, + QuantType, + load_model_with_shape_infer, + model_has_pre_process_metadata, + save_and_reload_model_with_shape_infer, +) +from onnxruntime.quantization.quantize import check_static_quant_arguments +from onnxruntime.quantization.registry import QDQRegistry, QLinearOpsRegistry +from tqdm import tqdm + +if ort.__version__ >= "1.18": + from onnxruntime.quantization.qdq_quantizer import QDQQuantizer + + +def _collect_value(histogram_collector, name_to_arr): + """Collect histogram on real value.""" + for tensor, data_arr in tqdm(name_to_arr.items()): + data_arr = np.asarray(data_arr) # noqa: PLW2901 + data_arr = data_arr.flatten() # noqa: PLW2901 + + if data_arr.size > 0: + min_value = np.min(data_arr) + max_value = np.max(data_arr) + else: + min_value = np.array(0, dtype=data_arr.dtype) + max_value = np.array(0, dtype=data_arr.dtype) + + # Change the inf and nan values to meaningful min/max + min_value = ( + np.finfo(np.float32).tiny if np.isinf(min_value) or np.isnan(min_value) else min_value + ) + max_value = ( + np.finfo(np.float32).max if np.isinf(max_value) or np.isnan(max_value) else max_value + ) + + threshold = max(abs(min_value), abs(max_value)) + + if tensor in histogram_collector.histogram_dict: + old_histogram = histogram_collector.histogram_dict[tensor] + histogram_collector.histogram_dict[tensor] = histogram_collector.merge_histogram( + old_histogram, data_arr, min_value, max_value, threshold + ) + else: + hist, hist_edges = np.histogram( + data_arr, histogram_collector.num_bins, range=(-threshold, threshold) + ) + histogram_collector.histogram_dict[tensor] = ( + hist, + hist_edges, + min_value, + max_value, + threshold, + ) + + +def _check_opset_version(onnx_quantizer): + ai_onnx_domain = [ + opset + for opset in onnx_quantizer.model.model.opset_import + if not opset.domain or opset.domain in ["ai.onnx", "ai.onnx.contrib"] + ] + opset_version = ai_onnx_domain[0].version + + if opset_version == 10: + return 10 + + if opset_version < 10: + onnx_quantizer.model.model.opset_import.remove(ai_onnx_domain[0]) + onnx_quantizer.model.model.opset_import.extend([onnx.helper.make_opsetid("", 11)]) + opset_version = 11 + + if opset_version < 19 and onnx_quantizer.weight_qType == onnx_pb.TensorProto.FLOAT8E4M3FN: + onnx_quantizer.model.model.opset_import.remove(ai_onnx_domain[0]) + onnx_quantizer.model.model.opset_import.extend([onnx.helper.make_opsetid("", 19)]) + onnx_quantizer.model.model.ir_version = 9 + opset_version = 19 + + onnx_quantizer.fuse_dynamic_quant = True + return opset_version + + +def _select_tensors_to_calibrate(calibrator, model: onnx.ModelProto): + """Select input/output tensors of candidate nodes to calibrate. + + Returns: + tensors (set): set of tensor name. + value_infos (dict): tensor name to value info. + """ + value_infos = {vi.name: vi for vi in model.graph.value_info} + value_infos.update({ot.name: ot for ot in model.graph.output}) + value_infos.update({it.name: it for it in model.graph.input}) + initializer = {init.name for init in model.graph.initializer} + + tensors_to_calibrate = set() + tensor_type_to_calibrate = {onnx_pb.TensorProto.FLOAT, onnx_pb.TensorProto.FLOAT16} + + for node in model.graph.node: + # Hack: in calibrator.op_types_to_calibrate we pass nodes_to_quantize + if node.name in calibrator.op_types_to_calibrate: + for tensor_name in node.input: + if tensor_name in value_infos: + vi = value_infos[tensor_name] + if ( + vi.type.HasField("tensor_type") + and (vi.type.tensor_type.elem_type in tensor_type_to_calibrate) + and (tensor_name not in initializer) + ): + tensors_to_calibrate.add(tensor_name) + + return tensors_to_calibrate, value_infos + + +def _create_inference_session_with_ep_config(calibrator, **kwargs): + """Create an ORT InferenceSession.""" + sess_options = ort.SessionOptions() + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + providers = kwargs.get("execution_providers", []) + + # Note. This path can be an empty string, which denotes that the model has custom ops and TRT EP is needed. + trt_extra_plugin_lib_paths = kwargs.get("trt_extra_plugin_lib_paths", None) + + if trt_extra_plugin_lib_paths is not None: + # Assert supported configuration + assert ( + ort.__version__ >= "1.18" + ), "Plugin support is only available with ORT 1.18, which only supports TRT 10." + if "TensorrtExecutionProvider" not in ort.get_available_providers(): + raise RuntimeError( + "Could not find `TensorrtExecutionProvider`, only {}".format( + ort.get_available_providers() + ) + ) + trt_ep_options = ( + {"trt_extra_plugin_lib_paths": trt_extra_plugin_lib_paths} + if trt_extra_plugin_lib_paths + else {} + ) + # Set GPU memory usage limit + trt_ep_options["trt_max_workspace_size"] = 80 * 1024 * 1024 * 1024 + + # Update the TRT EP with the plugin path + if "TensorrtExecutionProvider" in providers: + providers.remove("TensorrtExecutionProvider") + providers.insert(0, ("TensorrtExecutionProvider", trt_ep_options)) + + # Create the inference session with EP configuration + calibrator.infer_session = ort.InferenceSession( + calibrator.augmented_model_path, + sess_options=sess_options, + providers=providers, + ) + + +def _collect_data_minmax_calibrator(calibrator, data_reader: CalibrationDataReader): + """This function overwrite is needed to solve OOM issue due to the unlimited accumulation of intermediate_outputs. + + Support for: MinMax Calibrator. + Modification: indented the last lines of code inside the while loop in order to run compute_data for each sample + batch individually instead of the entire data at once. The assumption here is that the ONNX file has bs=N + and the calibration data size is M (where M is a multiple of N). So the calibrator is a sequence of M/N + samples with bs=N. + """ + while True: + inputs = data_reader.get_next() + if not inputs: + break + calibrator.intermediate_outputs.append(calibrator.infer_session.run(None, inputs)) + + # ======== Modification: block is indentend in ======== + if len(calibrator.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + t = calibrator.compute_data() + if not isinstance(t, TensorsData): + raise TypeError(f"compute_data must return a TensorsData not {type(t)}.") + calibrator.clear_collected_data() + # ===================================================== + + +def _merge_range_minmax_calibrator(calibrator, old_range: TensorsData, new_range: TensorsData): + """This function is an auxilliary function of collect_data to solve the OOM issue in the MinMax Calibrator. + + Issue fixed with this function: old_range is not a dictionary, but old_range.data is. + TODO: create an MR in the ORT repository for this function. Alternatively, we can also file the MR fixing + TensorData (need to at least add items() function there). + """ + if not old_range: + return new_range + + for key, value in old_range.data.items(): + value_tuple = value.range_value + new_range_tuple = new_range.data[key].range_value + if calibrator.moving_average: + min_value = value_tuple[0] + calibrator.averaging_constant * ( + new_range_tuple[0] - value_tuple[0] + ) + max_value = value_tuple[1] + calibrator.averaging_constant * ( + new_range_tuple[1] - value_tuple[1] + ) + else: + min_value = min(value_tuple[0], new_range_tuple[0]) + max_value = max(value_tuple[1], new_range_tuple[1]) + new_range.data[key] = TensorData(lowest=min_value, highest=max_value) + + return new_range + + +def _collect_data_histogram_calibrator(calibrator, data_reader: CalibrationDataReader): + """This function overwrite is needed to solve OOM issue due to the unlimited accumulation of intermediate_outputs. + + Support for: Histogram Calibrator (which affects Entropy, Percentile, and DIstribution Calibrators). + Modification: indented the last lines of code inside the while loop in order to run compute_data for each sample + batch individually instead of the entire data at once. + """ + while True: + inputs = data_reader.get_next() + if not inputs: + break + calibrator.intermediate_outputs.append(calibrator.infer_session.run(None, inputs)) + + # ======== Modification: block is indentend in ======== + # Here, compute_date is calculated for every sample batch instead of the entire data at once. + if len(calibrator.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + output_names = [ + calibrator.infer_session.get_outputs()[i].name + for i in range(len(calibrator.intermediate_outputs[0])) + ] + output_dicts_list = [ + dict(zip(output_names, intermediate_output)) + for intermediate_output in calibrator.intermediate_outputs + ] + + merged_dict = {} + for d in output_dicts_list: + for k, v in d.items(): + merged_dict.setdefault(k, []).append(v) + + clean_merged_dict = { + i: merged_dict[i] for i in merged_dict if i in calibrator.tensors_to_calibrate + } + + if not calibrator.collector: + calibrator.collector = HistogramCollector( + method=calibrator.method, + symmetric=calibrator.symmetric, + num_bins=calibrator.num_bins, + num_quantized_bins=calibrator.num_quantized_bins, + percentile=calibrator.percentile, + scenario=calibrator.scenario, + ) + calibrator.collector.collect(clean_merged_dict) + + calibrator.clear_collected_data() + # ===================================================== + + +def _adjust_tensor_ranges(base_quantizer): + if base_quantizer.tensors_range is None: + return + + for node in base_quantizer.model.nodes(): + # adjust tensor_ranges for input of Clip and Relu node + if node.op_type in ["Clip", "Relu"]: + if base_quantizer.is_activation_symmetric: + continue + if not base_quantizer.should_quantize_node(node): + continue + if len(base_quantizer.model.input_name_to_nodes()[node.input[0]]) != 1: + continue + if ( + node.input[0] not in base_quantizer.tensors_range + or node.output[0] not in base_quantizer.tensors_range + ): + continue + td = base_quantizer.tensors_range[node.output[0]] + if not isinstance(td, TensorData): + raise TypeError(f"Unexpected type {type(td)} for {node.output[0]!r}.") + base_quantizer.tensors_range[node.input[0]] = td + # Adjust Softmax to range from 0.0 to 1.0 + elif node.op_type == "Softmax": + if node.output[0] not in base_quantizer.tensors_range: + continue + base_quantizer.tensors_range[node.output[0]] = TensorData( + lowest=np.float32(0.0), + highest=np.float32(1.0), + avg=np.float32(0.0), + std=np.float32(1.0), + ) + + # Patching nan values in TensorData + # These nan values should not appear in calibration with real inputs + for tensor_name in base_quantizer.tensors_range: + td = base_quantizer.tensors_range[tensor_name] + if np.isnan(td.range_value).any(): + base_quantizer.tensors_range[tensor_name] = TensorData( + lowest=np.float32(0.0), + highest=np.float32(448.0), + ) + + +def _create_calibrator_with_extra_options( + model: Union[str, Path], + op_types_to_calibrate: Optional[Sequence[str]] = None, + augmented_model_path="augmented_model.onnx", + calibrate_method=CalibrationMethod.MinMax, + use_external_data_format=False, + extra_options={}, # noqa: B006 +): + """This function overwrite is needed to pass the TRT plugin path and EP list to the inference session.""" + calibrator = None + if calibrate_method == CalibrationMethod.MinMax: + # default settings for min-max algorithm + symmetric = extra_options.get("symmetric", False) + moving_average = extra_options.get("moving_average", False) + averaging_constant = extra_options.get("averaging_constant", 0.01) + max_intermediate_outputs = extra_options.get("max_intermediate_outputs", None) + calibrator = MinMaxCalibrater( + model, + op_types_to_calibrate, + augmented_model_path, + use_external_data_format=use_external_data_format, + symmetric=symmetric, + moving_average=moving_average, + averaging_constant=averaging_constant, + max_intermediate_outputs=max_intermediate_outputs, + ) + elif calibrate_method == CalibrationMethod.Entropy: + # default settings for entropy algorithm + num_bins = extra_options.get("num_bins", 128) + num_quantized_bins = extra_options.get("num_quantized_bins", 128) + symmetric = extra_options.get("symmetric", False) + calibrator = EntropyCalibrater( + model, + op_types_to_calibrate, + augmented_model_path, + use_external_data_format=use_external_data_format, + symmetric=symmetric, + num_bins=num_bins, + num_quantized_bins=num_quantized_bins, + ) + elif calibrate_method == CalibrationMethod.Percentile: + # default settings for percentile algorithm + num_bins = extra_options.get("num_bins", 2048) + percentile = extra_options.get("percentile", 99.999) + symmetric = extra_options.get("symmetric", True) + calibrator = PercentileCalibrater( + model, + op_types_to_calibrate, + augmented_model_path, + use_external_data_format=use_external_data_format, + symmetric=symmetric, + num_bins=num_bins, + percentile=percentile, + ) + + elif calibrate_method == CalibrationMethod.Distribution: + # default settings for percentile algorithm + num_bins = extra_options.get("num_bins", 2048) + scenario = extra_options.get("scenario", "same") + + calibrator = DistributionCalibrater( + model, + op_types_to_calibrate, + augmented_model_path, + use_external_data_format=use_external_data_format, + num_bins=num_bins, + scenario=scenario, + ) + + if calibrator: + calibrator.augment_graph() + # ======== Modification: additional parameter with TRT plugin path ======== + calibrator.create_inference_session(**extra_options) + # ========================================================================= + return calibrator + + raise ValueError(f"Unsupported calibration method {calibrate_method}") + + +def _quantize_static( + model_input: Union[str, Path, onnx.ModelProto], + model_output: Union[str, Path], + calibration_data_reader: CalibrationDataReader, + quant_format=QuantFormat.QDQ, + op_types_to_quantize=None, + per_channel=False, + reduce_range=False, + activation_type=QuantType.QInt8, + weight_type=QuantType.QInt8, + nodes_to_quantize=None, + nodes_to_exclude=None, + use_external_data_format=False, + calibrate_method=CalibrationMethod.MinMax, + extra_options=None, +): + """Modification: enables TRT custom ops in the calibrator via 'TrtExtraPluginLibraryPaths' in extra_options. + + See ort.quantization.quantize.quantize_static for full function description. Additional info: + + extra_options: + key value pair dictionary for various options in different case. Current used: + ... + TrtExtraPluginLibraryPaths = string : + Default is None. Set TensorRT plugin paths if required. + ExecutionProviders = list[string] : + Default is [("CUDAExecutionProvider", {"device_id": 0}), "CPUExecutionProvider", + "TensorrtExecutionProvider"] + """ + if activation_type == QuantType.QFLOAT8E4M3FN or weight_type == QuantType.QFLOAT8E4M3FN: + if calibrate_method != CalibrationMethod.Distribution: + raise ValueError( + "Only Distribution calibration method is supported for float quantization." + ) + + extra_options = extra_options or {} + nodes_to_exclude = nodes_to_exclude or [] + nodes_to_quantize = nodes_to_quantize or [] + op_types_to_quantize = op_types_to_quantize or [] + mode = QuantizationMode.QLinearOps + + if not op_types_to_quantize or len(op_types_to_quantize) == 0: + q_linear_ops = list(QLinearOpsRegistry.keys()) + qdq_ops = list(QDQRegistry.keys()) + op_types_to_quantize = list(set(q_linear_ops + qdq_ops)) + + model = ( + save_and_reload_model_with_shape_infer(model_input) + if isinstance(model_input, onnx.ModelProto) + else load_model_with_shape_infer(Path(model_input)) + ) + + pre_processed: bool = model_has_pre_process_metadata(model) + if not pre_processed: + logging.warning( + "Please consider to run pre-processing before quantization. Refer to example: " + "https://github.com/microsoft/onnxruntime-inference-examples/blob/main/quantization/image_classification" + "/cpu/ReadMe.md " + ) + + calib_extra_options_keys = [ + ("CalibTensorRangeSymmetric", "symmetric"), + ("CalibMovingAverage", "moving_average"), + ("CalibMovingAverageConstant", "averaging_constant"), + ("CalibMaxIntermediateOutputs", "max_intermediate_outputs"), + # ====================== Modification ====================== + ("TrtExtraPluginLibraryPaths", "trt_extra_plugin_lib_paths"), + ("ExecutionProviders", "execution_providers"), + # ========================================================== + ] + calib_extra_options = { + key: extra_options.get(name) + for (name, key) in calib_extra_options_keys + if name in extra_options + } + + with tempfile.TemporaryDirectory(prefix="ort.quant.") as quant_tmp_dir: + if isinstance(model_input, onnx.ModelProto): + output_path = str(Path(quant_tmp_dir) / "model_input.onnx") + onnx.save_model( + model_input, + output_path, + save_as_external_data=True, + ) + model_input = output_path + + calibrator = calibrate.create_calibrator( + Path(model_input), + # ======== Modification ======== + nodes_to_quantize, + # ======== Modification ======== + augmented_model_path=Path(quant_tmp_dir).joinpath("augmented_model.onnx").as_posix(), + calibrate_method=calibrate_method, + use_external_data_format=use_external_data_format, + extra_options=calib_extra_options, + ) + calibrator.collect_data(calibration_data_reader) + tensors_range = calibrator.compute_data() + if not isinstance(tensors_range, TensorsData): + raise TypeError( + f"Unexpected type {type(tensors_range)} for tensors_range and calibrator={type(calibrator)}." + ) + del calibrator + + check_static_quant_arguments(quant_format, activation_type, weight_type) + + if quant_format is QuantFormat.QOperator: + quantizer = QDQQuantizer( + model, + per_channel, + reduce_range, + mode, + True, # static + weight_type, + activation_type, + tensors_range, + nodes_to_quantize, + nodes_to_exclude, + op_types_to_quantize, + extra_options, + ) + else: + quantizer = QDQQuantizer( + model, + per_channel, + reduce_range, + weight_type, + activation_type, + tensors_range, + nodes_to_quantize, + nodes_to_exclude, + op_types_to_quantize, + extra_options, + ) + + quantizer.quantize_model() + quantizer.model.save_model_to_file(model_output, use_external_data_format) + if not pre_processed: + logging.warning( + "Please consider pre-processing before quantization. See " + "https://github.com/microsoft/onnxruntime-inference-examples/blob/main/quantization/image_classification" + "/cpu/ReadMe.md " + ) + + +def patch_ort_modules(): + """Patches the ORT modules.""" + HistogramCollector.collect_value = _collect_value + calibrate.create_calibrator = _create_calibrator_with_extra_options + CalibraterBase.create_inference_session = _create_inference_session_with_ep_config + CalibraterBase.select_tensors_to_calibrate = _select_tensors_to_calibrate + QDQQuantizer.check_opset_version = _check_opset_version + MinMaxCalibrater.collect_data = _collect_data_minmax_calibrator + MinMaxCalibrater.merge_range = _merge_range_minmax_calibrator + HistogramCalibrater.collect_data = _collect_data_histogram_calibrator + + if ort.__version__ >= "1.18": + from onnxruntime.quantization.base_quantizer import BaseQuantizer + + BaseQuantizer.adjust_tensor_ranges = _adjust_tensor_ranges diff --git a/modelopt/onnx/quantization/ort_utils.py b/modelopt/onnx/quantization/ort_utils.py new file mode 100755 index 0000000000..e13cb2ee0e --- /dev/null +++ b/modelopt/onnx/quantization/ort_utils.py @@ -0,0 +1,137 @@ +# 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. + +"""Provides basic ORT inference utils, shoule be replaced by modelopt.torch.ort_client.""" + +from typing import Union + +import onnxruntime as ort +from onnxruntime.quantization.operators.qdq_base_operator import QDQOperatorBase +from onnxruntime.quantization.registry import QDQRegistry, QLinearOpsRegistry + +from modelopt.onnx.quantization.operators import QDQConvTranspose, QDQNormalization +from modelopt.onnx.quantization.ort_patching import patch_ort_modules + + +def _prepare_ep_list(calibration_eps: list[str]): + """Prepares the EP list for ORT from the given user input.""" + providers: list[Union[str, tuple[str, dict]]] = [] + for ep in calibration_eps: + if "cuda" in ep: + device_id = int(ep.split(":")[1]) if ":" in ep else 0 + providers.append(("CUDAExecutionProvider", {"device_id": device_id})) + if "dml" in ep: + device_id = int(ep.split(":")[1]) if ":" in ep else 0 + providers.append(("DmlExecutionProvider", {"device_id": device_id})) + if "cpu" in ep: + providers.append("CPUExecutionProvider") + if "trt" in ep: + providers.append("TensorrtExecutionProvider") + + return providers + + +def create_inference_session(onnx_path_or_model: Union[str, bytes], calibration_eps: list[str]): + """Create an ORT InferenceSession.""" + sess_options = ort.SessionOptions() + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + return ort.InferenceSession( + onnx_path_or_model, + sess_options=sess_options, + providers=_prepare_ep_list(calibration_eps), + ) + + +def get_quantizable_op_types(op_types_to_quantize: list[str]) -> list[str]: + """Returns a set of quantizable op types. + + Note. This function should be called after quantize._configure_ort() is called once. + This returns quantizable op types either from the user supplied parameter + or from modelopt.onnx's default quantizable ops setting. + """ + op_types_to_quantize = op_types_to_quantize or [] + + if not op_types_to_quantize or len(op_types_to_quantize) == 0: + q_linear_ops = list(QLinearOpsRegistry.keys()) + qdq_ops = list(QDQRegistry.keys()) + op_types_to_quantize = list(set(q_linear_ops + qdq_ops)) + + return op_types_to_quantize + + +def configure_ort( + op_types: list[str], + op_types_to_quantize: list[str], + trt_extra_plugin_lib_paths: str = None, + calibration_eps: list[str] = None, +): + """Configure and patches ORT to support ModelOpt ONNX quantization.""" + # Register some new QDQ operators on top of ORT + QDQRegistry["BatchNormalization"] = QDQNormalization + QDQRegistry["ConvTranspose"] = QDQConvTranspose + QDQRegistry["LRN"] = QDQNormalization # Example: caffenet-12.onnx + QDQRegistry["HardSwish"] = ( + QDQOperatorBase # Example: mobilenet_v3_opset17, efficientvit_b3_opset17 + ) + + # Patch ORT modules to fix bugs and support some edge cases + patch_ort_modules() + + # Remove copy, reduction and activation ops from ORT QDQ registry + for op_type in [ + "ArgMax", + "Concat", + "EmbedLayerNormalization", + "Gather", + "InstanceNormalization", + "LeakyRelu", + "Pad", + "Relu", + "Reshape", + "Sigmoid", + "Softmax", + "Split", + "Squeeze", + "Transpose", + "Unsqueeze", + "Where", + ]: + if op_type in QLinearOpsRegistry: + del QLinearOpsRegistry[op_type] + if op_type in QDQRegistry: + del QDQRegistry[op_type] + + # Prepare TensorRT friendly quantization settings + trt_guided_options = { + "QuantizeBias": False, + "ActivationSymmetric": True, + "OpTypesToExcludeOutputQuantization": op_types, # No output quantization + "AddQDQPairToWeight": True, # Instead of quantizing the weights, add QDQ node + "QDQOpTypePerChannelSupportToAxis": { + "Conv": 0, + "ConvTranspose": 1, + }, # per_channel should be True + "DedicatedQDQPair": True, + "ForceQuantizeNoInputCheck": ( + # By default, for some latent operators like MaxPool, Transpose, etc., + # ORT does not quantize if their input is not quantized already. + True + ), + "TrtExtraPluginLibraryPaths": trt_extra_plugin_lib_paths, + "ExecutionProviders": _prepare_ep_list(calibration_eps), + } + + quantizable_op_types = get_quantizable_op_types(op_types_to_quantize) + return trt_guided_options, quantizable_op_types diff --git a/modelopt/onnx/quantization/partitioning.py b/modelopt/onnx/quantization/partitioning.py new file mode 100644 index 0000000000..14e7ae9a3d --- /dev/null +++ b/modelopt/onnx/quantization/partitioning.py @@ -0,0 +1,411 @@ +# 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. + +"""Utilities related to partitioning the ONNX model to place QDQ nodes.""" + +from onnx_graphsurgeon.ir.graph import Graph +from onnx_graphsurgeon.ir.node import Node + +from modelopt.onnx.op_types import ( + is_copy_op, + is_linear_op, + is_pointwise_or_elementwise_op, + is_pooling_or_window_op, +) +from modelopt.onnx.quantization.graph_utils import ( + get_fusible_backbone, + has_const_input, + has_path_type, + is_const_input, +) +from modelopt.onnx.utils import get_child_nodes, get_variable_inputs + + +def _build_fusible_partition( + cur_node: Node, + fusible_partition: list[Node], + partitioned_nodes: set[str], + non_residual_inputs: dict[str, str], + graph: Graph, +) -> None: + """Traverses the graph starting from cur_node and updates the fusible_partition list. + + Add a nodes to the partition if any of these holds: + 1. The node is a unary or binary pointwise operation and fusible by cask + 2. The node is BN and/or Relu and fusible with preceding Conv op + 3. The node is a residual Add and fusible with current partition + + Args: + cur_node: Current candidate node for the partition. + fusible_partition: Current fusible partition. + partitioned_nodes: Set of already partitioned nodes. + non_residual_inputs: Non-residual input map. + graph: ONNX model graph. + + Returns: + Backbone node of the given pointwise op, None if not found. + """ + + def _is_on_non_residual_path(node: Node) -> bool: + if ( + node.op == "Add" # Input node should be an Add node + # The Add node should have a non-residual input + and non_residual_inputs[node.name] + # Input from the current node is non-residual + and cur_node.outputs[0].name == non_residual_inputs[node.name] + ): + return True + return False + + def _get_partition_node_outputs() -> list[str]: + # Collect tensor names produced by nodes in fusible_partition + # TODO: cache sub-partition outputs and append after them + partition_node_outputs = [] + for partition_node in fusible_partition: + for node_output in partition_node.outputs: + partition_node_outputs.append(node_output.name) + + return partition_node_outputs + + def _is_cask_fusible(node: Node, partition_node_outputs: list[str]) -> bool: + for tensor in node.inputs: + if tensor.name not in partition_node_outputs: + if not is_const_input(tensor): + return False + return True + + def _is_fusible_mul(mul_node: Node) -> bool: + # Don't consider Mul as fusible if it has the indicated ancestors. + # Otherwise, this causes regressions in: + # - densenet-12 and inception-v2-9 (dangling constants): [mul_node.op, "Unsqueeze"] + # - faster_vit: ["Mul", "Add", "Tanh", "Mul", "Add", "Mul", "Pow"] + # This improves perf in various models: mobilenet_v3, vovnet19b, coatnet_0, regnety_040. + + var_inps = get_variable_inputs(mul_node) + if len(var_inps) <= 1: + return True + + # Conv-Sigmoid-Mul chain is fusible + fusible_patterns = [ + ["Mul", "Sigmoid", "Conv"], + ["Mul", "Sigmoid", "BatchNormalization", "Conv"], + ] + if any([has_path_type(mul_node, graph, p, is_forward=False) for p in fusible_patterns]): + return True + + non_fusible_patterns = [["Mul", "Sigmoid"], ["Mul", "HardSigmoid"]] + if any([has_path_type(mul_node, graph, p, is_forward=False) for p in non_fusible_patterns]): + return False + + return True + + # Check the Mul nodes for their fusion compatibility + if cur_node.op == "Mul" and not _is_fusible_mul(cur_node): + return + + # Add current node to the partition + fusible_partition.append(cur_node) + partitioned_nodes.add(cur_node.name) + + # If on non-residual path, return after adding the node to the partition + # TODO: can the compiler fuse pointwise ops followed by residual Add? + if cur_node.op == "Add" and non_residual_inputs[cur_node.name]: + return + + consumer_nodes = get_child_nodes(cur_node) + partition_node_outputs = _get_partition_node_outputs() + + # TODO: traverse consumer nodes in topologically sorted order + for consumer_node in consumer_nodes: + if consumer_node.name in partitioned_nodes: + continue + + if ( + ( + is_pointwise_or_elementwise_op(consumer_node.op) + and _is_cask_fusible(consumer_node, partition_node_outputs) + ) + or ( + consumer_node.op in ["BatchNormalization", "Relu"] + and get_fusible_backbone(consumer_node, graph) + ) + or _is_on_non_residual_path(consumer_node) + ): + # DFS with the consumer and find more nodes for the partition + _build_fusible_partition( + consumer_node, + fusible_partition, + partitioned_nodes, + non_residual_inputs, + graph, + ) + + +def _find_quantizable_kgen_partitions( + graph: Graph, + cask_fusible_partitions: list[list[Node]], +) -> list[list[Node]]: + """Collect the partitions that can have quantized inputs, e.g. `XYZ` in `Conv-[epilogue]-XYZ`. + + Args: + graph: Onnx model graph. + cask_fusible_partitions: List of partitions that are fusible by CASK with Conv/MatMul backbone. + + Returns: + List of KGEN partitions that prefer quantized inputs. + """ + node_type_prefer_quantized_input = [ + "Resize", + "Upsample", + ] + + # collect the outputs from the cask fusible partitions. + # assumption: the nodes are topologily sorted in the lists. + cask_fusion_outputs = [partition[-1].outputs[0] for partition in cask_fusible_partitions] + + all_fusible_partitions = [] + for node in graph.nodes: + if node.op not in node_type_prefer_quantized_input: + continue + if node.inputs[0] in cask_fusion_outputs: + all_fusible_partitions.append([node]) + + return all_fusible_partitions + + +def find_fusible_partitions( + graph: Graph, + partitioned_nodes: set[str], + non_residual_inputs: dict[str, str], +) -> tuple[list[list[Node]], list[list[Node]]]: + """Traverses the graph and collects all cask/kgen fusible partitions. + + Args: + graph: Onnx model graph. + partitioned_nodes: Set of already partitioned nodes. + non_residual_inputs: Non-residual input map. + + Returns: + List of partitions that are fusible by CASK with Conv/MatMul backbone. + List of KGEN partitions with pointwise ops only. + """ + + def _partition_helper(fusible_root_type_checker): + all_fusible_partitions = [] # Collects all individual partitions + for node in graph.nodes: + # Check if the node is already in some other partition + if node.name in partitioned_nodes: + continue + + # Start a partition with a linear op + if not fusible_root_type_checker(node.op): + continue + + # Try building a partition starting with this current linear op + fusible_partition = [] + _build_fusible_partition( + node, + fusible_partition, + partitioned_nodes, + non_residual_inputs, + graph, + ) + + # Gather the non-empty partitions + if fusible_partition: + partitioned_nodes.update([node.name for node in fusible_partition]) + all_fusible_partitions.append(fusible_partition) + + return all_fusible_partitions + + cask_fusible_partitions = _partition_helper(is_linear_op) + kgen_partitions = _partition_helper(is_pointwise_or_elementwise_op) + kgen_partitions += _find_quantizable_kgen_partitions(graph, cask_fusible_partitions) + + return cask_fusible_partitions, kgen_partitions + + +def get_skiped_output_layers(graph: Graph, paritially_quantizable_nodes: list[Node]) -> list[str]: + """Returns the name of the non-quantizable output layers.""" + # TODO: see if input producer is already quantized or not + # TODO: filter out input layers if consumer is not quantized already + output_layers = [] + paritially_quantizable_node_names = set([node.name for node in paritially_quantizable_nodes]) + graph_output_names = [tensor.name for tensor in graph.outputs] + + for node in graph.nodes: + for tensor in node.outputs: + if tensor.name in graph_output_names: + if ( + node.op not in ["Conv", "Gemm", "MatMul"] + and node.name not in paritially_quantizable_node_names + ): + output_layers.append(node.name) + + return output_layers + + +def find_quantizable_nodes( + graph: Graph, + nodes_to_quantize: list[Node], + partitioned_nodes: set[str], + quantizable_op_types: list[str], +) -> list[Node]: + """Return the graph ops which are quantizable but not partitioned yet.""" + + # TODO: Check if any BatchNormalization is un-partitioned + # Note. Maxpool quantization has +/- + # Note. Prologue fusion is not handled, so some pointwise ops might be unnecessarily quantized + def _has_quantizable_consumer(node: Node, quantizable_node_set: set[str]) -> bool: + children = get_child_nodes(node) + for child_node in children: + if (child_node.name in quantizable_node_set) or ( + is_copy_op(child_node.op) + and _has_quantizable_consumer(child_node, quantizable_node_set) + ): + return True + + return False + + quantizable_nodes = [] + pooling_and_window_ops = [] + for node in graph.nodes: + if node.name in partitioned_nodes or node.op not in quantizable_op_types: + continue + + # Collect pooling and window ops for second pass + # as they need to check their neighbor's quantization status + if is_pooling_or_window_op(node.op): + pooling_and_window_ops.append(node) + continue + + if is_pointwise_or_elementwise_op(node.op) and has_const_input(node): + continue + + quantizable_nodes.append(node) + + quantizable_node_set = set( + [node.name for node in nodes_to_quantize] + [node.name for node in quantizable_nodes] + ) + for node in pooling_and_window_ops: + # TODO: Add or _has_quantizable_producer, ex. inception-v1-12.onnx + if _has_quantizable_consumer(node, quantizable_node_set): + quantizable_nodes.append(node) + + return quantizable_nodes + + +def find_hardcoded_patterns(graph: Graph) -> list[list[Node]]: + """Finds some non-quantizable pre-defined patterns!. + + Note. matching this tail pattern causes MTL_v1 -5.5% + ["ReduceSum", "Add", "Div", "Mul", "ReduceSum", "Sub", "Pow", "Mul", "ReduceSum", "Sqrt"] + """ + gelu = ["Div", "Erf", "Add", "Mul", "Mul"] + + matched_node_names = [] + for node in graph.nodes: + for path_type in [gelu]: + path_nodes = [] + if has_path_type( + node, + graph, + path_type, + is_forward=True, + wild_card_types=[], + path_nodes=path_nodes, + ): + matched_node_names.extend(path_nodes) + + return [matched_node_names] + + +def find_layer_norm_partitions(graph: Graph) -> list[list[Node]]: + """Finds the layer norm patterns in the graph.""" + # The most common LayerNorm implementation looks like this: + # t -> ReduceMean -> Sub -> Pow -> ReduceMean -> Add -> Sqrt -> Div -> Mul -> Add -> output + # \________________/ \______________________________________/ + # For simplicity, we do not match the connection between Sub and Div nodes. + layer_norm_chain_types = [ + "ReduceMean", + "Sub", + "Pow", + "ReduceMean", + "Add", + "Sqrt", + "Div", + "Mul", + "Add", + ] + mean_var_norm_chain_types = layer_norm_chain_types[:-2] + wild_card_types = ["Cast"] + layer_norm_partitions = [] + + for node in graph.nodes: + layer_norm_partition = [] + if node.op == "LayerNormalization": + layer_norm_partitions.append([node]) + elif node.op == layer_norm_chain_types[0] and has_path_type( + node, graph, layer_norm_chain_types, True, wild_card_types, layer_norm_partition + ): + layer_norm_partitions.append(layer_norm_partition) + elif node.op == mean_var_norm_chain_types[0] and has_path_type( + node, graph, mean_var_norm_chain_types, True, wild_card_types, layer_norm_partition + ): + layer_norm_partitions.append(layer_norm_partition) + + return layer_norm_partitions + + +def find_mha_partitions(graph: Graph) -> list[list[Node]]: + """Finds the MHA patterns in the graph that should not be quantized. + + A common MHA implementation looks like this: + t -> MatMul -> (optional) Pointwise ops (such as Add, Mul, Sub) -> Softmax -> MatMul -> output + Patterns that do not look like that should not be quantized (at least for now). + """ + mha_chain_types = [ + [ + "MatMul", + "MatMul", + ], + ] + mha_partitions = [] + + for node in graph.nodes: + if node.op == "MatMul": + for chain_type in mha_chain_types: + mha_partition = [] + if has_path_type(node, graph, chain_type, True, [], mha_partition): + mha_partitions.append(mha_partition) + + return mha_partitions + + +def find_non_quantizable_partitions_from_patterns(graph: Graph) -> list[list[str]]: + """Finds fusible partition from fixed patterns. + + Certain fused kernel counterpart is often a subgraph of native ops in onnx. + Those patterns are identified here and quantized to match compiler expectation. + """ + hard_coded_partitions = find_hardcoded_patterns(graph) + layer_norm_partitions = find_layer_norm_partitions(graph) + mha_partitions = find_mha_partitions(graph) + + partitions = [] + for partition_nodes in hard_coded_partitions + layer_norm_partitions + mha_partitions: + partitions.append([node.name for node in partition_nodes]) + + return partitions diff --git a/modelopt/onnx/quantization/qdq_utils.py b/modelopt/onnx/quantization/qdq_utils.py new file mode 100644 index 0000000000..67ccc74951 --- /dev/null +++ b/modelopt/onnx/quantization/qdq_utils.py @@ -0,0 +1,822 @@ +# 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. + +"""Various utils to support inserting Q/DQ nodes.""" + +import logging +from typing import Any, Sequence, Union + +import numpy as np +import onnx +import onnx_graphsurgeon as gs +from onnx.reference.custom_element_types import float8e4m3fn +from onnx.reference.ops.op_cast import Cast_19 as Cast + +from modelopt.onnx.quantization.graph_utils import ( + get_tensor_consumer_nodes, + get_tensor_producer_nodes, +) +from modelopt.onnx.quantization.quant_utils import ( + get_weights_scaling_factor, + get_weights_scaling_factor_2, + quantize, +) + +QUANTIZE_NODE_NAME = "QuantizeLinear" +DEQUANTIZE_NODE_NAME = "DequantizeLinear" + +onnx_dtype_map = { + "BFloat16": onnx.TensorProto.BFLOAT16, + "Float": onnx.TensorProto.FLOAT, + "Float8": onnx.TensorProto.FLOAT8E4M3FN, + "Half": onnx.TensorProto.FLOAT16, + "INT8": onnx.TensorProto.INT8, + "UINT8": onnx.TensorProto.UINT8, +} + + +def use_trt_qdq_ops(): + """Globally set node names to TRT custom names.""" + global QUANTIZE_NODE_NAME + QUANTIZE_NODE_NAME = "TRT_INT4QuantizeLinear" + global DEQUANTIZE_NODE_NAME + DEQUANTIZE_NODE_NAME = "TRT_INT4DequantizeLinear" + + +def _wq_name(name: str): + return name + "_i4" + + +def _scale_name(name: str): + return name + "_scale" + + +def _awq_scale_name(name: str): + return name + "_awq_scale" + + +def _zp_name(name: str): + return name + "_zp" + + +def _q_name(name: str): + return name + "_QuantizeLinear" + + +def _q_out_name(name: str): + return name + "_QuantizeLinear_Output" + + +def _dq_name(name: str): + return name + "_DequantizeLinear" + + +def _pqs_name(name: str): + return name + "_PQS" + + +def _dq_out_name(name: str): + return name + "_DequantizeLinear_Output" + + +def _pqs_out_name(name: str): + return name + "_PQS_Tensor" + + +def make_gs_quantized_weight(name: str, wq: np.ndarray, dtype) -> gs.Constant: + """Create a GraphSurgeon tensor from a quantized weight tensor. + + `name` is the desired _basename_ of the tensor. + """ + return gs.make_constant(_wq_name(name), np.asarray(wq), dtype) + + +def make_gs_zp(name: str, shape: Sequence[int], dtype) -> gs.Constant: + """Create a GraphSurgeon zero-point tensor of all zeroes with the given shape. + + `name` is the desired _basename_ of the tensor. + """ + return gs.make_constant( + _zp_name(name), + np.zeros(shape, dtype=onnx.mapping.TENSOR_TYPE_MAP[int(dtype)].np_dtype), + dtype, + ) + + +def make_gs_scale(name: str, scale: np.ndarray) -> gs.Constant: + """Create a GraphSurgeon scale tensor from the given numpy array. + + `name` is the desired _basename_ of the tensor. + """ + return gs.Constant(_scale_name(name), np.asarray(scale)) + + +def make_gs_awq_scale(name: str, scale: np.ndarray) -> gs.Constant: + """Create a GraphSurgeon scale tensor from the given numpy array. + + `name` is the desired _basename_ of the tensor. + """ + return gs.Constant(_awq_scale_name(name), np.asarray(scale)) + + +def make_gs_quantize_output( + name: str, shape: Sequence[int], dtype: onnx.TensorProto.DataType +) -> gs.Variable: + """Create a GraphSurgeon variable representing the output of a quantize node. + + `name` is the desired _basename_ of the node. + """ + return gs.make_variable(_q_out_name(name), dtype=dtype, shape=shape) + + +def make_gs_quantize_node( + name: str, inputs: Sequence[gs.Tensor], outputs: Sequence[gs.Tensor] +) -> gs.Node: + """Create a GraphSurgeon Quantize node. + + `name` is the desired _basename_ of the node. + """ + return gs.Node( + QUANTIZE_NODE_NAME, + name=_q_name(name), + inputs=inputs, + outputs=outputs, + ) + + +def make_gs_pre_quant_scale_output( + name: str, + shape: Sequence[int], + dtype: np.dtype, +) -> gs.Variable: + """Create a GraphSurgeon variable representing the output of a quantize node. + + `name` is the desired _basename_ of the node. + """ + return gs.Variable(_pqs_out_name(name), dtype=dtype, shape=shape) + + +def make_gs_dequantize_output( + name: str, + shape: Sequence[int], + dtype: np.dtype, +) -> gs.Variable: + """Create a GraphSurgeon variable representing the output of a quantize node. + + `name` is the desired _basename_ of the node. + """ + return gs.Variable(_dq_out_name(name), dtype=dtype, shape=shape) + + +def make_gs_pre_quant_scale_node( + name: str, inputs: Sequence[gs.Tensor], outputs: Sequence[gs.Tensor] +) -> gs.Node: + """Create a GraphSurgeon Dequantize node. + + `name` is the desired _basename_ of the node. + """ + return gs.Node( + "Mul", + name=_pqs_name(name), + inputs=inputs, + outputs=outputs, + ) + + +def make_gs_dequantize_node( + name: str, + inputs: Sequence[gs.Tensor], + outputs: Sequence[gs.Tensor], + attributes: dict[str, Any] = None, +) -> gs.Node: + """Create a GraphSurgeon Dequantize node. + + `name` is the desired _basename_ of the node. + """ + return gs.Node( + DEQUANTIZE_NODE_NAME, + name=_dq_name(name), + inputs=inputs, + outputs=outputs, + attrs=attributes, + ) + + +def _postprocess_qdq( + graph: gs.Graph, + orig_weight_names: set[str], + q_nodes: dict[str, gs.Node] = {}, + dq_nodes: dict[str, gs.Node] = {}, +): + # Inserts all newly created nodes to graph. + # Update all consumers of original initializers to point to the DQ nodes. + for node in graph.nodes: + for i in range(len(node.inputs)): + key = node.inputs[i].name + if key not in orig_weight_names: + continue + node.inputs[i] = dq_nodes[key].outputs[0] + + # Insert new nodes. + graph.nodes.extend(q_nodes.values()) + graph.nodes.extend(dq_nodes.values()) + + graph.cleanup() + graph.toposort() + + +def insert_pre_quant_scale_nodes( + graph: gs.Graph, input_tensors: dict[str, str], pre_quant_scale: dict[str, np.ndarray] +): + """Insert new mul nodes into graph. + + Args: + graph: The graph to modify. + input_tensors: A dictionary of weight tensor names mapped to corresponding input tensor names + pre_quant_scale: A map from ONNX input tensor name to corresponding pre-quant scale. + """ + + def _insert_helper( + weight_tensor_name: str, + input_tensor_name: str, + scale: np.ndarray, + mul_nodes: dict[str, gs.Node], + ): + pre_quant_scale_tensor = make_gs_awq_scale(weight_tensor_name, scale) + # TODO: Study effects of caching Gemm/Matmul nodes on perf and mem usage. + gemm_nodes = [node for node in graph.nodes if node.op in ["Gemm", "MatMul"]] + for node in gemm_nodes: + input_set = set([input.name for input in node.inputs]) + input_idxs = {input.name: idx for idx, input in enumerate(node.inputs)} + if _dq_out_name(weight_tensor_name) in input_set and input_tensor_name in input_set: + pqs_in = node.inputs[input_idxs[input_tensor_name]] + pqs_out = make_gs_pre_quant_scale_output( + weight_tensor_name, shape=pqs_in.shape, dtype=scale.dtype + ) + mul_node = make_gs_pre_quant_scale_node( + weight_tensor_name, inputs=[pqs_in, pre_quant_scale_tensor], outputs=[pqs_out] + ) + node.inputs[input_idxs[input_tensor_name]] = mul_node.outputs[0] + mul_nodes[weight_tensor_name] = mul_node + + mul_nodes = {} + for w_name, scale in pre_quant_scale.items(): + inv_scale = 1.0 / scale + _insert_helper(w_name, input_tensors[w_name], inv_scale, mul_nodes) + + graph.nodes.extend(mul_nodes.values()) + + graph.cleanup() + graph.toposort() + + +def insert_dq_nodes( + graph: gs.Graph, + scales: dict[str, np.ndarray], + quantized_weights: dict[str, np.ndarray], + attributes: dict[str, Any] = None, + zero_points: Union[dict[str, np.ndarray], None] = None, +): + """Insert new initializers and DQ nodes into graph. + + Args: + graph: The graph to modify. + weights: A map from ONNX initializer name to tensor. + scales: A map from ONNX initializer name to desired scale factor for that initializer. + dq_only: Whether to only insert dq nodes. + """ + + def _insert_helper( + name: str, + wq: np.ndarray, + scale: np.ndarray, + dq_nodes: dict[str, gs.Node], + attrs: dict[str, Any], + zp: np.ndarray, + ): + tensor_dtype = onnx.TensorProto.INT4 if zp is None else onnx.TensorProto.UINT4 + wq_tensor = make_gs_quantized_weight(name, wq, tensor_dtype) + scale_tensor = make_gs_scale(name, scale) + dq_out = make_gs_dequantize_output(name, shape=wq.shape, dtype=scale.dtype) + inputs = [wq_tensor, scale_tensor] + if zp is not None: + zp_tensor = gs.make_constant(_zp_name(name), zp, tensor_dtype) + inputs.append(zp_tensor) + dq_node = make_gs_dequantize_node( + name, + inputs=inputs, + outputs=[dq_out], + attributes=attrs, + ) + dq_nodes[name] = dq_node + + dq_nodes = {} + for name, scale in scales.items(): + zp = None + if zero_points is not None: + zp = zero_points.get(name) + assert zp is not None, "zero-point is enabled but zero-point values not found" + _insert_helper(name, quantized_weights[name], scale, dq_nodes, attributes, zp) + + _postprocess_qdq( + graph, + orig_weight_names=set(scales.keys()), + dq_nodes=dq_nodes, + ) + + +def insert_qdq_nodes( + graph: gs.Graph, + scales: dict[str, np.ndarray], + weight_map: dict[str, gs.Tensor], +): + """Insert scales and QDQ nodes into graph. + + Args: + graph: The graph to modify. + scales: A map from ONNX initializer name to desired scale factor for that initializer. + weight_map: A map from ONNX initializer name to graphsurgeon tensor. + """ + + def _insert_helper( + name: str, + weight_to_quantize: gs.Tensor, + scale: np.ndarray, + q_nodes: dict[str, gs.Node], + dq_nodes: dict[str, gs.Node], + ): + scale_tensor = make_gs_scale(name, scale) + zp_tensor = make_gs_zp(name, scale.shape, onnx.TensorProto.INT4) + q_out = make_gs_quantize_output(name, weight_to_quantize.shape, onnx.TensorProto.INT4) + q_node = make_gs_quantize_node( + name, inputs=[weight_to_quantize, scale_tensor, zp_tensor], outputs=[q_out] + ) + dq_out = make_gs_dequantize_output(name, shape=weight_to_quantize.shape, dtype=scale.dtype) + dq_node = make_gs_dequantize_node( + name, inputs=[q_out, scale_tensor, zp_tensor], outputs=[dq_out] + ) + q_nodes[name] = q_node + dq_nodes[name] = dq_node + + q_nodes, dq_nodes = {}, {} + for name, scale in scales.items(): + _insert_helper(name, weight_map[name], scale, q_nodes, dq_nodes) + + _postprocess_qdq( + graph, + orig_weight_names=set(scales.keys()), + q_nodes=q_nodes, + dq_nodes=dq_nodes, + ) + + +def replace_scale_values(graph: onnx.onnx_ml_pb2.GraphProto, act_scales_dict: dict[str, float]): + """Replaces the scales values from calibration cache.""" + initializers = graph.initializer + initializer_indices = { + initializer.name: idx for idx, initializer in enumerate(graph.initializer) + } + + for node in graph.node: + if node.op_type == "QuantizeLinear": + scale_input_name = node.input[1] + if scale_input_name in act_scales_dict: + idx = initializer_indices.get(scale_input_name, None) + assert ( + idx is not None + ), f"Expected '{scale_input_name}' to be found in 'graph.initializer', but it was not present." + scale = onnx.numpy_helper.from_array( + np.float32(act_scales_dict[scale_input_name]), scale_input_name + ) + initializers[idx].CopyFrom(scale) + else: + # If the scale is not present in the act_scales_dict + # then the current node must be an weight quantizer and + # the weight should be available in the graph initializer + assert ( + initializer_indices.get(node.input[0], None) is not None + ), f"Tensor {node.input[0]} not found in initializers." + + +def qdq_to_dq( + onnx_model: onnx.onnx_pb.ModelProto, verbose: bool = False +) -> onnx.onnx_pb.ModelProto: + """Convert FP32/FP16 weights of the given ONNX model to INT8/FP8 weights. + + Q nodes will get removed from the weights and have only DQ nodes with those converted INT8/FP8 + weights in the output model. Also dangling Q nodes get fused and update its consumer's weight. + + Args: + onnx_model: ONNX model protobuf. + + Returns: + ONNX model protobuf with only DQ nodes for weights and QDQ nodes for activations. + """ + graph = onnx_model.graph + initializers = graph.initializer + initializer_indices = { + initializer.name: idx for idx, initializer in enumerate(graph.initializer) + } + + def _get_tensor_type(tensor_name): + for value_info in graph.value_info: + if value_info.name == tensor_name: + return value_info.type.tensor_type.elem_type + return None + + def _remove_unnecessary_cast(): + # Remove two pattern of unnecessary Cast node + cast_indices = [] + + tensor_consumers = get_tensor_consumer_nodes(graph) + + # find all Cast node with same input and output type + for node_idx, node in enumerate(graph.node): + if node.op_type != "Cast": + continue + + # if input type matches attribute "to", this is a useless Cast node + assert len(node.input) == 1 + input_name = node.input[0] + idx = initializer_indices.get(input_name, None) + if idx is not None: + data_type = initializers[idx].data_type + else: + data_type = _get_tensor_type(input_name) + + attr = node.attribute[0] + assert attr.name == "to" + + # Pattern 1: Input and Output Type are the same. + if data_type == attr.i: + cast_indices.append(node_idx) + else: + # Pattern 2: Input and Output Type differ but Cast node doesn't have a producer + # We do the conversion and fuse Cast node. + if idx is not None: + cast_indices.append(node_idx) + # Replace Q node input with new input + cast_input = onnx.numpy_helper.to_array(initializers[idx]) + + dtype = onnx.helper.tensor_dtype_to_np_dtype(attr.i) + converted_tensor = onnx.numpy_helper.from_array( + cast_input.astype(dtype), input_name + ) + + initializers[idx].CopyFrom(converted_tensor) + else: + continue + + # Renew input of consumer nodes + output_name = node.output[0] + consumers = tensor_consumers[output_name] + for q_node in consumers: + for i in range(len(q_node.input)): + if q_node.input[i] == output_name: + q_node.input[i] = input_name + break + + # Delete Cast node + for node_idx in sorted(cast_indices, reverse=True): + del graph.node[node_idx] + + def _convert(node: onnx.onnx_ml_pb2.NodeProto): + if verbose: + logging.info(f"Processing {node.name}") + + idx1 = initializer_indices.get(node.input[0], None) + assert ( + idx1 is not None + ), f"Expected '{node.input[0]}' to be found in 'graph.initializer', but it was not present." + w = initializers[idx1] + + w32 = onnx.numpy_helper.to_array(w) + + idx2 = initializer_indices.get(node.input[1], None) + if idx2 is not None: + y_scale = initializers[idx2] + else: + producer_node = tensor_producers[node.input[1]] + attr = producer_node.attribute[0] + assert attr.name == "value" + y_scale = attr.t + + np_y_scale = onnx.numpy_helper.to_array(y_scale) + + idx3 = initializer_indices.get(node.input[2], None) + if idx3 is not None: + zero_point = initializers[idx3] + else: + producer_node = tensor_producers[node.input[2]] + attr = producer_node.attribute[0] + assert attr.name == "value" + zero_point = attr.t + + np_zero_point = onnx.numpy_helper.to_array(zero_point) + + dq_node = tensor_consumers[node.output[0]][0] + next_node = tensor_consumers[dq_node.output[0]][0] + + # No transpose is needed for 2D "MatMul", only for 3D (fails with PETR otherwise) + transpose_nodes = ["Conv", "Transpose", "Gemm"] + is_3d_matmul = next_node.op_type in "MatMul" and len(np.shape(w32)) == 3 + do_transpose = next_node.op_type in transpose_nodes or is_3d_matmul + + if do_transpose: + w32 = np.transpose(w32, axes=[0, 2, 1]) if is_3d_matmul else np.transpose(w32) + + # Scale should be a scaler or vector with the same length as the last dimension of the weight + assert not np_y_scale.shape or w32.shape[-1] == np_y_scale.shape[0] + + fp8 = np_zero_point.dtype == float8e4m3fn + + if fp8: + scaled = np.asarray(w32 / np_y_scale) + np_zero_point + else: + scaled = np.asarray((w32 / np_y_scale).round()) + np.clip(scaled + np_zero_point, -128, 127, out=scaled) + + if do_transpose: + scaled = np.transpose(scaled, axes=[0, 2, 1]) if is_3d_matmul else np.transpose(scaled) + + if fp8: + w8 = onnx.numpy_helper.from_array( + Cast.eval(scaled, to=onnx.TensorProto.FLOAT8E4M3FN), w.name + ) + else: + w8 = onnx.numpy_helper.from_array(scaled.astype("int8"), w.name) + + initializers[idx1].CopyFrom(w8) + + return idx2, idx3 + + _remove_unnecessary_cast() + + tensor_producers = get_tensor_producer_nodes(graph) + tensor_consumers = get_tensor_consumer_nodes(graph) + + dangling_q_indices = [] + dangling_init_indices = [] + + for node_idx, node in enumerate(graph.node): + if node.op_type == "QuantizeLinear": + weight_name = node.input[0] + + # Const input to quantize linear means weighted layer + if weight_name not in tensor_producers: + scale_init_idx, zp_init_idx = _convert(node) + dangling_q_indices.append(node_idx) + dangling_init_indices.extend([scale_init_idx, zp_init_idx]) + + # Update following DQ nodes input name, each q should only have one dq consumer + consumers = tensor_consumers[node.output[0]] + assert len(consumers) == 1 + dq_node = consumers[0] + assert dq_node.op_type == "DequantizeLinear" + dq_node.input[0] = weight_name + + # Remove Q nodes + for node_idx in sorted(dangling_q_indices, reverse=True): + del graph.node[node_idx] + + return onnx_model + + +def replace_fp4qdq_with_2dq( + graph: onnx.onnx_ml_pb2.GraphProto, + node: onnx.onnx_ml_pb2.NodeProto, + initializer_indices: dict[str, int], + value_info_map: dict[str, onnx.onnx_ml_pb2.ValueInfoProto], + graph_inputs: set[str], + w_f4: np.ndarray, + sw_f32_per_tensor: np.ndarray, + sw_f8_per_block: np.ndarray, + precision_dtype: str, + block_size: int, +): + """Replaces the given node in the ONNX graph with a subgraph consisting of two DequantizeLinear nodes. + + Args: + graph: The ONNX graph containing the node to replace. + node: The node to be replaced. + initializer_indices: A dictionary mapping initializer names to their indices in the graph. + value_info_map: A dictionary mapping value info names to their ValueInfoProto objects. + graph_inputs: A set of graph input names. + w_f4: NumPy array for w_f4. + sw_f32_per_tensor: NumPy array for sw_f32_per_tensor. + sw_f8_per_block: NumPy array for sw_f8_per_block. + precision_dtype: The precision of the weights. + block_size: Block size used in block quantization. + """ + + def _add_initializer(initializer): + if initializer.name not in initializer_indices: + graph.initializer.append(initializer) + + def _add_input_value_info(graph, tensor_proto): + assert ( + tensor_proto.name not in graph_inputs + ), f"{tensor_proto.name} already in graph inputs." + assert ( + tensor_proto.name not in value_info_map + ), f"{tensor_proto.name} already in value info." + + value_info = onnx.helper.make_tensor_value_info( + tensor_proto.name, tensor_proto.data_type, tensor_proto.dims + ) + graph.input.append(value_info) + + # Remove the original node from the graph + graph.node.remove(node) + weight_name = node.input[0] + + # Generate unique names for the initializers + w_f4_name = weight_name + "_f4" + sw_f8_per_block_name = weight_name + "_f8_scale" + sw_f32_per_tensor_name = sw_f8_per_block_name + "_f32_scale" + + # Create TensorProto for initializers + w_f4_proto = onnx.numpy_helper.from_array(w_f4, w_f4_name) + sw_f32_per_tensor_proto = onnx.numpy_helper.from_array( + sw_f32_per_tensor, sw_f32_per_tensor_name + ) + sw_f8_per_block_proto = onnx.numpy_helper.from_array(sw_f8_per_block, sw_f8_per_block_name) + + # Add ValueInfo for the initializers if not present + _add_input_value_info(graph, w_f4_proto) + _add_input_value_info(graph, sw_f32_per_tensor_proto) + _add_input_value_info(graph, sw_f8_per_block_proto) + + # Add the initializers to the graph + _add_initializer(w_f4_proto) + _add_initializer(sw_f32_per_tensor_proto) + _add_initializer(sw_f8_per_block_proto) + + # Create DequantizeLinear_1 node: (sw_f8_per_block, sw_f32_per_tensor) -> sw_f16 + sw_f16_name = weight_name + "_f16_scale" + dequant1 = onnx.helper.make_node( + "DequantizeLinear", + inputs=[sw_f8_per_block_proto.name, sw_f32_per_tensor_proto.name], + outputs=[sw_f16_name], + name=weight_name + "_DequantizeLinear", + ) + + # Create DequantizeLinear_2 node: (w_f4, sw_f16) -> w_16 + w16_name = node.output[0] + dequant2 = onnx.helper.make_node( + "DequantizeLinear", + inputs=[w_f4_proto.name, sw_f16_name], + outputs=[w16_name], + name=weight_name + "_DequantizeLinear_1", + axis=-1, + block_size=block_size, + ) + + # Add value_info for sw_f16 + # Assuming sw_f16 has the same shape as sw_f8_per_block + sw_f16_type_proto = onnx.helper.make_tensor_type_proto( + elem_type=onnx_dtype_map[precision_dtype], shape=sw_f8_per_block.shape + ) + sw_f16_value_info = onnx.helper.make_value_info(name=sw_f16_name, type_proto=sw_f16_type_proto) + graph.value_info.append(sw_f16_value_info) + + # Change the data type of w16 (output of 2nd DQ) to model weight precision type + if w16_name in value_info_map: + value_info_map[w16_name].type.tensor_type.elem_type = onnx_dtype_map[precision_dtype] + else: + raise ValueError(f"ValueInfo for {w16_name} not found.") + + # Add the new nodes to the graph + graph.node.extend([dequant1, dequant2]) + + +def fp4qdq_to_2dq(onnx_model: onnx.onnx_pb.ModelProto) -> onnx.onnx_pb.ModelProto: + """Convert FP32/FP16 weights of the given ONNX model to FP4 weights and scaling factors. + + TRT_FP4QDQ nodes will get removed from the weights and have two DQ nodes with those converted FP4 + weights and scaling factors in the output model. + + Args: + onnx_model: ONNX model protobuf. + + Returns: + ONNX model protobuf with DQ nodes for weights and DynQ + DQ nodes for activations. + """ + graph = onnx_model.graph + initializers = graph.initializer + initializers_to_delete = [] + tensor_consumers = get_tensor_consumer_nodes(graph) + initializer_indices = { + initializer.name: idx for idx, initializer in enumerate(graph.initializer) + } + value_info_map = {vi.name: vi for vi in graph.value_info} + graph_inputs = {inp.name for inp in graph.input} + + def _cast_input_dtypes(node: onnx.onnx_ml_pb2.NodeProto, precision_dtype: str): + # Change the input types to match weight precision (precision_dtype) + if node.op_type == "Transpose": + maybe_matmul = tensor_consumers[node.output[0]][0] + assert maybe_matmul.op_type == "MatMul" + node = maybe_matmul + + # Create Cast nodes for each input of the target node except bias + for i, input_name in enumerate(node.input[:2]): + cast_output_name = input_name + "_f16" # Unique name for the cast output + + # Create a Cast node to convert the input to FP16/BF16 + cast_node = onnx.helper.make_node( + "Cast", + inputs=[input_name], # Original input of the target node + outputs=[cast_output_name], + to=onnx_dtype_map[precision_dtype], # Cast to FP16/BF16 + ) + + # Insert the Cast node into the graph + graph.node.extend([cast_node]) + + # Update the target node input to use the cast node output + node.input[i] = cast_output_name + + def _get_precision_dtype() -> str: + # Check initializers to determine the precision of the weights + precision_dtype = "Half" + for initializer in graph.initializer: + if initializer.data_type == 16: + precision_dtype = "BFloat16" + break # Assuming all weights are of the same precision + + return precision_dtype + + def _bfloat16_to_float32(bf16_array): + uint32_array = bf16_array.astype(np.uint32) << 16 + return uint32_array.view(np.float32) + + def _read_f16_tensor_as_fp32(tensor): + if tensor.data_type == onnx.TensorProto.BFLOAT16: + raw_data = tensor.raw_data + uint16_array = np.frombuffer(raw_data, dtype=np.uint16) + float32_array = _bfloat16_to_float32(uint16_array) + tensor_shape = tuple(dim for dim in tensor.dims) + return float32_array.reshape(tensor_shape) + + # Read FLOAT16 tensor and return + return onnx.numpy_helper.to_array(tensor).astype(np.float32) + + print("Post-processing TRT_FP4QDQ nodes for TRT deployment...") + precision_dtype = _get_precision_dtype() + fp4_qdq_nodes = [node for node in graph.node if node.op_type == "TRT_FP4QDQ"] + + for node in fp4_qdq_nodes: + idx1 = initializer_indices.get(node.input[0], None) + assert idx1 is not None, f"Initializer for weight '{node.input[0]}' not found." + block_size = node.attribute[0].i + initializers_to_delete.append(initializers[idx1].name) + + tensor = initializers[idx1] + w32 = _read_f16_tensor_as_fp32(tensor) + sw_f32_per_tensor = get_weights_scaling_factor_2(w32) + sw_f32_per_block = get_weights_scaling_factor(w32, block_size, sw_f32_per_tensor) + w_f32 = quantize(w32, block_size, sw_f32_per_block, sw_f32_per_tensor) + + # Real quantize the tensors + w_f4 = Cast.eval(w_f32, to=onnx.TensorProto.FLOAT4E2M1) + sw_f8_per_block = Cast.eval(sw_f32_per_block, to=onnx.TensorProto.FLOAT8E4M3FN) + + replace_fp4qdq_with_2dq( + graph, + node, + initializer_indices, + value_info_map, + graph_inputs, + w_f4, + sw_f32_per_tensor, + sw_f8_per_block, + precision_dtype, + block_size, + ) + + # We need to change the bias etc. type + next_node = tensor_consumers[node.output[0]][0] + _cast_input_dtypes(next_node, precision_dtype) + + print(f"Replaced {node.name} with 2 DQ nodes.") + + new_initializers = [ + init for init in graph.initializer if init.name not in initializers_to_delete + ] + graph.ClearField("initializer") + graph.initializer.extend(new_initializers) + + return onnx_model diff --git a/modelopt/onnx/quantization/quant_utils.py b/modelopt/onnx/quantization/quant_utils.py new file mode 100644 index 0000000000..a02565b7fc --- /dev/null +++ b/modelopt/onnx/quantization/quant_utils.py @@ -0,0 +1,174 @@ +# 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. + +"""Provides some basic utilities that can be used in quantize() methods.""" + +from typing import Sequence, Union + +import numpy as np + +INT4_MIN = -8 +INT4_MAX = 7 +UINT4_MIN = 0 +UINT4_MAX = 15 + + +def pack_float32_to_4bit_optimized(array: Union[np.ndarray, Sequence], signed: bool) -> np.ndarray: + """Convert an array of float32 value to a 4bit data-type and pack every two concecutive elements in a byte. + + This is the optimized version of pack_float32_to_4bit() utility in ONNX helper file. The basic optimizations + done here mainly rely on moving some common code out of the per-element function calls or loops, thereby making + them per-input-array, instead of per-input-element. The remaining logic should largely remain as is. + + Args: + array: array of float to convert and pack + signed: Whether the 4 bit variant is signed or unsigned + + Returns: + Packed array with size `ceil(array.size/2)` (single dimension). + """ + if not isinstance(array, np.ndarray): + array = np.asarray(array, dtype=np.float32) + + array_flat = array.ravel() + is_odd_volume = np.prod(array.shape) % 2 == 1 + if is_odd_volume: + array_flat = np.append(array_flat, np.array([0])) + + inp_arr_len = array_flat.size + dtype = np.int8 if signed else np.uint8 + clip_low = INT4_MIN if signed else UINT4_MIN + clip_high = INT4_MAX if signed else UINT4_MAX + array_flat = np.clip(array_flat, clip_low, clip_high) + array_flat = np.rint(array_flat).astype(dtype) + assert len(array_flat) % 2 == 0, "array length must be even at this point" + assert len(array_flat) == inp_arr_len, "output-length must match the input-length" + output_list = [] + for i in range(0, inp_arr_len, 2): + output_list.append((array_flat[i + 1] << 4) | (array_flat[i] & 0x0F)) + arr = np.array(output_list) + return arr.astype(np.uint8) + + +def pack_float32_to_4bit_cpp_based(array: Union[np.ndarray, Sequence], signed: bool) -> np.ndarray: + """Convert an array of float32 value to a 4bit data-type and pack every two concecutive elements in a byte. + + This is the optimized version of pack_float32_to_4bit() utility in ONNX helper file. The basic optimizations + here is to implement this round_and_pack logic in C++, which is supposed to be faster. + + Args: + array: array of float to convert and pack + signed: Whether the 4 bit variant is signed or unsigned + + Returns: + Packed array with size `ceil(array.size/2)` (single dimension). + """ + from .extensions import round_and_pack_ext + + if not isinstance(array, np.ndarray): + array = np.asarray(array, dtype=np.float32) + + # - Currently, FP32, FP64, UINT8 and INT8 dtypes have C++ implementation of the round-and-pack + # logic. + # - With above, FP16 should also get "implicitly" supported due to possible type promotion to higher + # precision float types (e.g. float32 or float64). So, it is mentioned as supported below. + # - We can add support for other dtypes as and when needed. + use_python_version = False + if round_and_pack_ext is None or array.dtype not in [ + "float", + "float16", + "float32", + "float64", + "int8", + "uint8", + ]: + use_python_version = True + + array_flat = array.ravel() + is_odd_volume = np.prod(array.shape) % 2 == 1 + if is_odd_volume: + array_flat = np.append(array_flat, np.array([0], array_flat.dtype)) + + inp_arr_len = array_flat.size + + assert inp_arr_len % 2 == 0, "input array length must be even at this point" + + if use_python_version: + print( + f"Using python optimized version for round_and_pack...input-array-dtype={array_flat.dtype}\n" + ) + numpy_out = pack_float32_to_4bit_optimized(array_flat, signed) + else: + numpy_out = np.zeros([1, int(inp_arr_len / 2)], dtype=np.int8) + numpy_out = numpy_out.ravel() + ret = round_and_pack_ext.round_and_pack( + signed, array_flat, array_flat.size, numpy_out, numpy_out.size + ) + assert ret == inp_arr_len / 2, "Unexpected output length" + numpy_out = numpy_out.astype(np.uint8) + + return numpy_out + + +def get_weights_scaling_factor_2(input: np.ndarray): + """Returns per tensor weight scaling factor.""" + per_block_scale_amax = np.max(np.abs(input)) / 6.0 + per_block_quant_scale = per_block_scale_amax / 448.0 + if per_block_quant_scale == 0: + per_block_quant_scale = 1.0 + return np.float32(per_block_quant_scale) + + +def get_weights_scaling_factor( + input: np.ndarray, block_size: int, weights_scaling_factor_2: np.float32 +): + """Returns quantized per block weight scaling factor.""" + # Get per_block amax + [n, k] = input.shape + assert block_size != 0, "Block size is zero. Cannot return per_block amax for given input." + + assert ( + k % block_size == 0 + ), "Weight shape is not divisible for block size for block quantiation." + + input = input.reshape(n, k // block_size, block_size) + # Get per block amax + per_block_amax = np.max(np.abs(input), axis=-1) + # Get per-block-scale + per_block_scale = per_block_amax / 6.0 + # Quantize per_block_scale to FP8 + q_per_block_scale = per_block_scale / weights_scaling_factor_2 + # Set all zero values in scale to 1.0 + q_per_block_scale[per_block_scale == 0] = 1.0 + return q_per_block_scale.astype(np.float32) + + +def quantize( + input: np.ndarray, + block_size: int, + weights_scaling_factor: np.ndarray, + weights_scaling_factor_2: np.ndarray, +): + """Converting a tensor to a quantized format based on NVFP4 quantization.""" + # Reshape the weight and scale factors + input = input.reshape(tuple(input.shape[:-1]) + (-1, block_size)) + + # Scale weights + scaled_weight = input / ( + np.expand_dims(weights_scaling_factor, axis=-1) * weights_scaling_factor_2 + ) + + # Reshape weights to original and return + return scaled_weight.reshape(scaled_weight.shape[0], -1) diff --git a/modelopt/onnx/quantization/quantize.py b/modelopt/onnx/quantization/quantize.py new file mode 100755 index 0000000000..7d316f8742 --- /dev/null +++ b/modelopt/onnx/quantization/quantize.py @@ -0,0 +1,369 @@ +# 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. + +"""This module is used to convert an ONNX model (w/o QDQ nodes) and calibration data into an ONNX model with QDQ nodes. + +Quantization typically targets linear operations like Convolution (Conv), Matrix Multiplication (MatMul), etc., as these +transformations often yield significant performance improvements. However, many other operations can also be quantized +(i.e., they have low-precision kernels available), which can provide optimal performance with a minimal drop in +accuracy. + +The default operation types that this ONNX post-training quantization (PTQ) tool quantizes in different quantization +modes are as follows: +- INT8 mode: Add, AveragePool, BatchNormalization, Clip, Conv, ConvTranspose, Gemm, GlobalAveragePool, MatMul, MaxPool, +Mul +- INT4 mode: Gemm, MatMul +- FP8 mode: Conv, Gemm, MatMul + +This tool inserts Quantize-Dequantize (QDQ) nodes following compiler-friendly patterns and generates an explicit ONNX +model. +""" + +import logging +import os +import platform +import shutil +import tempfile +from typing import Any, Optional + +import numpy as np +import onnx +import onnx.onnx_cpp2py_export.checker as C # noqa: N812 +import onnx_graphsurgeon as gs + +from modelopt.onnx.quantization.calib_utils import ( + CalibrationDataProvider, + CalibrationDataType, + RandomDataProvider, +) +from modelopt.onnx.quantization.fp8 import quantize as quantize_fp8 +from modelopt.onnx.quantization.graph_utils import ( + add_fp16_fp32_cast, + find_nodes_from_mha_to_exclude, + print_stat, +) +from modelopt.onnx.quantization.int4 import quantize as quantize_int4 +from modelopt.onnx.quantization.int8 import quantize as quantize_int8 +from modelopt.onnx.quantization.qdq_utils import qdq_to_dq +from modelopt.onnx.quantization.trt_utils import load_onnx_model +from modelopt.onnx.utils import duplicate_shared_constants, name_onnx_nodes, save_onnx + +__all__ = ["quantize"] + + +# Set logging level to info +logging.getLogger().setLevel(logging.INFO) + + +def _preprocess_onnx( + onnx_path: str, + use_external_data_format: bool, + output_path: str, + trt_plugins: Optional[str], + trt_plugins_precision: Optional[list[str]], +) -> tuple[str, list[str], bool]: + intermediate_generated_files = [] + output_dir = os.path.dirname(output_path) + model_name = os.path.splitext(os.path.basename(onnx_path))[0] + + # Load the model and weights + onnx_model, has_custom_op, custom_ops = load_onnx_model( + onnx_path, trt_plugins, use_external_data_format + ) + if has_custom_op: + onnx_path = os.path.join(output_dir, f"{model_name}_ort_support.onnx") + save_onnx(onnx_model, onnx_path, use_external_data_format) + logging.info( + f"Model with ORT support is saved to {onnx_path}. Model contains custom ops: {custom_ops}." + ) + intermediate_generated_files.append(onnx_path) + elif platform.system() != "Windows": + logging.warning( + "No custom ops found. If that's not correct, please make sure that the 'tensorrt' python " + "is correctly installed and that the path to 'libcudnn*.so' is in PATH or LD_LIBRARY_PATH. If the " + "custom op is not directly available as a plugin in TensorRT, please also make sure that " + "the path to the compiled '.so' TensorRT plugin is also being given via the " + "'--trt_plugins' flag (requires TRT 10+)." + ) + + # Check if there's a custom TensorRT op in the ONNX model + has_custom_op = np.any([node.domain == "trt.plugins" for node in onnx_model.graph.node]) + + # Per-Channel support with QDQ format requires onnx opset version 13 or above + ai_onnx_domain = [ + opset + for opset in onnx_model.opset_import + if not opset.domain or opset.domain in ["ai.onnx", "ai.onnx.contrib"] + ] + opset_version = ai_onnx_domain[0].version + logging.info(f"Model {onnx_path} with opset_version {opset_version} is loaded.") + + required_opset_version = 13 + if opset_version < required_opset_version and opset_version != 1: + opset_version = required_opset_version + onnx_model = onnx.version_converter.convert_version(onnx_model, opset_version) + onnx_path = os.path.join(output_dir, f"{model_name}_opset{opset_version}.onnx") + save_onnx(onnx_model, onnx_path, use_external_data_format) + logging.info(f"Model is cloned to {onnx_path} with opset_version {opset_version}.") + intermediate_generated_files.append(onnx_path) + + # Sometimes input onnx model does not contain the node names + # This tool depends on those names, so we assign names if needed + graph = onnx_model.graph + is_named = name_onnx_nodes(graph) + onnx_model, is_duplicated_constant = duplicate_shared_constants(onnx_model) # FasterViT-0, eef + + if is_named or is_duplicated_constant: + onnx_path = os.path.join(output_dir, f"{model_name}_named.onnx") + save_onnx(onnx_model, onnx_path, use_external_data_format) + logging.info(f"Model is cloned to {onnx_path} after naming the nodes.") + intermediate_generated_files.append(onnx_path) + + # If custom op precisions are given, check if they're fp16. If so, add cast_to_fp16 before all inputs and + # cast_to_fp32 after all outputs. + if trt_plugins_precision: + custom_ops_to_cast = [] + for trt_plugin_precision in trt_plugins_precision: + assert ":" in trt_plugin_precision, ( + "Plugin precision is incorrectly formatted." + " Please check that it's in the format :." + ) + op_type, precision = trt_plugin_precision.split(":") + if precision == "fp16": + custom_ops_to_cast.append(op_type) + if custom_ops_to_cast: + onnx_path = add_fp16_fp32_cast(onnx_path, custom_ops_to_cast) + logging.info("Adding cast nodes related to custom ops to match requested precisions.") + intermediate_generated_files.append(onnx_path) + return onnx_path, intermediate_generated_files, has_custom_op + + +def quantize( + onnx_path: str, + quantize_mode: str = "int8", + calibration_data: CalibrationDataType = None, + calibration_method: str = None, + calibration_cache_path: str = None, + calibration_shapes: str = None, + calibration_eps: list[str] = ["cuda:0", "cpu", "trt"], + op_types_to_quantize: list[str] = None, + op_types_to_exclude: list[str] = None, + nodes_to_quantize: list[str] = None, + nodes_to_exclude: list[str] = None, + use_external_data_format: bool = False, + keep_intermediate_files: bool = False, + output_path: str = None, + verbose: bool = False, + trt_plugins: str = None, + trt_plugins_precision: list[str] = None, + high_precision_dtype: str = None, + mha_accumulation_dtype: str = "fp32", + disable_mha_qdq: bool = False, + dq_only: bool = True, + block_size: Optional[int] = None, + use_zero_point: bool = False, + **kwargs: Any, +) -> None: + """Quantizes the provided ONNX model. + + Args: + onnx_path: + Path to the input ONNX model. + quantize_mode: + Quantization mode. One of 'int8' (default), 'int4' and 'fp8'. + calibration_data: + Calibration data, either a numpy array or list/dict of numpy arrays. + calibration_method: + Calibration method choices. Options are int8: 'entropy' (default) and 'max', + fp8: 'max' (default) and int4: 'awq_clip' (default), 'awq_lite', 'awq_full' and 'rtn_dq'. + calibration_cache_path: + Path to pre-calculated activation tensor ranges, also known as calibration cache. + calibration_eps: + Priority order for the execution providers (EP) to calibrate the model. + Any subset of ['cuda:x', 'dml:x', 'cpu', 'trt'], where 'x' is the device id. + op_types_to_quantize: + List of op types to quantize. If None (default), all supported operators are quantized. + This flag does not support regular expression. + op_types_to_exclude: + List of op types to exclude from quantization. This flag does not support regular expression. + nodes_to_quantize: + List of node names to quantize. If None (default), all supported nodes are quantized. + This flag supports regular expression. + nodes_to_exclude: + List of node names to exclude from quantization. This flag supports regular expression. + use_external_data_format: + If True, separate data path will be used to store the weights of the quantized model. + keep_intermediate_files: + If True, keep all intermediate files generated during the ONNX model's conversion/calibration. + output_path: + Output filename to save the quantized ONNX model. + If None, save in the same directory as the original ONNX model with .quant suffix. + verbose: + If True, print details of node partition, selection etc. throughout the quantization process. + trt_plugins: + Specifies custom TensorRT plugin library paths in .so format (compiled shared library). + For multiple paths, separate them with a semicolon, i.e.: "lib_1.so;lib_2.so". + If this is not None or the model has custom ops, TensorrtExecutionProvider becomes the first choice as + calibration execution provider, meaning that the TensorRT is a requirement. + trt_plugins_precision: + A space-separated list indicating the precision for each custom op. + Each item should have the format :, where precision can be fp32 (default) or fp16. + For example: op_type_1:fp16 op_type_2:fp32. + high_precision_dtype: + High precision data type, one of ['fp32', 'fp16']. If high_precision_dtype == 'fp16', model's weight and + activation will be converted to fp16. + mha_accumulation_dtype: + MHA accumulation dtype. One of ['fp32', 'fp16']. 'fp32' by default. + If quantize_mode == 'fp8' and high_precision_dtype == 'fp32', Cast nodes will be added to + MHA's bmm1 and bmm2's input and output tensors. + disable_mha_qdq: + Don't add Q/DQ layers to MatMuls in MHA pattern. + dq_only: + If True (default), only add DQ nodes to the model. If False, add Q/DQ nodes to the model. + block_size: + Block size parameter for int4 quantization. + use_zero_point: + Use zero-point based quantization, if True. + kwargs: + Additional keyword arguments for int4 quantization, including: + - awqlite_alpha_step (float): Alpha step for lite, range [0, 1]. + - awqclip_alpha_step (float): Min alpha step for clip, range [awqclip_alpha_step, 1]. + - awqclip_alpha_min (float): Alpha step to find best alpha for clip. + - awqclip_bsz_col (int): Batch size for processing the column dimension in clip. + + Returns: + None, writes the quantized onnx model in the supplied output_path + or writes to the same directory with filename like ".quant.onnx". + """ + # quantize_static creates a shape-inferred copy at the input model's directory + # Needs to check if we have write permission to this directory + assert onnx_path.endswith(".onnx") or onnx_path.endswith(".pb") + if not os.access(os.path.dirname(os.path.abspath(onnx_path)), os.W_OK): + old_dir = os.path.dirname(os.path.abspath(onnx_path)) + tmp_dir = tempfile.mkdtemp() + logging.info(f"Directory {old_dir} is not writable, store intermediate files in {tmp_dir}") + onnx_path = os.path.join(tmp_dir, os.path.basename(onnx_path)) + + # We assume that the model directory contains only model related weights and protobuf file + # Anything extra in the model directory will be copied unnecessarily + for file in os.listdir(old_dir): + old_file_path = os.path.join(old_dir, file) + new_file_path = os.path.join(tmp_dir, file) + if os.path.isfile(old_file_path): + shutil.copy(old_file_path, new_file_path) + + model_name = os.path.splitext(os.path.basename(onnx_path))[0] + if not output_path: + output_dir = os.path.dirname(onnx_path) + output_path = os.path.join(output_dir, f"{model_name}.quant.onnx") + logging.info(f"No output path specified, save quantized model to {output_path}") + + # We need to preprocess the model with naming, weight duplication etc. + onnx_path, intermediate_generated_files, has_custom_op = _preprocess_onnx( + onnx_path, use_external_data_format, output_path, trt_plugins, trt_plugins_precision + ) + # If the model has a custom op and no plugin path was given, assume that this custom op is being implemented + # by a TRT native plugin. In order to enable the TRT EP, 'trt_extra_plugin_lib_paths' needs to be != None. + if has_custom_op and not trt_plugins: + trt_plugins = "" + + # Use random scales if calibration data is not supplied + if calibration_data is None: + calibration_data_reader = RandomDataProvider(onnx_path, calibration_shapes) + else: + calibration_data_reader = CalibrationDataProvider( + onnx_path, calibration_data, calibration_shapes + ) + + nodes_to_quantize = nodes_to_quantize or [] + nodes_to_exclude = nodes_to_exclude or [] + + # (1) If disable_mha_qdq is set, don't add Q/DQ layers to MatMuls in MHA pattern. + # (2) else when quantize_mode == "int8", if seq_len > 512, don't add Q/DQ layers to + # MatMuls in MHA pattern. + # (2) else when quantize_mode == "fp8", if head_size > 256 or head_size % 16 != 0 + # or mha doesn't meet fp8 fMHA v2 pattern, don't add Q/DQ layers to MatMuls in MHA pattern. + nodes_to_exclude = find_nodes_from_mha_to_exclude( + onnx_path, + use_external_data_format, + nodes_to_exclude, + disable_mha_qdq, + quantize_mode, + intermediate_generated_files, + calibration_data_reader, + calibration_eps, + verbose, + ) + + if quantize_mode in ["fp8", "int8"]: + quantize_func = quantize_int8 if quantize_mode == "int8" else quantize_fp8 + default_calibration_method = "entropy" if quantize_mode == "int8" else "max" + onnx_model = quantize_func( + onnx_path=onnx_path, + calibration_method=calibration_method or default_calibration_method, + calibration_data_reader=calibration_data_reader, + calibration_cache_path=calibration_cache_path, + calibration_shapes=calibration_shapes, + calibration_eps=calibration_eps, + op_types_to_quantize=op_types_to_quantize, + op_types_to_exclude=op_types_to_exclude, + nodes_to_quantize=nodes_to_quantize, + nodes_to_exclude=nodes_to_exclude, + use_external_data_format=use_external_data_format, + intermediate_generated_files=intermediate_generated_files, + verbose=verbose, + trt_extra_plugin_lib_paths=trt_plugins, + high_precision_dtype=high_precision_dtype, + mha_accumulation_dtype=mha_accumulation_dtype, + ) + elif "int4" in quantize_mode: + onnx_model = quantize_int4( + onnx_path=onnx_path, + calibration_method=calibration_method or "awq_clip", + calibration_data_reader=calibration_data_reader, + calibration_eps=calibration_eps, + use_external_data_format=use_external_data_format, + block_size=block_size, + nodes_to_exclude=nodes_to_exclude, + use_zero_point=use_zero_point, + **kwargs, + ) + else: + raise RuntimeError(f"Invalid quantization mode choice: {quantize_mode}") + + if onnx_model: + # Fuse Q nodes for INT8/FP8 mode + if quantize_mode in ["int8", "fp8"] and dq_only: + onnx_model = qdq_to_dq(onnx_model, verbose=verbose) + + # Collect and print stats of the quantized model + print_stat(gs.import_onnx(onnx_model), verbose) + + # Save the quantized model to the output path + save_onnx(onnx_model, output_path, use_external_data_format) + logging.info(f"Quantized onnx model is saved as {output_path}") + + # Check if intermediate files should be deleted + if not keep_intermediate_files: + for file in intermediate_generated_files: + if os.path.exists(file): + os.remove(file) + + # Check if the quantized model is valid + try: + onnx.checker.check_model(output_path) + except C.ValidationError as e: + logging.warning("ONNX model checker failed, check your deployment status.") + logging.warning(e) diff --git a/modelopt/onnx/quantization/src/modelopt_round_and_pack_ext.cpp b/modelopt/onnx/quantization/src/modelopt_round_and_pack_ext.cpp new file mode 100644 index 0000000000..d9f8273784 --- /dev/null +++ b/modelopt/onnx/quantization/src/modelopt_round_and_pack_ext.cpp @@ -0,0 +1,127 @@ +/* + * 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. + */ + +#include +#include +#include +#include + +#include +#include +#include + +namespace py = pybind11; + +constexpr int8_t INT4_MIN = -8; +constexpr int8_t INT4_MAX = 7; +constexpr int8_t UINT4_MAX = 15; +constexpr int8_t UINT4_MIN = 0; + +namespace { +int8_t ClipAndRoundOneNumber(bool is_signed, const double num) { + + const int8_t upper = is_signed ? INT4_MAX : UINT4_MAX; + const int8_t lower = is_signed ? INT4_MIN : UINT4_MIN; + + // clip the number to range [lower, upper] + const double min_num = num < upper ? num : upper; + const double max_num = min_num >= lower ? min_num : lower; + + // round the number to nearest integer + return static_cast(round(max_num)); +} +} // namespace + +// +// Arguments: +// is_signed: bool, whether to convert to signed int4 or uint4. +// inputs[]: array containing the input elements +// input_array_size: number of elements in input array +// outputs[]: int8 array for storing the packed output +// output_array_size: size of the output array, +// - should be large enough to accomodate all final entries +// - should be at least input_array_size/2 +// Return value: +// the number of elements written in output array +// +// Note: +// (1) When input array is of odd length +// - 0 is used as 2nd element for packing, and it goes to MSB. +// - "(input_array_size/2)+1" is the required output array length. +// + +template +size_t RoundAndPack(const bool is_signed, const py::array_t &inputs_arg, + const py::size_t input_array_size, py::array_t &outputs_arg, + const py::size_t output_array_size) { + + if (!input_array_size || !output_array_size) { + return 0; + } + + if (output_array_size < + ceil(input_array_size / 2.0f)) // should be "==" for most/practical purposes + { + return 0; + } + + const auto buf = inputs_arg.request(); + const T *inputs = static_cast(buf.ptr); + + auto buf2 = outputs_arg.request(); + int8_t *outputs = static_cast(buf2.ptr); + + if (!inputs || !outputs || (buf.ndim != 1) || (buf2.ndim != 1)) { + return 0; + } + + size_t i, j; + for (i = 0, j = 0; (i < (input_array_size - 1)) && (j < output_array_size); i += 2, j++) { + int8_t v1 = ClipAndRoundOneNumber(is_signed, static_cast(inputs[i])); + int8_t v2 = ClipAndRoundOneNumber(is_signed, static_cast(inputs[i + 1])); + outputs[j] = + ((v2 << 4) | (v1 & 0x0F)); // index i (v1) goes to LSB and index i+1 (v2) goes to MSB + } + + if (input_array_size % 2) { + assert(i == (input_array_size - 1)); + assert(j <= (output_array_size - 1)); + // 4-bit number goes to LSB, MSB remains 0 + outputs[j] = 0x0F & ClipAndRoundOneNumber(is_signed, inputs[i]); + j++; + i++; + } + + assert(i == input_array_size); + assert(j == ceil(input_array_size / 2.0f)); + + return j; +} + +PYBIND11_MODULE(modelopt_round_and_pack_ext, m) { + m.doc() = "C++ implementation of a basic round_and_pack logic"; + m.def("round_and_pack", &RoundAndPack, "Round-And-Pack C++ Implementation"); + m.def("round_and_pack", &RoundAndPack, "Round-And-Pack C++ Implementation"); + m.def("round_and_pack", &RoundAndPack, "Round-And-Pack C++ Implementation"); + m.def("round_and_pack", &RoundAndPack, "Round-And-Pack C++ Implementation"); +} + +/* +<% +setup_pybind11(cfg) +%> +*/ diff --git a/modelopt/onnx/quantization/trt_utils.py b/modelopt/onnx/quantization/trt_utils.py new file mode 100644 index 0000000000..4b97589bc0 --- /dev/null +++ b/modelopt/onnx/quantization/trt_utils.py @@ -0,0 +1,116 @@ +# 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. + +"""This module contains TensorRT utils.""" + +import ctypes +import platform +from typing import Optional + +import numpy as np +import onnx +import onnx_graphsurgeon as gs + +try: + import tensorrt as trt + + TRT_PYTHON_AVAILABLE = True +except ImportError: + TRT_PYTHON_AVAILABLE = False + + +def get_custom_layers(onnx_path: str, trt_plugins: Optional[str]) -> list[str]: + """Gets custom layers in ONNX model. + + Args: + onnx_path: Path to the input ONNX model. + trt_plugins: Paths to custom TensorRT plugins. + + Returns: + List of custom layers. + """ + # Initialize TensorRT plugins + if trt_plugins is not None: + trt_plugins = trt_plugins.split(";") + for plugin in trt_plugins: + ctypes.CDLL(plugin) + + # Create builder and network + trt_logger = trt.Logger(trt.Logger.WARNING) + trt.init_libnvinfer_plugins(trt_logger, "") + builder = trt.Builder(trt_logger) + network = builder.create_network() + + # Parse ONNX file + parser = trt.OnnxParser(network, trt_logger) + if not parser.parse_from_file(onnx_path): + error_str = [str(parser.get_error(error)) for error in range(parser.num_errors)] + raise Exception(f"Failed to parse ONNX file: {''.join(error_str)}") + + # Obtain plugin layer names + custom_layers = [] + for layer_idx in range(network.num_layers): + layer = network.get_layer(layer_idx) + if "PLUGIN" in str(layer.type): + custom_layers.append(layer.name) + + return custom_layers + + +def load_onnx_model( + onnx_path: str, trt_plugins: Optional[str] = None, use_external_data_format: bool = False +) -> tuple[onnx.onnx_pb.ModelProto, bool, list[str]]: + """Load ONNX model. If 'tensorrt' is installed, check if the model has custom ops and ensure it's supported by ORT. + + Args: + onnx_path: Path to the input ONNX model. + trt_plugins: Paths to custom TensorRT plugins. + use_external_data_format: If True, separate data path will be used to store the weights of the quantized model. + + Returns: + Loaded ONNX model supported by ORT. + Boolean indicating whether the model has custom ops or not. + List of custom ops in the ONNX model. + """ + custom_ops = [] + has_custom_op = False + + # Load the model and weights + onnx_model = onnx.load(onnx_path, load_external_data=use_external_data_format) + + if TRT_PYTHON_AVAILABLE and platform.system() != "Windows": + # Check if there's a custom TensorRT op in the ONNX model. If so, make it ORT compatible by adding + # `trt.plugins to the ONNX graph. + trt_plugin_domain = "trt.plugins" + trt_plugin_version = 1 + + custom_layers = get_custom_layers(onnx_path, trt_plugins) + has_custom_op = True if custom_layers else False + + if has_custom_op: + graph = gs.import_onnx(onnx_model) + for node in graph.nodes: + if node.name in custom_layers: + custom_ops.append(node.op) + node.domain = trt_plugin_domain + custom_ops = np.unique(custom_ops) + + # TODO: Add type and shape inference. + onnx_model = gs.export_onnx(graph) + onnx_model.opset_import.append( + onnx.helper.make_opsetid(trt_plugin_domain, trt_plugin_version) + ) + + return onnx_model, has_custom_op, custom_ops diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py new file mode 100644 index 0000000000..915e2152e6 --- /dev/null +++ b/modelopt/onnx/utils.py @@ -0,0 +1,629 @@ +# 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. + +"""Utility functions related to onnx.""" + +import io +import os +from collections import defaultdict +from typing import Any, Optional, Union + +import numpy as np +import onnx +import onnx.onnx_cpp2py_export.checker as C # noqa: N812 +import onnx_graphsurgeon as gs +from onnx import numpy_helper +from onnx_graphsurgeon import Constant, Node, Variable + + +def get_input_names_from_bytes(model_bytes: bytes, external_inputs_only: bool = True) -> list[str]: + """This function returns the inputs names of the given onnx model in bytes. + + Args: + model_bytes: Onnx model in bytes. + + Returns: + List of input names of the model. + """ + model = onnx.load_from_string(model_bytes) + return get_input_names(model, external_inputs_only) + + +def get_all_input_names(model: onnx.onnx_ml_pb2.ModelProto) -> list[str]: + """This function returns the inputs names of the given onnx model.""" + return [graph_input.name for graph_input in model.graph.input] + + +def _get_initializer_names(model: onnx.onnx_ml_pb2.ModelProto) -> list[str]: + return [initializer.name for initializer in model.graph.initializer] + + +def get_input_names( + model: onnx.onnx_ml_pb2.ModelProto, external_inputs_only: bool = True +) -> list[str]: + """This function returns the external inputs names of the given onnx model. + + Note: external_input_names = input_names - initializer_names + + Args: + model: Loaded in-memory onnx ModelProto. + + Returns: + List of external input names of the model. + """ + input_names = get_all_input_names(model) + if not external_inputs_only: + return input_names + + initializer_names = _get_initializer_names(model) + external_input_names = list(np.setdiff1d(input_names, initializer_names)) + return external_input_names + + +def get_output_names_from_bytes(model_bytes: bytes) -> list[str]: + """This function returns the output names of the given onnx model in bytes. + + Args: + model_bytes: Onnx model in bytes. + + Returns: + List of output names of the model. + """ + model = onnx.load_from_string(model_bytes) + return get_output_names(model) + + +def get_output_names(model: onnx.onnx_ml_pb2.ModelProto) -> list[str]: + """This function returns the output names of the given onnx model. + + Args: + model: Loaded in-memory onnx ModelProto. + + Returns: + List of output names of the model. + """ + return [output.name for output in model.graph.output] + + +def get_node_names_from_bytes(model_bytes: bytes) -> list[str]: + """This function returns all node names from the given onnx model in bytes. + + Args: + model: onnx model in bytes. + + Returns: + List of node names of the model. + """ + model = onnx.load_from_string(model_bytes) + return get_node_names(model) + + +def get_node_names(model: onnx.onnx_ml_pb2.ModelProto) -> list[str]: + """This function returns all node names from the given onnx model. + + Args: + model: Loaded in-memory onnx ModelProto. + + Returns: + List of node names of the model. + """ + return [node.name for node in model.graph.node] + + +def _get_tensor_shape(tensor: onnx.onnx_ml_pb2.ValueInfoProto) -> list[int]: + """This function returns the shape of the input onnx tensor. + + Onnx tensors are of type ValueInfoProto and their dimensions are stored in a + RepeatedCompositeContainer. Each of these dimensions is of type onnx.onnx_ml_pb2.Dimension. + In a loop we access each of the Dimension object and create a shape list to return. + Dynamic dimensions (i.e. with "dim_param" field) are replaced with 1. + + Args: + tensor: Onnx tensor object for which the shape needs to be computed. + + Returns: + Shape of the input tensor. + """ + if not hasattr(tensor.type, "tensor_type"): + raise NotImplementedError("Only tensor type inputs are supported.") + + dimensions = tensor.type.tensor_type.shape.dim + shape = [] + for dim in dimensions: + if dim.HasField("dim_param"): + shape.append(1) + if dim.HasField("dim_value"): + if dim.dim_value == -1: + shape.append(1) + else: + shape.append(dim.dim_value) + + return shape + + +def _get_all_shapes(container: Any) -> dict[str, list[int]]: + """This method returns the shape of tensors within a RepeatedCompositeContainer. + + Args: + container: Model graph input/output container. + + Returns: + Dictionary of tensor names and shape of the tensors within the container. + """ + results = {} + for tensor in container: + results[tensor.name] = _get_tensor_shape(tensor) + return results + + +def _get_selected_shapes(container: Any, inputs_to_include: list[str]) -> dict[str, list[int]]: + """This method returns the shape tensors within a RepeatedCompositeContainer. + + It only computes the shape of the tensors with name containing in `inputs_to_include` list. + + Args: + container: Model graph input/output container. + + Returns: + Dictionary of tensor names in inputs_to_include and their shapes. + """ + results = {} + for tensor in container: + if tensor.name in inputs_to_include: + results[tensor.name] = _get_tensor_shape(tensor) + return results + + +def get_input_shapes_from_bytes(model_bytes: bytes) -> dict[str, list[int]]: + """This function returns the input shapes of the given onnx model in bytes. + + Args: + model_bytes: Onnx model in bytes. + + Returns: + Dictionary of inputs names and shapes. + """ + model = onnx.load_from_string(model_bytes) + return get_input_shapes(model) + + +def get_input_shapes( + model: onnx.onnx_ml_pb2.ModelProto, external_inputs_only: bool = True +) -> dict[str, list[int]]: + """This function returns the inputs shapes for the given onnx model.""" + if external_inputs_only: + return _get_selected_shapes(model.graph.input, get_input_names(model)) + return _get_all_shapes(model.graph.input) + + +def get_output_shapes(model: onnx.onnx_ml_pb2.ModelProto) -> dict[str, list[int]]: + """This function returns the output shapes for the given onnx model.""" + return _get_all_shapes(model.graph.output) + + +def parse_shapes_spec(shapes_spec: str) -> dict[str, list[int]]: + """Parse shapes spec and returns them in a dictionary. + + Example shapes spec: input0:1x3x256x256,input1:1x3x128x128 + """ + shapes = shapes_spec.split(",") + shape_dict = {} + + for shape_def in shapes: + tensor_name, shape_str = shape_def.rsplit(":", 1) + shape_list = list(map(int, shape_str.split("x"))) + shape_dict[tensor_name.strip()] = shape_list + + return shape_dict + + +def _get_tensor_type(tensor: onnx.onnx_ml_pb2.ValueInfoProto) -> int: + if not hasattr(tensor.type, "tensor_type"): + raise NotImplementedError("Only tensor type inputs are supported.") + type_ = tensor.type.tensor_type.elem_type + return type_ + + +def _get_container_types( + container, inputs_to_include: Union[None, list[str]] = None +) -> dict[str, int]: + results = {} + for tensor in container: + if inputs_to_include is not None: + if tensor.name not in inputs_to_include: + continue + t_type = _get_tensor_type(tensor) + results[tensor.name] = t_type + return results + + +def _get_input_types( + model: onnx.onnx_ml_pb2.ModelProto, external_inputs_only: bool = True +) -> dict[str, int]: + inputs_to_include = get_input_names(model, external_inputs_only) + return _get_container_types(model.graph.input, inputs_to_include) + + +def _get_output_types(model: onnx.onnx_ml_pb2.ModelProto) -> dict[str, int]: + results = _get_container_types(model.graph.output) + return results + + +def _convert_types_to_np(types: Union[dict[str, int], list[int], int]) -> Any: + if isinstance(types, dict): + types_np = {} + for name in types.keys(): + types_np[name] = onnx.helper.tensor_dtype_to_np_dtype(types[name]) + return types_np + elif isinstance(types, list): + return [onnx.helper.tensor_dtype_to_np_dtype(type_) for type_ in types] + else: + return onnx.helper.tensor_dtype_to_np_dtype(types) + + +def gen_random_inputs( + model: onnx.onnx_ml_pb2.ModelProto, shapes_spec: str = None +) -> dict[str, np.ndarray]: + """This function generates random inputs for an onnx model. + + Args: + model: Loaded in-memory onnx ModelProto. + shapes_spec: A string representing the shape of each input tensors. The format is + ":x,:,...". If the shape is not provided for an input tensor, the shape is + inferred from the onnx model directly, with all the unknown dims filled with 1. + + Returns: + Dictionary of numpy tensors. + """ + input_dict = {} + types = _get_input_types(model) + types_np = _convert_types_to_np(types) + input_shapes = {} if shapes_spec is None else parse_shapes_spec(shapes_spec) + + for graph_input in model.graph.input: + # Generate tensors for external inputs only + if graph_input.name not in types_np: + continue + + shape_arr = ( + input_shapes[graph_input.name] + if graph_input.name in input_shapes + else _get_tensor_shape(graph_input) + ) + + input_dict[graph_input.name] = np.random.uniform(size=shape_arr).astype( + types_np[graph_input.name] + ) + + return input_dict + + +def remove_weights_data(onnx_bytes: bytes) -> bytes: + """Removes raw weight data from the onnx model.""" + model = onnx.load_from_string(onnx_bytes) + inits = model.graph.initializer + + for idx, init in enumerate(inits): + # Only remove arrays with dimension larger than 1 + if len(init.dims) > 1: + dtype = onnx.helper.tensor_dtype_to_np_dtype(init.data_type) + if dtype in ["float16", "float32", "float64"]: + # Setting up some metadata to randomize the weights later + np_tensor = np.frombuffer(init.raw_data, dtype=dtype) + meta = model.metadata_props.add() + meta.key = init.name + "_avg" + meta.value = str(np.average(np_tensor)) + + meta = model.metadata_props.add() + meta.key = init.name + "_var" + meta.value = str(np.var(np_tensor)) + + # Note that, onnx.checker will fail due to data cleaning + # We should not check the model till weights are reassigned + model.graph.initializer[idx].raw_data = b"" + + buffer = io.BytesIO() + onnx.save_model(model, buffer) + buffer.seek(0, 0) + + return buffer.read() + + +def randomize_weights(onnx_path: str) -> None: + """Assigns random values to the onnx model weights.""" + with open(onnx_path, "rb") as f: + onnx_bytes = f.read() + onnx_bytes = randomize_weights_onnx_bytes(onnx_bytes) + + with open(onnx_path, "wb") as f: + # Write the modified onnx model to the same path + f.write(onnx_bytes) + + +def randomize_weights_onnx_bytes(onnx_bytes: bytes, seed: int = 0) -> bytes: + """Assigns random values to the onnx model weights.""" + model = onnx.load_from_string(onnx_bytes) + inits = model.graph.initializer + np.random.seed(seed) + weight_metadata = {item.key: item.value for item in model.metadata_props} + + for idx, init in enumerate(inits): + # Randomize only the arrays with dimension larger than 1 + if len(init.dims) > 1: + dtype = onnx.helper.tensor_dtype_to_np_dtype(init.data_type) + if dtype in ["float16", "float32", "float64"]: + avg = weight_metadata.get(init.name + "_avg", None) + var = weight_metadata.get(init.name + "_var", None) + if avg and var: + numpy_array = np.random.normal(float(avg), float(var), size=init.dims).astype( + dtype + ) + tensor = numpy_helper.from_array(numpy_array, init.name) + model.graph.initializer[idx].CopyFrom(tensor) + + buffer = io.BytesIO() + onnx.save_model(model, buffer) + buffer.seek(0, 0) + + return buffer.read() + + +def validate_onnx(onnx_bytes: bytes) -> bool: + """Returns True if the onnx_bytes is valid, else False.""" + if not onnx_bytes: + return False + + try: + onnx_model = onnx.load_from_string(onnx_bytes) + return onnx_model is not None + except Exception: + return False + + +def validate_batch_size(onnx_bytes: bytes, batch_size: int) -> bool: + """Returns True if all the model inputs has batch dimension equal to batch_size.""" + input_shapes = list(get_input_shapes_from_bytes(onnx_bytes).values()) + for shape in input_shapes: + if shape[0] != batch_size: + return False + + return True + + +def get_batch_size(model: onnx.onnx_ml_pb2.ModelProto) -> int: + """Returns the batch size of the given onnx model. + + Assertion will fail if batch size is not same over all the inputs. + """ + input_shapes = list(get_input_shapes(model).values()) + batch_size = input_shapes[0][0] + for shape in input_shapes: + if batch_size != shape[0]: + # The model does not have the batch dimension + return 1 + + return batch_size + + +def get_batch_size_from_bytes(onnx_bytes: bytes) -> int: + """Returns the batch size of the given onnx model. + + Assertion will fail if batch size is not same over all the inputs. + """ + model = onnx.load_from_string(onnx_bytes) + return get_batch_size(model) + + +def save_onnx_bytes_to_dir(onnx_bytes: bytes, onnx_dir: str, onnx_name: str) -> None: + """Saves the onnx bytes to a directory with specified file name.""" + os.makedirs(onnx_dir, exist_ok=True) + file_path = os.path.join(onnx_dir, onnx_name + ".onnx") + + try: + with open(file_path, "wb") as f: + f.write(onnx_bytes) + print(f"Onnx model saved as {file_path}") + except Exception as e: + print(f"Onnx model exporting as {file_path} failed, error {str(e)}") + + +def name_onnx_nodes(graph: onnx.onnx_ml_pb2.GraphProto) -> bool: + """Assigns name to the onnx nodes if not present and return the modified status.""" + is_modified = False + node_names = set([node.name for node in graph.node]) + start_id = len(node_names) + for node in graph.node: + if not node.name: + new_name = f"{node.op_type}_{start_id}" + while new_name in node_names: + start_id += 1 + new_name = f"{node.op_type}_{start_id}" + + node.name = new_name + node_names.add(new_name) + is_modified = True + + return is_modified + + +def duplicate_shared_constants(onnx_model: onnx.ModelProto) -> tuple[onnx.ModelProto, bool]: + """Duplicate constant tensors if they are shared.""" + graph = gs.import_onnx(onnx_model) + name_dict = defaultdict(lambda: 0) + + def _get_unique_name(old_name): + name_dict[old_name] += 1 + return old_name + "_" + str(name_dict[old_name]) + + # Get tensors with shared constant inputs + tensors = [] + for node in graph.nodes: + for inp_idx, tensor in enumerate(node.inputs): + if isinstance(tensor, Constant): + if len(tensor.outputs) > 1: # constant is shared across multiple nodes + tensors.append({"tensor": tensor, "inp_node": node, "inp_idx": inp_idx}) + + # Duplicate shared tensors + for tensor_dict in tensors: + tensor = tensor_dict["tensor"] + new_tensor = Constant( + name=_get_unique_name(tensor.name), + values=tensor.values, + ) + tensor_dict["inp_node"].inputs[tensor_dict["inp_idx"]] = new_tensor + + onnx_model = gs.export_onnx(graph) + onnx_model.ir_version = 9 + is_modified = True if tensors else False + return onnx_model, is_modified + + +def is_valid_onnx_model(file_path): + """Checks if the given file is a valid ONNX model.""" + if not os.path.exists(file_path): + print(f"No file found at {file_path}") + return False + + try: + # Load the ONNX model + model = onnx.load(file_path) + + # Check the model + onnx.checker.check_model(model) + print(f"ONNX model at {file_path} is valid.") + return True + except C.ValidationError as e: + print(f"The file is not a valid ONNX model. {e}") + except Exception as e: + print(f"An error occurred: {e}") + + return False + + +def find_lowest_common_ancestor(node1: Node, node2: Node) -> tuple[Optional[str], int, int]: + """Function to find the lowest common ancestor of two nodes. + + Args: + node1: First node name. + node2: Second node name. + + Returns: + LCA node. + Distance from first node. + Distance from second node. + """ + + def _find_ancestors(node: Node): + ancestors = {node.name: 0} + stack = [(node, 0)] + while stack: + cur_node, distance = stack.pop() + for parent_node in get_parent_nodes(cur_node): + if parent_node.name not in ancestors: + ancestors[parent_node.name] = distance + 1 + stack.append((parent_node, distance + 1)) + + return ancestors + + ancestors1 = _find_ancestors(node1) + ancestors2 = _find_ancestors(node2) + + # Find the lowest common ancestor + common_ancestors = set(ancestors1.keys()).intersection(ancestors2.keys()) + if common_ancestors: + lowest_common_ancestor = common_ancestors.pop() + distance = ancestors1[lowest_common_ancestor] + for t in common_ancestors: + if ancestors1[t] < distance: + distance = ancestors1[t] + lowest_common_ancestor = t + distance1 = ancestors1[lowest_common_ancestor] + distance2 = ancestors2[lowest_common_ancestor] + return lowest_common_ancestor, distance1, distance2 + else: + return None, -1, -1 # No common ancestor found + + +def get_parent_nodes(node: Node) -> list[Node]: + """Returns list of input producer nodes for the given node.""" + parents = [] + for tensor in node.inputs: + # If the tensor is not a constant or graph input and has a producer, + # the producer is a parent of node `node` + if len(tensor.inputs) == 1: + parents.append(tensor.inputs[0]) + + return parents + + +def get_child_nodes(node: Node) -> list[Node]: + """Returns list of output consumer nodes for the given node.""" + children = [] + for tensor in node.outputs: + for consumer in tensor.outputs: # Traverse all consumer of the tensor + children.append(consumer) + + return children + + +def get_variable_inputs(node: Node) -> list[Variable]: + """Returns the variable inputs of the given Node.""" + var_inputs = [] + for tensor in node.inputs: + if isinstance(tensor, Variable): + if not tensor.inputs or (tensor.inputs and tensor.inputs[0].op != "Constant"): + var_inputs.append(tensor) + return var_inputs + + +def save_onnx( + model: onnx.onnx_ml_pb2.ModelProto, onnx_path: str, save_as_external_data: bool = False +): + """Save an ONNX model to given path. If a model is larger than 2GB, will save with external data.""" + size_threshold = 2 * 1024 * 1024 * 1024 + try: + model_proto = model.SerializeToString() + model_size = len(model_proto) + save_as_external_data = save_as_external_data or model_size > size_threshold + + except ValueError as e: + if "Message onnx.ModelProto exceeds maximum protobuf size of 2GB" in str(e): + save_as_external_data = True + else: + raise + + if save_as_external_data: + external_data_path = os.path.basename(onnx_path) + "_data" + onnx.save_model( + model, + onnx_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=external_data_path, + size_threshold=1024, + ) + else: + onnx.save(model, onnx_path) + + +def udpate_domain( + onnx_model: onnx.onnx_ml_pb2.ModelProto, op_type: str, domain: str +) -> onnx.onnx_ml_pb2.ModelProto: + """Updates the domain of all the nodes of the specified op_type to the specified domain.""" + for node in onnx_model.graph.node: + if node.op_type == op_type: + node.domain = domain + + return onnx_model diff --git a/modelopt/torch/__init__.py b/modelopt/torch/__init__.py new file mode 100644 index 0000000000..ccb3ec6528 --- /dev/null +++ b/modelopt/torch/__init__.py @@ -0,0 +1,31 @@ +# 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. + +"""Model optimization and deployment subpackage for torch.""" + +import warnings as _warnings + +from packaging.version import Version as _Version +from torch import __version__ as _torch_version + +try: + from . import distill, nas, opt, prune, quantization, sparsity, speculative, utils # noqa: E402 +except ImportError as e: + raise ImportError(f"{e}\nPlease install optional ``[torch]`` dependencies.") + +if _Version(_torch_version) < _Version("2.2"): + _warnings.warn( + "nvidia-modelopt will drop torch 2.1 support in a future release.", DeprecationWarning + ) diff --git a/modelopt/torch/_deploy/__init__.py b/modelopt/torch/_deploy/__init__.py new file mode 100644 index 0000000000..deec6c177b --- /dev/null +++ b/modelopt/torch/_deploy/__init__.py @@ -0,0 +1,23 @@ +# 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. + +"""Interfaces to benchmark models and convert models into binaries for inference devices.""" + +try: + from . import device_model, utils + from .compilation import * + from .profiling import * +except ImportError as e: + raise ImportError(f"{e}\nPlease install optional ``[onnx,torch]`` dependencies.") diff --git a/modelopt/torch/_deploy/_runtime/__init__.py b/modelopt/torch/_deploy/_runtime/__init__.py new file mode 100644 index 0000000000..b51da68f4f --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/__init__.py @@ -0,0 +1,30 @@ +# 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. + +from .registry import * +from .runtime_client import * + +# no runtime_client_impl will be available if 'deploy' is not installed +try: + from .ort_client import * +except ImportError: + pass + +try: + from .trt_client import * +except ImportError: + # ImportError if tensorrt is not installed + # ModuleNotFoundError if .tensorrt/ is not available + pass diff --git a/modelopt/torch/_deploy/_runtime/common.py b/modelopt/torch/_deploy/_runtime/common.py new file mode 100644 index 0000000000..355de28568 --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/common.py @@ -0,0 +1,73 @@ +# 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 logging +import time +from typing import Callable + + +def timeit(method: Callable) -> Callable: + """This function is supposed to use as a decorator to measure the execution time of another function. + + If the decorator is applied and no changes are done at the call site, this will print out the + timing information on the log console. If the call site wants to get the time info returned, they + should pass a dictionary named log_time like below- + + (regular_returns, ...), func_exec_time = func(regular_params, ..., log_time={}) + """ + + def timed(*args, **kw): + ts = time.time() + result = method(*args, **kw) + te = time.time() + if "log_time" in kw: + name = kw.get("log_name", method.__name__.upper()) + kw["log_time"][name] = (te - ts) * 1000 + return result, kw["log_time"][name] + else: + logging.info(f"{method.__name__}: {(te - ts) * 1000} ms") + return result + + return timed + + +def init_logging() -> None: + logging.basicConfig( + format="%(asctime)s P%(process)d T%(thread)d %(levelname)-8s %(message)s", + level=logging.INFO, + datefmt="%Y-%m-%d %H:%M:%S", + ) + + +def read_bytes(file_path: str) -> bytes: + with open(file_path, "rb") as file: + file_bytes = file.read() + return file_bytes + + +def read_string(file_path: str) -> str: + with open(file_path, "r") as file: + data = file.read() + return data + + +def write_bytes(data: bytes, file_path: str) -> None: + with open(file_path, "wb") as file: + file.write(data) + + +def write_string(data: str, file_path: str) -> None: + with open(file_path, "w") as file: + file.write(data) diff --git a/modelopt/torch/_deploy/_runtime/ort_client.py b/modelopt/torch/_deploy/_runtime/ort_client.py new file mode 100644 index 0000000000..a9b704559a --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/ort_client.py @@ -0,0 +1,190 @@ +# 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 +import os +import tempfile +import time +from typing import Any, Optional, Union + +import numpy as np +import onnxruntime as ort +import torch + +from .registry import RuntimeRegistry +from .runtime_client import Deployment, DeploymentTable, DetailedResults, RuntimeClient + +__all__ = ["ORTLocalClient"] + + +@RuntimeRegistry.register("ORT") +class ORTLocalClient(RuntimeClient): + """A client for using the local onnx runtime with CPU backend.""" + + @property + def _profile_defaults(self) -> dict[str, Union[int, float]]: + """Default profiling parameters.""" + return { + "iterations": 10, + "iterations_max": 1000, + "warm_up": 0.2, + "duration": 3.0, + } + + @property + def _accelerator_to_provider(self) -> dict[str, str]: + """Maps accelerator to ORT execution provider.""" + return { + "CPU": "CPUExecutionProvider", + } + + @property + def default_deployment(self) -> Deployment: + return {k: v[0] for k, v in self.deployment_table.items()} + + @property + def deployment_table(self) -> DeploymentTable: + return { + "version": ["1.16"], + "accelerator": list(self._accelerator_to_provider.keys()), + "precision": ["fp32"], + "onnx_opset": ["13"], + } + + def _ir_to_compiled(self, ir_bytes: bytes, compilation_args: dict[str, Any] = None) -> bytes: + """Converts an ONNX model to a compiled device model.""" + return ir_bytes # ir_bytes (onnx) are also compiled model for ORT + + def _onnx_to_np_dtype(self, onnx_type: str) -> np.dtype: + """Maps an ONNX data type to a numpy data type.""" + return { + "tensor(float16)": np.float16, + "tensor(float)": np.float32, + "tensor(double)": np.float64, + "tensor(int8)": np.int8, + "tensor(int16)": np.int16, + "tensor(int32)": np.int32, + "tensor(int64)": np.int64, + "tensor(uint8)": np.uint8, + "tensor(uint16)": np.uint16, + "tensor(uint32)": np.uint32, + "tensor(uint64)": np.uint64, + }[onnx_type] + + def _init_session( + self, compiled_model: bytes, session_options: Optional[ort.SessionOptions] = None + ) -> ort.InferenceSession: + """Initializes an inference session with the compiled model and returns the session.""" + provider = self._accelerator_to_provider[self.deployment["accelerator"]] + return ort.InferenceSession(compiled_model, session_options, providers=[provider]) + + def _profile( + self, compiled_model: bytes, compilation_args: dict[str, Any] = {} + ) -> tuple[float, DetailedResults]: + """Profiles a compiled device model and returns the latency & detailed profiling results.""" + # use a temp folder for profiling results + with tempfile.TemporaryDirectory() as temp_dir: + # initialize session + options + session_options = ort.SessionOptions() + session_options.enable_profiling = True + session_options.profile_file_prefix = os.path.join(temp_dir, "ort_profile") + ort_session = self._init_session(compiled_model, session_options) + + # generate dummy inputs from ort_session.get_inputs() + inputs = [ + np.asarray(np.random.rand(*x.shape)).astype(self._onnx_to_np_dtype(x.type)) + for x in ort_session.get_inputs() + ] + + # run session with dummy inputs + self._run_session(ort_session, inputs, **self._profile_defaults) # type: ignore[arg-type] + + # end profiling and load results + prof_file = ort_session.end_profiling() + with open(prof_file, "r") as p_file: + results = json.load(p_file) + + # get latency from profiling results (latencies are in nano-seconds) + # We generally use milliseconds for latency, so divide by 1e3 + latencies = [x["dur"] for x in results if x.get("name") == "model_run"] + avg_latency = np.mean(latencies) / 1e3 + + # return latency & detailed profiling results + return avg_latency, {"ort_results": results} + + def _inference( + self, compiled_model: bytes, inputs: list[torch.Tensor], io_shapes: dict[str, list] = {} + ) -> list[torch.Tensor]: + """Run inference with the compiled model and return the output as list of numpy arrays.""" + # initialize session, run session, and return session outputs + assert io_shapes == {}, "ORT does not support specifying IO shapes." + ort_session = self._init_session(compiled_model) + inputs = [input.numpy() for input in inputs] + return self._run_session(ort_session, inputs) + + def _run_session( + self, + ort_session: ort.InferenceSession, + inputs: list[np.ndarray], + iterations: int = 1, + iterations_max: int = 1000, + warm_up: float = 0.0, + duration: float = 0.0, + ) -> list[torch.Tensor]: + """Run inference with the compiled model and return the output as list of numpy arrays. + + Args: + ort_session: The ONNX runtime inference session. + inputs: The input tensors to the session. + iterations: The minimum number of iterations to run the session for. + iterations_max: The maximum number of iterations to run the session for. + warm_up: The number of seconds to warm up the session for. + duration: The minimum number of seconds to run the session for. + + Returns: + The outputs of the session as a list of numpy arrays. + """ + assert iterations > 0, "Number of iterations must be positive!" + + # basics of the session + ort_inputs = {ort_session.get_inputs()[i].name: inp for i, inp in enumerate(inputs)} + time_elapsed = 0.0 + + # small utility for running the session once with time measurement + def _run_session_once(): + nonlocal time_elapsed + time_elapsed -= time.perf_counter() + ort_outputs = ort_session.run(None, ort_inputs) + time_elapsed += time.perf_counter() + return ort_outputs + + # warm-up + while time_elapsed < warm_up: + _run_session_once() + + # run inference + time_elapsed = 0.0 + iter = 0 + total_time = -time.time() + while (iter < iterations or time_elapsed < duration) and iter < iterations_max: + ort_outputs = _run_session_once() + iter += 1 + total_time += time.time() + print(f"ORT: {iter} iterations in {total_time:.3f} seconds ({total_time/iter:.3f} s/iter)") + + ort_outputs = [torch.tensor(out) for out in ort_outputs] + + # return session outputs of final iteration + return ort_outputs diff --git a/modelopt/torch/_deploy/_runtime/registry.py b/modelopt/torch/_deploy/_runtime/registry.py new file mode 100644 index 0000000000..815c8e889b --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/registry.py @@ -0,0 +1,83 @@ +# 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. + +from typing import Callable + +from .runtime_client import Deployment, RuntimeClient + +__all__ = ["RuntimeRegistry"] + + +class RuntimeRegistry: + """Registry to store and retrieve various runtime client implementations.""" + + _runtime_client_lookup: dict[str, type[RuntimeClient]] = {} + + @classmethod + def register(cls, runtime: str) -> Callable[[type[RuntimeClient]], type[RuntimeClient]]: + """A decorator to register a RuntimeClient with its relevant runtime. + + For example: + + .. code-block:: python + + @RuntimeRegistry.register("my_runtime") + class MyRuntimeClient(RuntimeClient): + pass + """ + + def _register_runtime_client(new_type: type[RuntimeClient]) -> type[RuntimeClient]: + cls._runtime_client_lookup[runtime] = new_type + new_type._runtime = runtime + return new_type + + return _register_runtime_client + + @classmethod + def unregister(cls, runtime: str) -> None: + """A helper to unregister a RuntimeClient + + For example: + + .. code-block:: python + + @RuntimeRegistry.register("my_runtime") + class MyRuntimeClient(RuntimeClient): + pass + + + # later + RuntimeRegistry.unregister("my_runtime") + """ + cls._runtime_client_lookup.pop(runtime) + + @classmethod + def get(cls, deployment: Deployment) -> RuntimeClient: + """Get the runtime client for the given deployment. + + Args: + deployment: Deployment configuration. + + Returns: + The runtime client for the given deployment. + """ + # check for valid runtime + if "runtime" not in deployment: + raise KeyError("Runtime was not set.") + if deployment["runtime"] not in cls._runtime_client_lookup: + raise ValueError(f"Runtime {deployment['runtime']} is not supported.") + + # initialize runtime client + return cls._runtime_client_lookup[deployment["runtime"]](deployment) diff --git a/modelopt/torch/_deploy/_runtime/runtime_client.py b/modelopt/torch/_deploy/_runtime/runtime_client.py new file mode 100644 index 0000000000..f9eb9ecc34 --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/runtime_client.py @@ -0,0 +1,158 @@ +# 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. + +from abc import ABC, abstractmethod +from typing import Any + +import torch + +__all__ = ["Deployment", "DeploymentTable", "DetailedResults", "RuntimeClient"] + +Deployment = dict[str, str] +DeploymentTable = dict[str, list[str]] +DetailedResults = dict[str, Any] + + +class RuntimeClient(ABC): + """A abstract client class for implementing various runtimes to be used for deployment. + + The RuntimeClient defines a common interfaces for accessing various runtimes within modelopt. + """ + + _runtime: str # runtime of the client --> set by RuntimeRegistry.register + + def __init__(self, deployment: Deployment): + super().__init__() + self.deployment = self.sanitize_deployment_config(deployment) + + @property + def runtime(self) -> str: + return self._runtime + + def sanitize_deployment_config(self, deployment: Deployment) -> Deployment: + """Cleans/checks the deployment config & fills in runtime-specific default values. + + Args: + deployment: Deployment config with at least the ``runtime`` key specified. + + Returns: + The sanitized deployment config with all runtime-specific default values filled + in for missing keys. + """ + # check runtime + assert self.runtime == deployment["runtime"], "Runtime mismatch!" + + # check that version was provided + if "version" not in deployment: + raise KeyError("Runtime version must be provided!") + + # fill in default values and update + deployment = {**self.default_deployment, **deployment} + + # sanity check on keys (inverse doesn't have to be checked since we fill in defaults) + table = self.deployment_table + extra_keys = deployment.keys() - table.keys() - {"runtime"} + assert not extra_keys, f"Invalid deployment config keys detected: {extra_keys}!" + + # sanity checks on values + invalid_values = {(k, deployment[k]): t for k, t in table.items() if deployment[k] not in t} + assert not invalid_values, f"Invalid deployment config values detected: {invalid_values}!" + + return deployment + + def ir_to_compiled(self, ir_bytes: bytes, compilation_args: dict[str, Any] = {}) -> bytes: + """Converts a model from its intermediate representation (IR) to a compiled device model. + + Args: + ir_bytes: Intermediate representation (IR) of the model. + compilation_args: Additional arguments for the compilation process. + + Returns: The compiled device model that can be used for further downstream tasks such as + on-device inference and profiling, + """ + # run model compilation + compiled_model = self._ir_to_compiled(ir_bytes, compilation_args) + assert compiled_model, "Device conversion failed!" + + return compiled_model + + def profile( + self, compiled_model: bytes, compilation_args: dict[str, Any] = {} + ) -> tuple[float, DetailedResults]: + """Profiles a compiled device model and returns the latency & detailed profiling results. + + Args: + compiled_model: Compiled device model from compilation service. + compilation_args: Additional arguments for the compilation process. + + Returns: A tuple (latency, detailed_result) where + ``latency`` is the latency of the compiled model in ms, + ``detailed_result`` is a dictionary containing additional benchmarking results + """ + # get latency & detailed results from client + latency, detailed_result = self._profile(compiled_model, compilation_args) + assert latency > 0.0, "Profiling failed!" + + return latency, detailed_result + + def inference( + self, compiled_model: bytes, inputs: list[torch.Tensor], io_shapes: dict[str, list] = {} + ) -> list[torch.Tensor]: + """Run inference with the compiled model and return the output as list of torch Tensors. + + Args: + compiled_model: Compiled device model from compilation service. + inputs: Inputs to do inference on server. + io_shapes: Defines the input and output shapes of the model. These shapes must be provided by the user + when TensorRT cannot automatically infer them, such as in cases where dynamic dimensions are used for + model outputs. To specify the shapes, use the following format: + io_shapes = {"out.0": [2, 4, 64, 64]} + + Returns: + A list of torch tensors from the inference outputs + """ + # run inference + outputs = self._inference(compiled_model, inputs, io_shapes) + assert len(outputs) > 0, "Inference failed!" + + return outputs + + @property + @abstractmethod + def default_deployment(self) -> Deployment: + """Provides the default deployment config without the device key.""" + raise NotImplementedError + + @property + @abstractmethod + def deployment_table(self) -> DeploymentTable: + """Provides a set of supported values for each deployment config key.""" + raise NotImplementedError + + @abstractmethod + def _ir_to_compiled(self, ir_bytes: bytes, compilation_args: dict[str, Any] = {}) -> bytes: + """Converts a model from its intermediate representation (IR) to a compiled device model.""" + + @abstractmethod + def _profile( + self, compiled_model: bytes, compilation_args: dict[str, Any] = {} + ) -> tuple[float, DetailedResults]: + """Profiles a compiled device model and returns the latency & detailed profiling results.""" + + @abstractmethod + def _inference( + self, compiled_model: bytes, inputs: list[torch.Tensor], io_shapes: dict[str, list] = {} + ) -> list[torch.Tensor]: + """Run inference with the compiled model and return the output as list of torch tensors.""" diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/constants.py b/modelopt/torch/_deploy/_runtime/tensorrt/constants.py new file mode 100644 index 0000000000..d56180278c --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/tensorrt/constants.py @@ -0,0 +1,107 @@ +# 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. + +"""TensorRT server specific constants. + +If any constant is shared by both the TensorRT client and TensorRT server, +that should go into client/constants.py module and server should import from there. +Also if we open some of these settings for the clients, they should be moved there. +""" + +# Versions +TENSORRT_7_MAJOR_VERSION = 7 +TENSORRT_8_MAJOR_VERSION = 8 + +# Sizes (unit: MB) +ONE_MEBI = 1 +ONE_GIBI = 1024 + +ONE_MEBI_IN_BYTES = 1 << 20 +ONE_GIBI_IN_BYTES = 1 << 30 + +# TensorRT conversion tool names +TRTEXEC = "trtexec" + +# trtexec path within docker +TRTEXEC_PATH = "trtexec" + +# Default conversion params +DEFAULT_VALIDATION_THRESHOLD = 1e-4 +DEFAULT_MAX_BATCH_SIZE = 1 +DEFAULT_ACCELERATOR = "GPU" + +# With empty tactic source string trtexec will use all the available sources +DEFAULT_TACTIC_SOURCES = "" + +# NVTX annotation verbosity (default, verbose or none) +DEFAULT_NVTX_MODE = "none" + +# The minimum number of iterations used in kernel selection +DEFAULT_MIN_TIMING = 1 + +# The number of times averaged in each iteration for kernel selection +DEFAULT_AVG_TIMING = 1 + +# Default batch size for inference +DEFAULT_BATCH_SIZE = 1 + +# Default GPU Id +DEFAULT_GPU_ID = 0 + +# Default maximum workspace size +DEFAULT_MAX_WORKSPACE_SIZE = 256 + +# trtexec settings. +# Run for N milliseconds to warmup before measuring performance +WARMUP_TIME_MS = 500 + +# Profiling parameters +DEFAULT_PROFILING_RUNS = 1 +DEFAULT_NUM_INFERENCE_PER_RUN = 100 + +# Layerwise profiling parameters +INPUT_DATA_KEY = "input_data" +OUTPUT_DATA_KEY = "output_data" +UNNAMED_LAYER_KEY = "other" +NODE_NAME_DELIMITER = " + " + + +# Class representing all TensorRT modes expressed as text +class TRTMode: + FLOAT32 = "fp32" + FLOAT16 = "fp16" + BFLOAT16 = "bf16" + FLOAT8 = "fp8" + INT8 = "int8" + INT8_IQ = "int8_iq" + INT4 = "int4" + STRONGLY_TYPED = "stronglyTyped" + BEST = "best" + + +TRT_MODE_FLAGS = { + TRTMode.FLOAT32: [], + TRTMode.FLOAT16: ["--fp16"], + TRTMode.BFLOAT16: ["--bf16"], + TRTMode.FLOAT8: ["--fp16", "--fp8"], + TRTMode.INT8: ["--fp16", "--int8"], + TRTMode.INT8_IQ: ["--int8"], + TRTMode.INT4: ["--fp16", "--int4"], + TRTMode.STRONGLY_TYPED: ["--stronglyTyped"], + TRTMode.BEST: ["--best"], +} + + +SHA_256_HASH_LENGTH = 32 diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py b/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py new file mode 100644 index 0000000000..42ac07663b --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py @@ -0,0 +1,332 @@ +# 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 logging +import os +import subprocess # nosec +from tempfile import NamedTemporaryFile, TemporaryDirectory +from typing import Optional + +from ..._runtime.common import read_bytes, timeit, write_bytes, write_string +from ..._runtime.tensorrt.layerwise_profiling import process_layerwise_result +from ...utils import OnnxBytes +from .constants import ( + DEFAULT_NUM_INFERENCE_PER_RUN, + SHA_256_HASH_LENGTH, + TRT_MODE_FLAGS, + TRTEXEC_PATH, + WARMUP_TIME_MS, + TRTMode, +) +from .tensorrt_utils import convert_shape_to_string, prepend_hash_to_bytes + +try: + from trex import EnginePlan, layer_type_formatter, render_dot, to_dot + + HAVE_TREX = True +except ImportError: + HAVE_TREX = False + + +# TODO: Get rid of this function or get approval for `# nosec` usage if we want to include this +# as a non-compiled python file in the release. +def _run_command(cmd: str, cwd: str = None) -> tuple[int, str]: + """Util function to execute a command. + + This util will not direct stdout and stderr to console if the cmd succeeds. + + Args: + cmd: the command line string + cwd: current working directory + + Return: + return code: O means successful. + log_string: the stdout and stderr output as a string. + + """ + logging.info(cmd) + with NamedTemporaryFile("w+") as log: + p = subprocess.Popen(cmd.split(), stdout=log, stderr=log, cwd=cwd) # nosec + p.wait() + log.seek(0) + log_string = log.read() + if p.returncode != 0: + logging.error(log_string) + return p.returncode, log_string + + +def _get_profiling_params(profiling_runs: int) -> list[str]: + return [ + f"--warmUp={WARMUP_TIME_MS}", + f"--avgRuns={DEFAULT_NUM_INFERENCE_PER_RUN}", + f"--iterations={profiling_runs * DEFAULT_NUM_INFERENCE_PER_RUN}", + ] + + +def _get_trtexec_params( + engine_path: str, + builder_optimization_level: str, + timing_cache_file: str = None, +) -> list[str]: + cmd = ["--saveEngine=" + engine_path] + + cmd.append("--skipInference") + + if timing_cache_file: + cmd.append("--timingCacheFile=" + timing_cache_file) + + # set the builder optimization level to 4 + cmd.append("--builderOptimizationLevel=" + builder_optimization_level) + + # Generates an engine layer graph file for visualization. + cmd += [ + "--verbose", + f"--exportLayerInfo={engine_path}.graph.json", + ] + + return cmd + + +def _draw_engine(engine_json_fname: str) -> Optional[bytes]: + if not HAVE_TREX: + logging.exception("`trex` is not installed. Skipping engine graph drawing.") + return None + + plan = EnginePlan(engine_json_fname) + formatter = layer_type_formatter + display_regions = True + expand_layer_details = False + + graph = to_dot( + plan, formatter, display_regions=display_regions, expand_layer_details=expand_layer_details + ) + render_dot(graph, engine_json_fname, "svg") + + with open(f"{engine_json_fname}.svg", "rb") as svg_file: + return svg_file.read() + + +def _is_low_bit_mode(trt_mode: str) -> bool: + return trt_mode in [ + TRTMode.INT8, + TRTMode.INT8_IQ, + TRTMode.INT4, + TRTMode.FLOAT8, + TRTMode.BEST, + TRTMode.STRONGLY_TYPED, + ] + + +def _update_dynamic_shapes(dynamic_shapes: dict, cmd: list[str]) -> None: + """Update the dynamic shapes in the command list.""" + shapes = ["minShapes", "optShapes", "maxShapes"] + for shape in shapes: + if shape in dynamic_shapes: + cmd.extend( + [ + f"--{shape}=" + convert_shape_to_string(dynamic_shapes[shape]), + ] + ) + + +@timeit +def build_engine( + onnx_bytes: OnnxBytes, + trt_mode: str = TRTMode.FLOAT32, + calib_cache: str = None, + dynamic_shapes: dict = None, + plugin_config: dict = None, + builder_optimization_level: str = "3", + draw_engine: bool = False, +) -> tuple[Optional[bytes], bytes, Optional[bytes]]: + """This method produces serialized TensorRT engine from an ONNX model. + + Args: + onnx_bytes: Data of the ONNX model stored as an OnnxBytes object. + trt_mode: The precision with which the TensorRT engine will be built. Supported modes are: + - TRTMode.FLOAT32 + - TRTMode.FLOAT16 + - TRTMode.BFLOAT16 + - TRTMode.INT8 + - TRTMode.FLOAT8 + - TRTMode.INT4 + - TRTMode.STRONGLY_TYPED + - TRTMode.BEST + For more details about TensorRT modes, please refer to: + https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#trtexec-flags + calib_cache: Calibration data cache. + dynamic_shapes: Dictionary of dynamic shapes for the input tensors. Example is as follows: + { + "minShapes": {"input": [1,3,244,244]}, + "optShapes": {"input": [16,3,244,244]}, + "maxShapes": {"input": [32,3,244,244]} + } + plugin_config: Dictionary of plugin configurations. Example is as follows: + { + "staticPlugins": ["staticPluginOne.so", "staticPluginTwo.so", ...], + "dynamicPlugins": ["dynamicPluginOne.so", "dynamicPluginTwo.so", ...], + "setPluginsToSerialize": ["dynamicPluginOne.so", "pluginSerializeOne.so", ...], + "ignoreParsedPluginLibs": False + } + builder_optimization_level: Optimization level for the TensorRT builder. + For more details, please refer to: + https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/BuilderConfig.html + draw_engine: True or False based on whether to draw the engine graph or not. + + Returns: + The generated engine file data. + + The stdout log produced by trtexec tool. \ + [If there is subprocess.CalledProcessError, this byte variable is transferred to str] + + The generated TensorRT engine graph as svg. + """ + with TemporaryDirectory() as working_dir: + onnx_path = os.path.join(working_dir, "onnx") + calib_cache_path = os.path.join(working_dir, "calib_cache") + + onnx_bytes.write_to_disk(onnx_path) + engine_path = os.path.join( + working_dir, f"{onnx_bytes.model_name}/{onnx_bytes.model_name}.engine" + ) + # TODO: Enable timing cache + + cmd = [ + TRTEXEC_PATH, + "--onnx=" + onnx_path + f"/{onnx_bytes.model_name}.onnx", + ] + cmd.extend(TRT_MODE_FLAGS[trt_mode]) + if trt_mode == TRTMode.INT8 and calib_cache: + write_string(calib_cache, calib_cache_path) + cmd.append(f"--calib={calib_cache_path}") + + if dynamic_shapes: + _update_dynamic_shapes(dynamic_shapes, cmd) + + if plugin_config: + if "staticPlugins" in plugin_config: + for plugin in plugin_config["staticPlugins"]: + cmd.append(f"--staticPlugins={plugin}") + if "dynamicPlugins" in plugin_config: + for plugin in plugin_config["dynamicPlugins"]: + cmd.append(f"--dynamicPlugins={plugin}") + if "setPluginsToSerialize" in plugin_config: + for plugin in plugin_config["setPluginsToSerialize"]: + cmd.append(f"--setPluginsToSerialize={plugin}") + if "ignoreParsedPluginLibs" in plugin_config: + cmd.append("--ignoreParsedPluginLibs") + + builder_optimization_level = ( + "4" if _is_low_bit_mode(trt_mode) else builder_optimization_level + ) + cmd += _get_trtexec_params( + engine_path, + builder_optimization_level, + ) + + try: + engine_directory = os.path.dirname(engine_path) + os.makedirs(engine_directory, exist_ok=True) + trtexec_cmd = " ".join(cmd) + ret_code, out = _run_command(trtexec_cmd) + if ret_code != 0: + return None, out.encode(), None + + # _run_command will handle the exceptions from trtexec + # For other errors, they will be caught in the block below + except Exception as e: + out = str(e) + logging.exception(out) + return None, out.encode(), None + + engine_bytes = read_bytes(engine_path) + engine_graph_bytes = _draw_engine(f"{engine_path}.graph.json") if draw_engine else None + + engine_bytes = prepend_hash_to_bytes(engine_bytes) + + return engine_bytes, out.encode(), engine_graph_bytes + + +@timeit +def profile_engine( + engine_bytes: bytes, + profiling_runs: int = 1, + onnx_node_names: list[str] = None, + enable_layerwise_profiling: bool = False, + dynamic_shapes: dict = None, +) -> tuple[Optional[dict[str, float]], bytes]: + """This method produces profiles a TensorRT engine and returns the detailed results. + + Args: + engine_bytes: Bytes of the serialized TensorRT engine, prepended with a SHA256 hash. + onnx_node_names: List of node names in the onnx model. + profiling_runs: number of profiling runs. Each run runs `DEFAULT_NUM_INFERENCE_PER_RUN` inferences. + enable_layerwise_profiling: + True or False based on whether layerwise profiling is required or not. + dynamic_shapes: Dictionary of dynamic shapes for the input tensors. Example is as follows: + { + "minShapes": {"input": [1,3,244,244]}, + "optShapes": {"input": [16,3,244,244]}, + "maxShapes": {"input": [32,3,244,244]} + } + + Returns: + Layerwise profiling output as a json string. + Stdout log produced by trtexec tool. + """ + with TemporaryDirectory() as working_dir: + engine_path = os.path.join(working_dir, "engine") + profile_path = os.path.join(working_dir, "profile") + + engine_bytes = engine_bytes[SHA_256_HASH_LENGTH:] # Remove the hash + + write_bytes(engine_bytes, engine_path) + + cmd = [TRTEXEC_PATH, "--loadEngine=" + engine_path] + + cmd += _get_profiling_params(profiling_runs) + + if enable_layerwise_profiling: + cmd += [ + "--dumpProfile", + "--separateProfileRun", + "--exportProfile=" + profile_path, + ] + + if dynamic_shapes: + _update_dynamic_shapes(dynamic_shapes, cmd) + + try: + _, out = _run_command(" ".join(cmd)) + + layerwise_results = {} # empty dictionary + if enable_layerwise_profiling and os.path.exists(profile_path): + layerwise_results = process_layerwise_result(profile_path, onnx_node_names) + + return layerwise_results, out.encode() + # If _run_command has error, + # the error will be subprocess.CalledProcessError + # We will catch this error in this block and send back the str version of the error message(e.output) + # The reason we have this independent block is that the error message, e.output, + # is byte type with "utf-8" coding, different from other error. + except subprocess.CalledProcessError as e: + out = str(e) + logging.exception(out) + return None, out.encode() + # For all other errors, they will be caught in the block below, error info str(e) will be sent back + except Exception as e: + out = str(e) + logging.exception(out) + return None, out.encode() diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/hw_param_config.py b/modelopt/torch/_deploy/_runtime/tensorrt/hw_param_config.py new file mode 100644 index 0000000000..d123e4247c --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/tensorrt/hw_param_config.py @@ -0,0 +1,51 @@ +# 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. + +"""Hardware specific parameters. + +All the hardware parameters will be treated as dictionary type for omniengine client APIs. +The key name and the suggested value range are listed with name TENSORRT_HW_PARAMS_SUGGESTED_OPTIONS. +The key name and the optimum value are listed with name TENSORRT_HW_PARAMS_OPT_OPTIONS. +""" + +from .constants import ( + DEFAULT_AVG_TIMING, + DEFAULT_MAX_WORKSPACE_SIZE, + DEFAULT_MIN_TIMING, + DEFAULT_TACTIC_SOURCES, +) + +# Key names +# Workspace unit: MB +MAX_WORKSPACE_SIZE = [16, 32, 64, 128, 256, 512, 1024] +TACTIC_SOURCES = ["cublasLt", "cublas", "cudnn"] +ALL_TATIC_SOURCES_COMPONENT = [] +for source in TACTIC_SOURCES: + ALL_TATIC_SOURCES_COMPONENT.append("+" + source) + ALL_TATIC_SOURCES_COMPONENT.append("-" + source) + +TENSORRT_HW_PARAMS_SUGGESTED_OPTIONS = { + "tacticSources": ALL_TATIC_SOURCES_COMPONENT, + "minTiming": range(1, 5), + "avgTiming": range(1, 5), + "workspace": MAX_WORKSPACE_SIZE, +} + +TENSORRT_HW_PARAMS_OPT_OPTIONS = { + "tacticSources": DEFAULT_TACTIC_SOURCES, + "minTiming": str(DEFAULT_MIN_TIMING), + "avgTiming": str(DEFAULT_AVG_TIMING), + "workspace": str(DEFAULT_MAX_WORKSPACE_SIZE), +} diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/layerwise_profiling.py b/modelopt/torch/_deploy/_runtime/tensorrt/layerwise_profiling.py new file mode 100644 index 0000000000..2fd0480f5c --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/tensorrt/layerwise_profiling.py @@ -0,0 +1,175 @@ +# 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 +import logging + +from ..common import read_string +from .constants import INPUT_DATA_KEY, NODE_NAME_DELIMITER, OUTPUT_DATA_KEY, UNNAMED_LAYER_KEY + + +def _merge_reformatters(layer_latency_dict: dict[str, float]) -> dict[str, float]: + reformatter_delimiters = [" input reformatter ", " output reformatter ", " to "] + keys2merge = [k for k in layer_latency_dict if any(rn in k for rn in reformatter_delimiters)] + + for key in keys2merge: + layer_latency = layer_latency_dict.pop(key) + for delimiter in reformatter_delimiters: + if delimiter in key: + split_key = key.split(delimiter) + if len(split_key) != 2: + logging.error( + f"Key splitting in layerwise profiling failed for reformatter node: {key}." + ) + else: + key = split_key[1] if delimiter == " to " else split_key[0] + break + + layer_latency_dict[key] = layer_latency_dict.get(key, 0.0) + layer_latency + + return layer_latency_dict + + +def process_layerwise_result(profile_path: str, onnx_node_names: list[str]) -> dict[str, float]: + """This module process the layerwise profiling result to make them mappable to PyTorch. + + Args: + profile_path: Path to the layerwise profiling dump from trtexec. + onnx_node_names: List of node names in the onnx model. + + Returns: + A dictionary mapping layer names to latency in ms. For example, + + .. code-block:: python + + { + "Conv_0 + Clip_3": 0.016, + "Conv_4": 0.008, + "Gemm_169": 0.015, + "input_data": 0.010, + "output_data": 0.040, + } + """ + layerwise_profiling = read_string(profile_path) + layerwise_results = json.loads(layerwise_profiling)[1:] + + # Just keep averageMs from layerwise profiling result + layer_latency_dict = {} + for results in layerwise_results: + layer_latency_dict[results["name"].replace("onnx::", "")] = results["averageMs"] + + # Merge the input/output reformatter op into its following op + layer_latency_dict = _merge_reformatters(layer_latency_dict) + + return map_trt_layers_to_onnx(layer_latency_dict, onnx_node_names) + + +def map_trt_layers_to_onnx( + layerwise_result: dict[str, float], onnx_node_names: list[str] +) -> dict[str, float]: + """This module maps the TensorRT layers from profiling result with onnx nodes. + + Args: + layerwise_result: Layerwise profiling result. + onnx_node_names: Onnx node names. + + Returns: + A dictionary mapping layer names to latency in ms. + """ + + def _group_split(key: str, delimiter: str): + parenthesis_balance = 0 + key_group = [] + current_key = "" + + for ch in key: + if ch == "(": + parenthesis_balance += 1 + elif ch == ")": + parenthesis_balance -= 1 + elif ch == delimiter and parenthesis_balance == 0: + key_group.append(current_key.strip()) + current_key = "" + continue + + current_key += ch + + # Insert the last key in group + key_group.append(current_key.strip()) + + return key_group + + def _remove_non_onnx_nodes(layer: str) -> tuple[str]: + keys = [] + + # To match the longer name first, match from the last node + for node_name in reversed(onnx_node_names): + if node_name in layer: + keys.append(node_name) + layer = layer.replace(node_name, "") # Skip from further matching + + cleaned_layer = NODE_NAME_DELIMITER.join(reversed(keys)) + + # Lets check if the layer has any input/output data timing + if not cleaned_layer and "input" in layer: + return (INPUT_DATA_KEY,) + if not cleaned_layer and "output" in layer: + return (OUTPUT_DATA_KEY,) + + if not onnx_node_names: # empty list if omnimizer<=v0.4.1 + return (layer,) + + return (cleaned_layer,) if len(cleaned_layer) > 0 else () # type: ignore[return-value] + + def _iterative_split_key(key: str) -> tuple: + if ": " in key: + # e.g., 2-layer MLP: Conv_68 + Relu_69 -> Conv_70 + if len(key.split(": ")) != 2: + logging.error(f"Key splitting in layerwise profiling failed for key: {key}.") + return _remove_non_onnx_nodes(key) + return _iterative_split_key(key.split(": ")[1]) + elif key.startswith("PWN"): + # e.g., PWN(Clip_589 + (Unnamed Layer* 16) [Shuffle], PWN(PWN(PWN(Add_586 + (Unnamed Layer* 11) [Shuffle] + # + Add_43, Clip_46), Mul_47), Div_49)) + if not key.endswith(")"): + logging.error(f"Key: {key} starts with 'PWN' but does not end with a ')'.") + key = key[4:] + else: + key = key[4:-1] + return sum((_iterative_split_key(k) for k in _group_split(key, ",")), ()) + elif " + " in key: + # e.g. Add_586 + (Unnamed Layer* 11) [Shuffle] + Add_43 + return sum((_iterative_split_key(k) for k in _group_split(key, "+")), ()) + elif " -> " in key: + # e.g. Conv_68 + Relu_69 -> Conv_70 + key = key.split(" -> ") + if len(key) != 2: + logging.error(f"Key splitting in layerwise profiling failed for key: {key}.") + return _remove_non_onnx_nodes(key) + + covered_keys = _group_split(key[0], "+") + [key[1]] + if onnx_node_names: + idx1, idx2 = onnx_node_names.index(key[0]), onnx_node_names.index(key[1]) + covered_keys = tuple(onnx_node_names[idx1 : idx2 + 1]) + return sum((_iterative_split_key(k) for k in covered_keys), ()) + + return _remove_non_onnx_nodes(key) + + mapped_layerwise_result = {} + for k, v in layerwise_result.items(): + new_key = NODE_NAME_DELIMITER.join(_iterative_split_key(k)) or UNNAMED_LAYER_KEY + mapped_layerwise_result[new_key] = mapped_layerwise_result.get(new_key, 0.0) + v + + return mapped_layerwise_result diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/parse_trtexec_log.py b/modelopt/torch/_deploy/_runtime/tensorrt/parse_trtexec_log.py new file mode 100644 index 0000000000..ba7e3cf3c0 --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/tensorrt/parse_trtexec_log.py @@ -0,0 +1,154 @@ +# 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. + +# Reference: https://github.com/NVIDIA/TensorRT/tree/release/9.1/tools/experimental/trt-engine-explorer/utils + + +"""trtexec log file parsing""" + +import re +from typing import Any + + +def __to_float(line: str) -> float: + """Scan the input string and extract the first float instance.""" + # https://docs.python.org/3/library/re.html#simulating-scanf + float_match = re.search(r"[-+]?(\d+(\.\d*)?|\.\d+)([eE][-+]?\d+)?", line) + if float_match is None: + raise ValueError + start, end = float_match.span() + return float(line[start:end]) + + +def __get_stats(line: str) -> list[float]: + """Parse a string containing pairs of "key = value" and return the list of values. + + Here's a sample input line: "min = 0.87854 ms, max = 0.894043 ms, mean = 0.881251 ms" + The values are expected to be floats. + Split the kv list to "k = v" substrings, then split each substring to + k, v and return float(v) + """ + return [__to_float(substr.split("=")[1]) for substr in line.split(",")] + + +class FileSection: + def __init__(self, section_header: str): + self.section_header = section_header + self.dict = {} + + def entered_section(self, line: str): + s = re.search(self.section_header, line) + return s is not None + + def parse_line(self, line: str): + def parse_kv_line(line: str) -> tuple[Any, Any]: + """Parse a log line that reports a key-value pair. + + The log line has this format: [mm/dd/yyyy-hh:mm:ss] [I] key_name: key_value + """ + match = re.search(r"(\[\d+/\d+/\d+-\d+:\d+:\d+\] \[I\] )", line) + if match is not None: + match_end = match.span()[1] + kv_line = line[match_end:].strip() + kv = kv_line.split(": ") + if len(kv) > 1: + return kv[0], kv[1] + return None, None + + k, v = parse_kv_line(line) + if k is not None and v is not None: + self.dict[k] = v + return True + if k is not None: + return True + return False + + +def __parse_log_file(trtexec_log: str, sections: list) -> list[dict]: + current_section = None + for line in trtexec_log.split("\n"): + if current_section is None: + for section in sections: + if section.entered_section(line): + current_section = section + break + else: + if not current_section.parse_line(line): + current_section = None + dicts = [section.dict for section in sections] + return dicts + + +def parse_build_log(trtexec_log: str) -> dict[str, Any]: + """Parse the TensorRT engine build log and extract the builder configuration. + + Args: + trtexec_log: The log file generated by trtexec. + + Returns: + A dictionary containing the model options and build options. + """ + model_options = FileSection("=== Model Options ===") + build_options = FileSection("=== Build Options ===") + sections = [model_options, build_options] + __parse_log_file(trtexec_log, sections) + return { + "model_options": model_options.dict, + "build_options": build_options.dict, + } + + +def parse_profiling_log(trtexec_log: str) -> dict[str, Any]: + """Parse the TensorRT engine profiling log and extract the performance summary. + + Args: + trtexec_log: The log file generated by trtexec. + + Returns: + A dictionary containing the performance summary, inference options and device information. + """ + performance_summary = FileSection("=== Performance summary ===") + inference_options = FileSection("=== Inference Options ===") + device_information = FileSection("=== Device Information ===") + sections = [performance_summary, inference_options, device_information] + __parse_log_file(trtexec_log, sections) + + def post_process_perf(perf_summary: dict): + """Normalize the log results to a standard format""" + for k, v in perf_summary.items(): + if k in ["Throughput", "Total Host Walltime", "Total GPU Compute Time"]: + perf_summary[k] = __to_float(v) + if k in ["Latency", "Enqueue Time", "H2D Latency", "GPU Compute Time", "D2H Latency"]: + perf_summary[k] = __get_stats(v) + return perf_summary + + def post_process_device_info(device_info: dict): + """Convert some value fields to float""" + for k, v in device_info.items(): + if k in [ + "Compute Clock Rate", + "Memory Bus Width", + "Memory Clock Rate", + "Compute Capability", + "SMs", + ]: + device_info[k] = __to_float(v) + return device_info + + return { + "performance_summary": post_process_perf(performance_summary.dict), + "inference_options": inference_options.dict, + "device_information": post_process_device_info(device_information.dict), + } diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/tensorrt_utils.py b/modelopt/torch/_deploy/_runtime/tensorrt/tensorrt_utils.py new file mode 100644 index 0000000000..54f85605ec --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/tensorrt/tensorrt_utils.py @@ -0,0 +1,198 @@ +# 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 hashlib +import logging + +import numpy as np +import onnx +import tensorrt as trt +import torch + +from modelopt.onnx.utils import get_batch_size +from modelopt.onnx.utils import get_input_names as get_onnx_input_names + +from .constants import TENSORRT_8_MAJOR_VERSION, TRTMode + + +def is_trt8(): + return int(trt.__version__.split(".", maxsplit=1)[0]) == TENSORRT_8_MAJOR_VERSION + + +class HostDeviceMem(object): + """Simple helper data class to hold host and device memory pointers.""" + + def __init__(self, host_mem, device_mem, dtype): + self.host = host_mem + self.device = device_mem + self.dtype = dtype + + def __str__(self): + return ( + "Host:\n" + + str(self.host) + + "\nDevice:\n" + + str(self.device) + + "\nType:\n" + + str(self.dtype) + ) + + def __repr__(self): + return self.__str__() + + def __del__(self): + del self.host + del self.device + + +def get_engine_bytes(engine: trt.tensorrt.ICudaEngine) -> bytes: + """Return serialized TensorRT engine bytes.""" + return bytearray(engine.serialize()) + + +def load_engine(buffer: bytes, log_level: int = trt.Logger.ERROR) -> trt.tensorrt.ICudaEngine: + """Load a TensorRT engine from engine data and return.""" + try: + trt_logger = trt.Logger(log_level) + with trt.Runtime(trt_logger) as runtime: + return runtime.deserialize_cuda_engine(buffer), "" + except Exception as e: + logging.exception(str(e)) + return None, str(e) + + +def get_input_names(engine: trt.tensorrt.ICudaEngine) -> list[str]: + """Gather the input names from an ICudaEngine. + + Args: + engine: TensorRT engine object. + + Returns: + List of engine input names. + """ + input_names = [] + for binding_index in range(engine.num_bindings): + if engine.binding_is_input(binding_index): + input_names.append(engine.get_binding_name(binding_index)) + return input_names + + +def get_output_names(engine: trt.tensorrt.ICudaEngine) -> list[str]: + """Gather the output names from an ICudaEngine. + + Args: + engine: TensorRT engine object. + + Returns: + List of engine output names. + """ + output_names = [] + for binding_index in range(engine.num_bindings): + if not engine.binding_is_input(binding_index): + output_names.append(engine.get_binding_name(binding_index)) + return output_names + + +def get_output_shapes( + engine: trt.tensorrt.ICudaEngine, + context: trt.tensorrt.IExecutionContext, +) -> list[list[int]]: + """Gather the output shapes from an ICudaEngine. + + Args: + engine: TensorRT engine object. + context: Current execution context for the inference. + + Returns: + List of shapes of outputs which are list of integers. + """ + assert context.all_binding_shapes_specified + assert context.all_shape_inputs_specified + + output_shapes = [] + for binding_index in range(engine.num_bindings): + if not engine.binding_is_input(binding_index): + shape = context.get_binding_shape(binding_index) + output_shapes.append(shape) + return output_shapes + + +def validate_precision(precision: str) -> bool: + """Returns whether an input precision is in supported set.""" + return precision in [TRTMode.FLOAT32, TRTMode.FLOAT16, TRTMode.INT8] + + +def calib_data_generator(onnx_bytes: bytes, input_tensors: list[np.ndarray]): + """The calibation data generator that yields calibration feed_dict to tensorrt.""" + input_names = get_onnx_input_names(onnx.load_from_string(onnx_bytes)) + + batch_size = get_batch_size(onnx.load_from_string(onnx_bytes)) + if not batch_size or batch_size <= 0: + batch_size = 1 + # If input tensor batch % batch_size != 0, we don't use all input tensors for calibration. + num_batches = int(input_tensors[0].shape[0] / batch_size) + + for i in range(num_batches): + feed_dict = {} + tensor_batch_dim_index = i * batch_size + for idx, input_name in enumerate(input_names): + feed_dict[input_name] = input_tensors[idx][ + tensor_batch_dim_index : (tensor_batch_dim_index + batch_size) + ] + yield feed_dict + + +def convert_trt_dtype_to_torch(trt_dtype: trt.tensorrt.DataType) -> torch.dtype: + """Convert TensorRT data type to torch data type.""" + trt_to_torch_dtype_map = { + trt.DataType.FLOAT: torch.float32, + trt.DataType.HALF: torch.float16, + trt.DataType.BF16: torch.bfloat16, + trt.DataType.INT8: torch.int8, + trt.DataType.INT32: torch.int32, + trt.DataType.INT64: torch.int64, + trt.DataType.BOOL: torch.bool, + } + + assert trt_dtype in trt_to_torch_dtype_map, f"Unsupported TensorRT data type: {trt_dtype}" + return trt_to_torch_dtype_map[trt_dtype] + + +def prepend_hash_to_bytes(engine_bytes: bytes) -> bytes: + """Prepend the engine bytes with the SHA256 hash of the engine bytes + This has will serve as a unique identifier for the engine and will be used to manage + TRTSessions in the TRTClient. + """ + hash_object = hashlib.sha256(engine_bytes) + hash_object.update(engine_bytes) + hash_bytes = hash_object.digest() + engine_bytes = hash_bytes + engine_bytes + return engine_bytes + + +def convert_shape_to_string(shape: dict[str, list]) -> str: + """Convert a shape dictionary to a string. + For example, if the shape is: + { + "input": [1, 3, 224, 224], + "output": [1, 1000] + }. + The output string will be: + input:1x3x244x244,output:1x1000 + """ + result = "" + for key, value in shape.items(): + result += f"{key}:{'x'.join(map(str, value))}," + return result[:-1] diff --git a/modelopt/torch/_deploy/_runtime/trt_client.py b/modelopt/torch/_deploy/_runtime/trt_client.py new file mode 100644 index 0000000000..fc56d8fc63 --- /dev/null +++ b/modelopt/torch/_deploy/_runtime/trt_client.py @@ -0,0 +1,217 @@ +# 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. + +from typing import Any + +import tensorrt as trt +import torch + +from modelopt.onnx.utils import get_node_names_from_bytes + +from ..utils import OnnxBytes +from .registry import RuntimeRegistry +from .runtime_client import Deployment, DeploymentTable, DetailedResults, RuntimeClient +from .tensorrt.constants import SHA_256_HASH_LENGTH +from .tensorrt.engine_builder import build_engine, profile_engine +from .tensorrt.parse_trtexec_log import parse_profiling_log +from .tensorrt.tensorrt_utils import convert_trt_dtype_to_torch + +__all__ = ["TRTLocalClient"] + + +@RuntimeRegistry.register("TRT") +class TRTLocalClient(RuntimeClient): + """A client for using the local TRT runtime with GPU backend.""" + + @property + def default_deployment(self) -> Deployment: + return {k: v[0] for k, v in self.deployment_table.items()} + + @property + def deployment_table(self) -> DeploymentTable: + return { + "version": ["8.6", "9.1", "9.2", "9.3", "10.0", "10.1", "10.2", "10.3"], + "accelerator": ["GPU"], + "precision": [ + "fp32", + "fp16", + "bf16", + "fp8", + "int8", + "int8_iq", + "int4", + "stronglyTyped", + "best", + ], + # Support ONNX opsets 13-20 + "onnx_opset": [str(i) for i in range(13, 21)], + } + + def __init__(self, deployment: Deployment): + """Initialize a TRTLocalClient with the given deployment.""" + super().__init__(deployment) + self.inference_sessions = {} + logger = trt.Logger(trt.Logger.WARNING) + self.trt_runtime = trt.Runtime(logger) + assert trt.init_libnvinfer_plugins(logger, ""), "Failed to initialize nvinfer plugins." + self.stream = torch.cuda.Stream() + + def _ir_to_compiled(self, ir_bytes: bytes, compilation_args: dict[str, Any] = None) -> bytes: + """Converts an ONNX model to a compiled TRT engine. + + Args: + ir_bytes: The ONNX model bytes. + compilation_args: A dictionary of compilation arguments. Supported args: dynamic_shapes, plugin_config. + + Returns: + The compiled TRT engine bytes. + """ + onnx_bytes = OnnxBytes.from_bytes(ir_bytes) + onnx_model_file_bytes = onnx_bytes.get_onnx_model_file_bytes() + self.node_names = get_node_names_from_bytes(onnx_model_file_bytes) + engine_bytes, _, _ = build_engine( + onnx_bytes, + dynamic_shapes=compilation_args.get("dynamic_shapes", None), + plugin_config=compilation_args.get("plugin_config", None), + trt_mode=self.deployment["precision"], + ) + self.engine_bytes = engine_bytes + return engine_bytes + + def _profile( + self, compiled_model: bytes, compilation_args: dict[str, Any] = None + ) -> tuple[float, DetailedResults]: + node_names = None + enable_layerwise_profiling = False + if hasattr(self, "node_names"): + node_names = self.node_names + enable_layerwise_profiling = True + _, trtexec_log = profile_engine( + compiled_model, + onnx_node_names=node_names, + enable_layerwise_profiling=enable_layerwise_profiling, + dynamic_shapes=compilation_args.get("dynamic_shapes", None), + ) + profiling_results = parse_profiling_log(trtexec_log.decode()) + latency = 0.0 + detailed_results = {} + if profiling_results is not None: + gpu_compute_time = profiling_results["performance_summary"]["GPU Compute Time"] + latency = sum(gpu_compute_time) / len(gpu_compute_time) + detailed_results = profiling_results + return latency, detailed_results + + def _inference( + self, compiled_model: bytes, inputs: list[torch.Tensor], io_shapes: dict[str, list] = {} + ) -> list[torch.Tensor]: + """Run inference with the compiled model and return the output as list of torch tensors.""" + assert compiled_model is not None, "Engine bytes are not set." + + model_hash = compiled_model[:SHA_256_HASH_LENGTH] + if model_hash not in self.inference_sessions: + model_bytes = compiled_model[SHA_256_HASH_LENGTH:] + self.inference_sessions[model_hash] = self.TRTSession( + model_bytes, self.trt_runtime, self.stream, inputs[0].device.type, io_shapes + ) + return self.inference_sessions[model_hash].run(inputs) + + def _teardown_all_sessions(self): + """Clean up all TRT sessions.""" + for session in self.inference_sessions.values(): + del session + self.inference_sessions = {} + + class TRTSession: + def __init__(self, compiled_model, trt_runtime, stream, device, io_shapes): + self.engine = trt_runtime.deserialize_cuda_engine(compiled_model) + assert self.engine is not None, "Engine deserialization failed." + self.execution_context = self.engine.create_execution_context() + self.stream = stream + self.device = device + self.io_shapes = io_shapes + self.input_tensors, self.output_tensors = self.initialize_input_output_tensors( + self.engine + ) + + def initialize_input_output_tensors(self, engine): + # Allocate torch tensors for inputs and outputs + input_tensors = [] + output_tensors = [] + for idx in range(engine.num_io_tensors): + tensor_name = engine.get_tensor_name(idx) + if tensor_name not in self.io_shapes: + try: + tensor_shape = engine.get_tensor_profile_shape(tensor_name, 0)[0] + except Exception as e: + raise ValueError( + f"Shape for tensor {tensor_name} is not specified. " + f"Please provide the shape with the io_shapes input. Error: {e}" + ) + + else: + tensor_shape = trt.Dims(self.io_shapes[tensor_name]) + tensor_dtype = convert_trt_dtype_to_torch(engine.get_tensor_dtype(tensor_name)) + torch_tensor = torch.empty(tuple(tensor_shape), dtype=tensor_dtype, device="cuda") + if engine.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT: + input_tensors.append(torch_tensor) + self.execution_context.set_tensor_address( + tensor_name, input_tensors[idx].data_ptr() + ) + self.execution_context.set_input_shape(tensor_name, tensor_shape) + else: + output_tensors.append(torch_tensor) + self.execution_context.set_tensor_address( + tensor_name, + output_tensors[idx - len(input_tensors)].data_ptr(), + ) + assert ( + self.execution_context.all_shape_inputs_specified + ), "Not all shape inputs are specified." + + # Set selected profile idx + self.execution_context.set_optimization_profile_async(0, self.stream.cuda_stream) + + # Assertion: to ensure all the shapes can be inferred when user does not provide them + if self.io_shapes == {}: + assert ( + len(self.execution_context.infer_shapes()) == 0 + ), "Shapes of all the bindings cannot be inferred." + + return input_tensors, output_tensors + + def run(self, inputs): + assert self.engine is not None, "Engine is not set." + + # Copy inputs to GPU + with torch.cuda.stream(self.stream): + for i, input_t in enumerate(inputs): + if self.input_tensors[i].shape == input_t.shape: + self.input_tensors[i].copy_(input_t, non_blocking=True) + else: + # Zero out the self.input_tensor[i] + self.input_tensors[i].zero_() + + # Create slices for each dimension of input_t + slices = tuple(slice(0, dim) for dim in input_t.shape) + + # Assign the portion of input_t to the self.input_tensor[i] + self.input_tensors[i][slices] = input_t + + # Run inference + self.execution_context.execute_async_v3(stream_handle=self.stream.cuda_stream) + + self.stream.synchronize() + + return self.output_tensors diff --git a/modelopt/torch/_deploy/compilation.py b/modelopt/torch/_deploy/compilation.py new file mode 100644 index 0000000000..811c7f25a4 --- /dev/null +++ b/modelopt/torch/_deploy/compilation.py @@ -0,0 +1,143 @@ +# 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. + +"""Module to compile a model for a target device.""" + +from typing import Any, Union + +import torch.nn as nn + +from modelopt.torch.utils import is_channels_last + +from ._runtime import Deployment, RuntimeRegistry +from .device_model import DeviceModel +from .utils import DEFAULT_ONNX_OPSET, OnnxBytes, get_onnx_bytes_and_metadata + +__all__ = ["compile"] + + +def compile( + model: nn.Module, + dummy_input: Union[Any, tuple], + deployment: Deployment, + dynamic_axes: dict = {}, + compilation_args: dict = {}, +) -> DeviceModel: + """Compile a given torch model into a device model according to the provided deployment. + + Args: + model: PyTorch model to compile for target device. + dummy_input: Arguments of ``model.forward()``. This is used for exporting and calculating + inference-based metrics, such as latency/FLOPs. The format of ``dummy_inputs`` follows + the convention of the ``args`` argument in + `torch.onnx.export `_. + Specifically, ``dummy_input`` can be: + + #. a single argument (``type(dummy_input) != tuple``) corresponding to + + .. code-block:: python + + model.forward(dummy_input) + + #. a tuple of arguments corresponding to + + .. code-block:: python + + model.forward(*dummy_input) + + #. a tuple of arguments such that ``type(dummy_input[-1]) == dict`` corresponding to + + .. code-block:: python + + model.forward(*dummy_input[:-1], **dummy_input[-1]) + + .. warning:: + + In this case the model's ``forward()`` method **cannot** contain keyword-only + arguments (e.g. ``forward(..., *, kw_only_args)``) or variable keyword arguments + (e.g. ``forward(..., **kwargs)``) since these cannot be sorted into positional + arguments. + + .. note:: + + In order to pass a dict as last non-keyword argument, you need to use a tuple as + ``dummy_input`` and add an *empty* dict as the last element, e.g., + + .. code-block:: python + + dummy_input = (x, {"y": y, "z": z}, {}) + + The empty dict at the end will then be interpreted as the keyword args. + + See `torch.onnx.export `_ + for more info. + + Note that if you provide a ``{arg_name}`` with batch size ``b``, the results will be + computed based on batch size ``b``. + deployment: Deployment configuration with keys as specified below. + + * ``runtime``: the desired runtime for deployment (*required*); + * ``accelerator``: the accelerator on the device to be used (*optional*); + * ``version``: the version of runtime to be used (*optional*); + * ``precision``: the desired precision (*optional*); + * ``onnx_opset``: the opset version (*optional*). + + An example of a deployment configuration is: + + .. code-block:: python + + deployment = { + "runtime": "ORT", + "accelerator": "CPU", + "version": "1.11", + "precision": "fp32", + "onnx_opset": 13, + } + dynamic_axes: A dictionary that maps input names to the dynamic axes of the model. This is + useful when the model has dynamic input shapes. The keys are the input names and the + values are the dynamic axes. For example, if the model has a dynamic batch size, the + dictionary would look like: + + .. code-block:: python + + dynamic_axes = {"input": {0: "batch", 1: "sequence"}} + + See `torch.onnx.export `_ + for more info. + compilation_args: Additional arguments for the compilation process. + Returns: + An instance of DeviceModel. + """ + # Add check for nhwc format before we formally support NHWC models + assert not is_channels_last(model), "Only NCHW models are supported!" + + # Export onnx model and get some required names from it + onnx_bytes, metadata = get_onnx_bytes_and_metadata( + model, + dummy_input, + dynamic_axes=dynamic_axes, + onnx_opset=int(deployment.get("onnx_opset", DEFAULT_ONNX_OPSET)), + ) + + client = RuntimeRegistry.get(deployment) + + # For the ORTLocalClient, we need to pass the bytes of the onnx model instead of the OnnxBytes object + if client.__class__.__name__ == "ORTLocalClient": + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + onnx_bytes = onnx_bytes_obj.onnx_model[f"{onnx_bytes_obj.model_name}.onnx"] + + compiled_model = client.ir_to_compiled(onnx_bytes, compilation_args) + + return DeviceModel(client, compiled_model, metadata) diff --git a/modelopt/torch/_deploy/device_model.py b/modelopt/torch/_deploy/device_model.py new file mode 100644 index 0000000000..f24cad91f3 --- /dev/null +++ b/modelopt/torch/_deploy/device_model.py @@ -0,0 +1,157 @@ +# 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. + +"""Class representing the compiled model for a particular device.""" + +import os +from typing import Any + +from modelopt.torch.utils import unflatten_tree + +from ._runtime import DetailedResults, RuntimeClient +from .utils import ModelMetadata, generate_onnx_input + + +class DeviceModel: + """On-device model with profiling functions and PyTorch-like inference interface. + + This object should be generated from + :meth:`compile `. + """ + + def __init__( + self, + client: RuntimeClient, + compiled_model: bytes, + metadata: ModelMetadata, + compilation_args: dict[str, Any] = {}, + io_shapes: dict[str, list] = {}, + ignore_nesting: bool = False, + ): + """Initialize a device model with the corresponding model, onnx, and engine model. + + Args: + client: the runtime client used to compile the model. + compiled_model: Compiled device model created during runtime compilation. + metadata: The model's metadata (needed for inference/profiling with compiled model). + compilation_args: Additional arguments for the compilation process. + io_shapes: Defines the input and output shapes of the model. These shapes must be provided by the user + when TensorRT cannot automatically infer them, such as in cases where dynamic dimensions are used for + model outputs. To specify the shapes, use the following format: + io_shapes = {"out.0": [2, 4, 64, 64]} + ignore_nesting: If True, only the last part of the nested input name will be considered during inference. + eg. if the input name is x.y.z, only z will be considered. + """ + self.client = client + self.compiled_model = compiled_model + self.model_metadata = metadata + self.config = self.model_metadata.get("config", {}) + self.compilation_args = compilation_args + self.io_shapes = io_shapes + self.ignore_nesting = ignore_nesting + + def __call__(self, *args, **kwargs): + """Execute forward function of the model on the specified deployment and return output.""" + return self.forward(*args, **kwargs) + + def get_latency(self) -> float: + """Profiling API to let user get model latency with the compiled device model. + + Returns: + The latency of the compiled model in ms. + """ + latency, _ = self._profile_device() + return latency + + def profile(self, verbose: bool = False) -> tuple[float, DetailedResults]: + """Inference API to let user do inference with the compiled device model. + + Args: + verbose: If True, print out the profiling results as a table. + + Returns: A tuple (latency, detailed_result) where + ``latency`` is the latency of the compiled model in ms, + ``detailed_result`` is a dictionary containing additional benchmarking results + """ + latency, detailed_result = self._profile_device() + + if verbose: + print(detailed_result) + + return latency, detailed_result + + def forward(self, *args, **kwargs) -> Any: + """Execute forward of the model on the specified deployment and return output. + + Arguments: + args: Non-keyword arguments to the model for inference. + kwargs: Keyword arguments to the model for inference. + + Returns: + The inference result in the same (nested) data structure as the original model. + + .. note:: + + This API let the users do inference with the compiled device model. + + .. warning:: + + All return values will be of type ``torch.Tensor`` even if the original model returned + native python types such as bool/int/float. + """ + if self.compiled_model is None: + raise AttributeError("Please compile the model first.") + + # Flatten all args, kwargs into a single list of tensors for onnx/device inference. + all_args = args + (kwargs,) if kwargs or (args and isinstance(args[-1], dict)) else args + + # If Model metadata is None then DeviceModel is instantiated with raw ONNX bytes instead of PyTorch module. + onnx_inputs = all_args[0] + if self.model_metadata: + onnx_inputs = list( + generate_onnx_input(self.model_metadata, all_args, self.ignore_nesting).values() + ) + + # run inference with the engine equivalent of the model + onnx_outputs = self.client.inference( + compiled_model=self.compiled_model, inputs=onnx_inputs, io_shapes=self.io_shapes + ) + + # Note that bool/float/ints will be returned as corresponding tensors + # TODO: maybe eventually we want to compare this against the original types + # generate expected returned data structure of the model + if not self.model_metadata: + return onnx_outputs + return unflatten_tree(onnx_outputs, self.model_metadata["output_tree_spec"]) + + def save_compile_model(self, path: str, remove_hash: bool = False): + """Saves the compiled model to a file. + + Args: + path: The path to save the compiled model. + remove_hash: If True, the hash will be removed from the saved model. + """ + compiled_model = self.compiled_model + if remove_hash: + compiled_model = compiled_model[32:] + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "wb") as f: + f.write(compiled_model) + + def _profile_device(self) -> tuple[float, DetailedResults]: + """Profile the device model stored in self and return latency results.""" + return self.client.profile( + compiled_model=self.compiled_model, compilation_args=self.compilation_args + ) diff --git a/modelopt/torch/_deploy/profiling.py b/modelopt/torch/_deploy/profiling.py new file mode 100644 index 0000000000..247396f296 --- /dev/null +++ b/modelopt/torch/_deploy/profiling.py @@ -0,0 +1,195 @@ +# 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. + +"""Module to profile a model on a target device.""" + +from typing import Any, Union + +import torch.nn as nn + +from ._runtime import Deployment, DetailedResults +from .compilation import compile + +__all__ = ["get_latency", "profile"] + + +def get_latency( + model: nn.Module, + dummy_input: Union[Any, tuple], + deployment: Deployment, +) -> float: + """Obtain deployment latency of model by compiling and sending model to engine for profiling. + + Args: + model: PyTorch model to compile for target device. + dummy_input: Arguments of ``model.forward()``. This is used for exporting and calculating + inference-based metrics, such as latency/FLOPs. The format of ``dummy_inputs`` follows + the convention of the ``args`` argument in + `torch.onnx.export `_. + Specifically, ``dummy_input`` can be: + + #. a single argument (``type(dummy_input) != tuple``) corresponding to + + .. code-block:: python + + model.forward(dummy_input) + + #. a tuple of arguments corresponding to + + .. code-block:: python + + model.forward(*dummy_input) + + #. a tuple of arguments such that ``type(dummy_input[-1]) == dict`` corresponding to + + .. code-block:: python + + model.forward(*dummy_input[:-1], **dummy_input[-1]) + + .. warning:: + + In this case the model's ``forward()`` method **cannot** contain keyword-only + arguments (e.g. ``forward(..., *, kw_only_args)``) or variable keyword arguments + (e.g. ``forward(..., **kwargs)``) since these cannot be sorted into positional + arguments. + + .. note:: + + In order to pass a dict as last non-keyword argument, you need to use a tuple as + ``dummy_input`` and add an *empty* dict as the last element, e.g., + + .. code-block:: python + + dummy_input = (x, {"y": y, "z": z}, {}) + + The empty dict at the end will then be interpreted as the keyword args. + + See `torch.onnx.export `_ + for more info. + + Note that if you provide a ``{arg_name}`` with batch size ``b``, the results will be + computed based on batch size ``b``. + + deployment: Deployment configuration with keys as specified below. + + * ``runtime``: the desired runtime for deployment (*required*); + * ``accelerator``: the accelerator on the device to be used (*optional*); + * ``version``: the version of runtime to be used (*optional*); + * ``precision``: the desired precision (*optional*); + * ``onnx_opset``: the opset version (*optional*). + + An example of a deployment configuration is: + + .. code-block:: python + + deployment = { + "runtime": "ORT", + "accelerator": "CPU", + "version": "1.11", + "precision": "fp32", + "onnx_opset": 13, + } + + Returns: + The latency of the compiled model in ms. + """ + device_model = compile(model, dummy_input, deployment) + return device_model.get_latency() + + +def profile( + model: nn.Module, + dummy_input: Union[Any, tuple], + deployment: Deployment, + verbose: bool = False, +) -> tuple[float, DetailedResults]: + """Model profile method to help user to profile their model on target device. + + Args: + model: PyTorch model to compile for target device. + dummy_input: Arguments of ``model.forward()``. This is used for exporting and calculating + inference-based metrics, such as latency/FLOPs. The format of ``dummy_inputs`` follows + the convention of the ``args`` argument in + `torch.onnx.export `_. + Specifically, ``dummy_input`` can be: + + #. a single argument (``type(dummy_input) != tuple``) corresponding to + + .. code-block:: python + + model.forward(dummy_input) + + #. a tuple of arguments corresponding to + + .. code-block:: python + + model.forward(*dummy_input) + + #. a tuple of arguments such that ``type(dummy_input[-1]) == dict`` corresponding to + + .. code-block:: python + + model.forward(*dummy_input[:-1], **dummy_input[-1]) + + .. warning:: + + In this case the model's ``forward()`` method **cannot** contain keyword-only + arguments (e.g. ``forward(..., *, kw_only_args)``) or variable keyword arguments + (e.g. ``forward(..., **kwargs)``) since these cannot be sorted into positional + arguments. + + .. note:: + + In order to pass a dict as last non-keyword argument, you need to use a tuple as + ``dummy_input`` and add an *empty* dict as the last element, e.g., + + .. code-block:: python + + dummy_input = (x, {"y": y, "z": z}, {}) + + The empty dict at the end will then be interpreted as the keyword args. + + See `torch.onnx.export `_ + for more info. + + Note that if you provide a ``{arg_name}`` with batch size ``b``, the results will be + computed based on batch size ``b``. + deployment: Deployment configuration with keys as specified below. + + * ``runtime``: the desired runtime for deployment (*required*); + * ``accelerator``: the accelerator on the device to be used (*optional*); + * ``version``: the version of runtime to be used (*optional*); + * ``precision``: the desired precision (*optional*); + * ``onnx_opset``: the opset version (*optional*). + + An example of a deployment configuration is: + + .. code-block:: python + + deployment = { + "runtime": "ORT", + "accelerator": "CPU", + "version": "1.11", + "precision": "fp32", + "onnx_opset": 13, + } + verbose: If True, print out the profiling results as a table. + + Returns: A tuple (latency, detailed_result) where + ``latency`` is the latency of the compiled model in ms, + ``detailed_result`` is a dictionary containing additional benchmarking results + """ + device_model = compile(model, dummy_input, deployment) + return device_model.profile(verbose) diff --git a/modelopt/torch/_deploy/utils/__init__.py b/modelopt/torch/_deploy/utils/__init__.py new file mode 100644 index 0000000000..4840ec0cc2 --- /dev/null +++ b/modelopt/torch/_deploy/utils/__init__.py @@ -0,0 +1,18 @@ +# 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. + +"""Utilities for modelopt.torch._deploy.""" + +from .torch_onnx import * diff --git a/modelopt/torch/_deploy/utils/onnx_optimizer.py b/modelopt/torch/_deploy/utils/onnx_optimizer.py new file mode 100644 index 0000000000..fba8306716 --- /dev/null +++ b/modelopt/torch/_deploy/utils/onnx_optimizer.py @@ -0,0 +1,118 @@ +# 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. + +"""Utility to optimze onnx graphs.""" + +import os +import tempfile + +import onnx +import onnx_graphsurgeon as gs + + +class Optimizer: + """Optimizer for onnx graphs.""" + + def __init__(self, onnx_graph, verbose=False): + """Initializes optimizer.""" + self.graph = gs.import_onnx(onnx_graph) + self.verbose = verbose + + def info(self, prefix): + """Prints the graph information.""" + if self.verbose: + print( + f"{prefix} .. {len(self.graph.nodes)} nodes," + f" {len(self.graph.tensors().keys())} tensors," + f" {len(self.graph.inputs)} inputs, {len(self.graph.outputs)} outputs" + ) + + def cleanup(self, return_onnx=False): + """Cleans the onnx graph. + + Args: + return_onnx (bool): If True, returns the onnx graph. + + Returns: + onnx_graph: The cleaned onnx graph. + """ + self.graph.cleanup().toposort() + if return_onnx: + return gs.export_onnx(self.graph) + + def select_outputs(self, keep, names=None): + """Selects the output nodes.""" + self.graph.outputs = [self.graph.outputs[o] for o in keep] + if names: + for i, name in enumerate(names): + self.graph.outputs[i].name = name + + def infer_shapes(self, return_onnx=False): + """Infers shapes of the onnx graph.""" + onnx_graph = gs.export_onnx(self.graph) + if onnx_graph.ByteSize() > 2147483648: + temp_dir = tempfile.TemporaryDirectory().name + os.makedirs(temp_dir, exist_ok=True) + onnx_orig_path = os.path.join(temp_dir, "model.onnx") + onnx_inferred_path = os.path.join(temp_dir, "inferred.onnx") + onnx.save_model( + onnx_graph, + onnx_orig_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + convert_attribute=False, + ) + onnx.shape_inference.infer_shapes_path(onnx_orig_path, onnx_inferred_path) + onnx_graph = onnx.load(onnx_inferred_path) + else: + onnx_graph = onnx.shape_inference.infer_shapes(onnx_graph) + + if return_onnx: + return onnx_graph + + # TODO: Move this fuctionality to the diffusion runner as it is specific to CLIP. + def clip_add_hidden_states(self, return_onnx=False): + """Adds hidden states to the CLIP model graph.""" + hidden_layers = -1 + onnx_graph = gs.export_onnx(self.graph) + for i in range(len(onnx_graph.graph.node)): + for j in range(len(onnx_graph.graph.node[i].output)): + name = onnx_graph.graph.node[i].output[j] + if "layers" in name: + hidden_layers = max(int(name.split(".")[1].split("/")[0]), hidden_layers) + for i in range(len(onnx_graph.graph.node)): + for j in range(len(onnx_graph.graph.node[i].output)): + if onnx_graph.graph.node[i].output[ + j + ] == "/text_model/encoder/layers.{}/Add_1_output_0".format(hidden_layers - 1): + onnx_graph.graph.node[i].output[j] = "hidden_states" + for j in range(len(onnx_graph.graph.node[i].input)): + if onnx_graph.graph.node[i].input[ + j + ] == "/text_model/encoder/layers.{}/Add_1_output_0".format(hidden_layers - 1): + onnx_graph.graph.node[i].input[j] = "hidden_states" + if return_onnx: + return onnx_graph + + def fold_constants(self, return_onnx=False): + """Folds constants in the onnx graph for multiple iterations till prefolded nodes == post folded nodes.""" + prefold_num_nodes = len(self.graph.nodes) + post_fold_num_nodes = -1 + while prefold_num_nodes != post_fold_num_nodes: + prefold_num_nodes = len(self.graph.nodes) + self.graph.fold_constants() + post_fold_num_nodes = len(self.graph.nodes) + if return_onnx: + return gs.export_onnx(self.graph) diff --git a/modelopt/torch/_deploy/utils/onnx_utils.py b/modelopt/torch/_deploy/utils/onnx_utils.py new file mode 100644 index 0000000000..a377afcb6b --- /dev/null +++ b/modelopt/torch/_deploy/utils/onnx_utils.py @@ -0,0 +1,58 @@ +# Adapted from: https://github.com/huggingface/optimum/blob/0d808ad/optimum/onnx/utils.py + +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# 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. + +# 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 onnx +from onnx.external_data_helper import ExternalDataInfo, _get_initializer_tensors + + +def _get_onnx_external_data_tensors(model: onnx.ModelProto) -> list[str]: + """ + Gets the paths of the external data tensors in the model. + Note: make sure you load the model with load_external_data=False. + """ + model_tensors = _get_initializer_tensors(model) + model_tensors_ext = [ + ExternalDataInfo(tensor).location + for tensor in model_tensors + if tensor.HasField("data_location") and tensor.data_location == onnx.TensorProto.EXTERNAL + ] + return model_tensors_ext + + +def check_model_uses_external_data(model: onnx.ModelProto) -> bool: + """ + Checks if the model uses external data. + """ + model_tensors = _get_initializer_tensors(model) + return any( + tensor.HasField("data_location") and tensor.data_location == onnx.TensorProto.EXTERNAL + for tensor in model_tensors + ) diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py new file mode 100644 index 0000000000..13b6d3c7c6 --- /dev/null +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -0,0 +1,561 @@ +# 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. + +"""Utility functions related to Onnx.""" + +import inspect +import os +import shutil +from typing import Any, Union + +import cloudpickle +import onnx +import torch +import torch.nn as nn +from onnx import ModelProto +from torch.nn.parallel import DataParallel, DistributedDataParallel + +from modelopt.onnx.quantization.qdq_utils import qdq_to_dq +from modelopt.onnx.utils import ( + get_input_names, + get_input_shapes, + get_node_names, + get_output_names, + get_output_shapes, +) +from modelopt.torch._deploy._runtime.common import write_bytes +from modelopt.torch._deploy.utils.onnx_optimizer import Optimizer +from modelopt.torch.utils import flatten_tree, standardize_named_model_args +from modelopt.torch.utils._pytree import TreeSpec + +from .onnx_utils import _get_onnx_external_data_tensors, check_model_uses_external_data + +ModelMetadata = dict[str, Any] +ModelType = Any +ValueInfoType = Any + +# a few constants... +DEFAULT_ONNX_OPSET = 13 +ONNX_EXPORT_OUT_PREFIX = "out" +TWO_GB = 2147483648 + + +class OnnxBytes: + """A class to save and load onnx models as bytes.""" + + def __init__(self, onnx_load_path: str) -> None: + """Loads the model from the specified path. + + If the model is loaded without external data format, then it is saved as a dictionary where + the key is the model name and the value is the model bytes. + If the model is loaded with external data format, then the model is saved as a dictionary + where the keys include all the file names in the model directory and the value are the corresponding file bytes. + For external data format, we assume that the external data for the model is saved in the same directory + as the model file. + + Args: + onnx_load_path: The path to load the .onnx model file. + """ + self.onnx_load_path = os.path.abspath(onnx_load_path) + self.onnx_model = {} + self.model_name = "" + print("Loading onnx model from path:", self.onnx_load_path) + onnx_model = onnx.load(self.onnx_load_path, load_external_data=False) + + # Check for external data + external_data_format = False + for initializer in onnx_model.graph.initializer: + if initializer.external_data: + external_data_format = True + + if external_data_format: + onnx_model_dir = os.path.dirname(self.onnx_load_path) + for onnx_model_file in os.listdir(onnx_model_dir): + with open(os.path.join(onnx_model_dir, onnx_model_file), "rb") as f: + self.onnx_model[onnx_model_file] = f.read() + if onnx_model_file.endswith(".onnx"): + if self.model_name != "": + raise ValueError("Multiple onnx files found in the directory") + self.model_name = onnx_model_file.replace(".onnx", "") + else: + onnx_model_file = os.path.basename(self.onnx_load_path) + if not onnx_model_file.endswith(".onnx"): + raise ValueError("The file should be a .onnx file") + with open(self.onnx_load_path, "rb") as f: + self.onnx_model[onnx_model_file] = f.read() + self.model_name = onnx_model_file.replace(".onnx", "") + + def write_to_disk(self, onnx_save_path: str) -> None: + """Writes the onnx model to the specified path.""" + if os.path.exists(onnx_save_path): + print(f"Removing existing directory: {onnx_save_path}") + shutil.rmtree(onnx_save_path) + os.makedirs(onnx_save_path) + print("Writing onnx model to path:", onnx_save_path) + for onnx_model_file, onnx_model_bytes in self.onnx_model.items(): + with open(os.path.join(onnx_save_path, onnx_model_file), "wb") as f: + f.write(onnx_model_bytes) + + def to_bytes(self) -> bytes: + """Returns the bytes of the object.""" + return cloudpickle.dumps(self) + + def get_onnx_model_file_bytes(self) -> bytes: + """Returns the bytes of the onnx model file.""" + return self.onnx_model[self.model_name + ".onnx"] + + @classmethod + def from_bytes(cls, onnx_bytes: bytes) -> "OnnxBytes": + """Returns the OnnxBytes object from the bytes.""" + return cloudpickle.loads(onnx_bytes) + + +def _to_expected_onnx_type(val: Any) -> Any: + """Convert the given value to the expected onnx type. + + During the onnx export process, plain numeric types (floats and ints) are converted to torch + tensors. This function pre-converts the given val to a tensor in case val is a int or float for + easier handling of such input values during the onnx export process. + """ + if isinstance(val, (int, float)): + return torch.tensor(val).to(type(val)) + return val + + +def generate_onnx_input( + model_metadata: ModelMetadata, input: Union[Any, tuple], ignore_nesting: bool = False +) -> dict[str, Any]: + """Generate input for onnx model from model's forward signature and provided input. + + Args: + model_metadata: The model's metadata. + input: A tuple of args/kwargs or torch.Tensor feed into the model's ``forward()`` method, + see :meth:`standardize_model_args() ` + for more info on the convention. + ignore_nesting: If True, only the last part of the nested input name will be considered. + eg. if the input name is x.y.z, only z will be considered. + + Returns: + Args flattened into one dictionary with serialized keys compatible with provided onnx. + + .. note:: + + This function performs a sanity check on the provided input data to filter out args that + are constants (instead of input nodes) in the onnx graph. + + + Some more relevant background of why we want to flatten the input pytree here: + + * In the onnx export process, nested python data structures (like nested lists, tuples, + dictionaries) are being recursed into until leaf objects corresponding to tensors are + encountered. + + * This is used to flatten the input in an onnx model to a list of tensors. + + * However, this is a fairly complex process for the user to understand in case their models + takes a nested data structure. They have to understand how to manually flatten the data + structure in the *correct* order in order for them to run inference on a device_model or + onnx model. + + * With this function this additional complexity can be abstracted away from the user. + + * Example: if the original model took ``[x, {"y":y, "z" : [z1,z2]}]`` they can still provide + this nested data structure instead of the expected onnx input list of ``[x, y, z1, z2]`` + --> flattening and unflattering is handled internally. + """ + # get named args and set of params where we added default values + named_args, args_with_default = standardize_named_model_args(model_metadata["signature"], input) + + # retrieve onnx input names + onnx_input_names = model_metadata["input_onnx_names"] + input_none_names = model_metadata["input_none_names"] + + # capture flattened names of args from default values + named_default_args = {k: v for k, v in named_args.items() if k in args_with_default} + _, tree_spec_default_args = flatten_tree(named_default_args) + + # capture flattened args without default args that do not appear in onnx graph + values, tree_spec = flatten_tree(named_args) + if not ignore_nesting: + flat_kv = {k: v for k, v in zip(tree_spec.names, values)} + else: + flat_kv = {k.split(".")[-1]: v for k, v in zip(tree_spec.names, values)} + + # We wanna consider four types of flattened args: + # 1. Args that appear in the onnx graph + # 2. Args that are not their default value + # 3. Args that were provided as None during conversion but are not None right now + # 4. Args that were provided as None during conversion and are None right now + + args_in_onnx = {k for k in flat_kv if k in onnx_input_names} + args_not_default = {k for k in flat_kv if k not in tree_spec_default_args.names} + args_not_none = {k for k, v in flat_kv.items() if k in input_none_names and v is not None} + args_none = {k for k, v in flat_kv.items() if k in input_none_names and v is None} + + # identify unexpected args from these 4 types + unexpected_args = ((args_not_default - args_none) | args_not_none) - args_in_onnx + if unexpected_args: + print( + "The following args were provided that do not appear in the onnx graph of your model " + "since they are treated as constants in the onnx graph:" + + "\t\n".join(unexpected_args) + + "\nConsider removing these args from your input that are constants in the onnx model " + "or recompiling your onnx model with new constant values!" + ) + + # return the args that are relevant for the onnx graph in the right type + return {k: _to_expected_onnx_type(v) for k, v in flat_kv.items() if k in args_in_onnx} + + +def optimize(name, onnx_graph, verbose=False): + """Optimizes onnx graph.""" + opt = Optimizer(onnx_graph, verbose=verbose) + opt.info(name + ": original") + opt.cleanup() + opt.info(name + ": cleanup") + # TODO: fold constants is not working for some models from deploy_models(NestedOutModel, ArgsKwargsModel1) + # opt.fold_constants() + # opt.info(name + ": fold_constants") + onnx_graph = opt.infer_shapes(return_onnx=True) + opt.info(name + ": shape inference") + return onnx_graph + + +def split_args_kwargs(args_tuple): + """Splits args_tuple into positional arguments and keyword arguments.""" + split_index = len(args_tuple) + + for i, item in enumerate(reversed(args_tuple)): + if not isinstance(item, dict): + split_index = len(args_tuple) - i + break + + pos_args = args_tuple[:split_index] + kw_args = {} + for d in args_tuple[split_index:]: + kw_args.update(d) + + kw_args = None if kw_args == {} else kw_args + + # remove empty dict if it is the last element + if pos_args[-1] == {}: + pos_args = pos_args[:-1] + + return pos_args, kw_args + + +def get_onnx_bytes_and_metadata( + model: nn.Module, + dummy_input: Union[Any, tuple], + onnx_load_path: str = "", + dynamic_axes: dict = {}, + remove_exported_model: bool = True, + dynamo_export: bool = False, + onnx_opset: int = DEFAULT_ONNX_OPSET, + dq_only: bool = False, +) -> tuple[bytes, ModelMetadata]: + """Get onnx model in bytes from input pytorch model together with the input/output of model. + + Arguments: + model: PyTorch model to export to onnx. + dummy_input: A tuple of args/kwargs or torch.Tensor, see + `torch.onnx.export `_ + for more info on the convention. + onnx_load_path: The path to load the onnx model. + dynamic_axes: A dictionary of dynamic shapes used for exporting the torch model to onnx. + remove_exported_model: If True, the onnx model will be cleared from the disk after the + export process + dynamo_export: If True, the model is exported using dynamo_export. See + `torch.onnx.dynamo_export https://pytorch.org/docs/stable/onnx_dynamo.html` for more details. + onnx_opset: The onnx opset version to use for exporting the model. + dq_only: If True, the exported ONNX model is converted to a dq_only model. + + Returns: + bytes: Onnx model in bytes. + ModelMetadata: The model's meta data. + + Raises: + ValueError: If nn.Module is not passed as model. + """ + if not isinstance(model, nn.Module): + raise ValueError("Only PyTorch model compilation is supported.") + + # unwrap DDP and DP models + if isinstance(model, (DataParallel, DistributedDataParallel)): + model = model.module + + # Standardize model args and also tensorize them so they also appear in the onnx graph! + # Floats/ints are tensorized when they are provided, but not tensorized when they are not + # provided which is somewhat inconsistent (we always tensorize them!) + named_args, _ = standardize_named_model_args(model, dummy_input) + named_args = {k: _to_expected_onnx_type(v) for k, v in named_args.items()} + + # Also standardize dummy_input again so we can use it + dummy_input = tuple(named_args.values()) + if dummy_input and isinstance(dummy_input[-1], dict): + dummy_input = dummy_input + ({},) # we need to add an extra dict for the fake kwargs! + + dummy_input_args, dummy_input_kwargs = (), None + if dynamo_export: + dummy_input_args, dummy_input_kwargs = split_args_kwargs(dummy_input) + + # Get input tree spec, see generate_onnx_input for more info as well on this + flat_input, tree_spec_input = flatten_tree(named_args) + + # input names are the names of the flattened input tree spec but without None values + input_names = [k for k, v in zip(tree_spec_input.names, flat_input) if v is not None] + + # we also want to record the input names that are None so we can remove them from the input + # during inference. + input_none_names = list(set(tree_spec_input.names) - set(input_names)) + + # Get output once (we export in inference mode - so also using inference mode here!) + with torch.inference_mode(): + # We enable mixed precision for cases where some outputs may have a different data type than others + with torch.autocast("cuda"): + output = model(*named_args.values()) + + # Get output tree spec + flat_output, tree_spec_output = flatten_tree(output, prefix=ONNX_EXPORT_OUT_PREFIX) + + # output names are the names of the flattened input tree spec but without None values + output_names = [k for k, v in zip(tree_spec_output.names, flat_output) if v is not None] + + model_name = model.__class__.__name__ + onnx_build_folder = os.path.abspath("./build/onnx/") + onnx_path = os.path.join(onnx_build_folder, model_name) + os.makedirs(onnx_path, exist_ok=True) + onnx_save_path = os.path.join(onnx_path, f"{model_name}.onnx") + + # If the onnx_load path is specified by the user or if an onnx model exists at the default path + # then the model is loaded from this path and returned along with the metadata + if os.path.exists(onnx_save_path) and onnx_load_path == "": + print(f"Overriding onnx load path to {onnx_save_path}") + onnx_load_path = onnx_save_path + + if onnx_load_path != "": + onnx_model = OnnxBytes(onnx_load_path) + onnx_model_graph = onnx.load(os.path.join(onnx_load_path)) + model_metadata = create_model_metadata( + tree_spec_input, tree_spec_output, input_none_names, onnx_model_graph, model + ) + return onnx_model.to_bytes(), model_metadata + + # Export onnx model from pytorch model + # As the maximum size of protobuf is 2GB, we cannot use io.BytesIO() buffer during export. + with torch.inference_mode(): + if not dynamo_export: + torch.onnx.export( + model, + dummy_input, + onnx_save_path, + input_names=input_names, + output_names=output_names, + opset_version=onnx_opset, + dynamic_axes=dynamic_axes, + ) + else: + export_options = torch.onnx.ExportOptions(dynamic_shapes=True) + if dummy_input_kwargs: + torch.onnx.dynamo_export( + model, *dummy_input_args, export_options=export_options, **dummy_input_kwargs + ).save(onnx_save_path) + else: + torch.onnx.dynamo_export( + model, + *dummy_input_args, + export_options=export_options, + ).save(onnx_save_path) + # Check that export worked + assert len(os.listdir(onnx_path)) > 0, "Torch to onnx export failed." + + # Load the onnx graph for optimizaiton + onnx_graph = onnx.load(onnx_save_path, load_external_data=True) + + # Optimize the onnx graph + onnx_opt_graph = optimize(model.__class__.__name__, onnx_graph) + + # Remove training_mode attribute from BatchNormalization nodes + onnx_opt_graph = remove_node_training_mode(onnx_opt_graph, "BatchNormalization") + + model_metadata = create_model_metadata( + tree_spec_input, tree_spec_output, input_none_names, onnx_opt_graph, model + ) + + # Change the ONNX IR version to 9 to be compatible with ONNXRuntime + onnx_opt_graph.ir_version = 9 + + if dq_only: + onnx_opt_graph = qdq_to_dq(onnx_opt_graph) + + # If the onnx model contains external data store the external tensors in one file and save the onnx model + if has_external_data(onnx_save_path): + tensor_paths = _get_onnx_external_data_tensors(onnx_opt_graph) + onnx.save_model( + onnx_opt_graph, + onnx_save_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=f"{model_name}.onnx_data", + size_threshold=1024, + ) + for tensor in tensor_paths: + tensor_path = os.path.join(onnx_path, tensor) + os.remove(tensor_path) + else: + onnx.save_model(onnx_opt_graph, onnx_save_path) + + onnx_bytes = OnnxBytes(onnx_save_path) + + if remove_exported_model: + shutil.rmtree(os.path.dirname(onnx_build_folder)) + return onnx_bytes.to_bytes(), model_metadata + + +def has_external_data(onnx_model_path: str): + """Check if the onnx model has external data.""" + onnx_model = onnx.load(onnx_model_path, load_external_data=False) + return check_model_uses_external_data(onnx_model) + + +def create_onnx_model_dict( + onnx_load_path: str = None, model_metadata: ModelMetadata = None, model_dict: dict = None +) -> tuple[bytes, ModelMetadata]: + """Creates onnx model dictionary and model metadata. + + Args: + onnx_load_path: The path to load the onnx model. + model_metadata: Metadata to be stored in the DeviceModel. + model_dict: A dictionaty containing the Pytorch model and the input and output treespec for the model. + + Returns: + A tuple of onnx model dictionary and model metadata. + """ + onnx_model_dict = {"onnx_load_path": onnx_load_path} + onnx_model = {} + for onnx_model_file in os.listdir(onnx_load_path): + with open(os.path.join(onnx_load_path, onnx_model_file), "rb") as f: + onnx_model[onnx_model_file] = f.read() + if onnx_model_file.endswith(".onnx") and model_metadata is None: + onnx_model_file_path = os.path.join(onnx_load_path, onnx_model_file) + onnx_graph = onnx.load(onnx_model_file_path) + onnx_input_names = [input.name for input in onnx_graph.graph.input] + torch_input_names = model_dict["tree_spec_input"].names + assert set(onnx_input_names).issubset( + set(torch_input_names) + ), "One or more inputs in the onnx model are not present in the torch model." + onnx_model_dict["onnx_model"] = onnx_model + if model_metadata is None: + model_metadata = create_model_metadata( + model_dict["tree_spec_input"], + model_dict["tree_spec_output"], + model_dict["input_none_names"], + onnx_graph, + model_dict["model"], + ) + return cloudpickle.dumps(onnx_model_dict), model_metadata + + +def create_model_metadata( + tree_spec_input: TreeSpec, + tree_spec_output: TreeSpec, + input_none_names: list[str], + onnx_graph: ModelProto, + model: nn.Module, +) -> ModelMetadata: + """Create model metadata from the given input. + + Args: + tree_spec_input: pytree spec describing the structure of the pytree for the model input. + tree_spec_output: pytree spec describing the structure of the pytree for the model output. + input_none_names: List of input names with values that are None. + onnx_opt_graph: Graph of the onnx model. + model: Pytorch model. + + Returns: + ModelMetadata: The DeviceModel metadata. + """ + return { + "input_tree_spec": tree_spec_input, + "input_shapes": get_input_shapes(onnx_graph), + "input_onnx_names": get_input_names(onnx_graph), + "input_none_names": input_none_names, + "output_tree_spec": tree_spec_output, + "output_shapes": get_output_shapes(onnx_graph), + "output_onnx_names": get_output_names(onnx_graph), + "signature": inspect.signature(model.forward), + "onnx_node_names": get_node_names(onnx_graph), + "is_bytes_pickled": onnx_graph.ByteSize() > TWO_GB, + "config": model.config if hasattr(model, "config") else None, + } + + +def write_onnx_bytes(onnx_model_bytes: Union[bytes, dict], onnx_save_path: str) -> tuple[str, str]: + """Write onnx bytes to the specified path. + + Args: + onnx_model_bytes: Can be one of the following: + 1. A pickled dictionary of onnx model files. The keys of this dictionary are file names and\ + the values are the corresponding bytes. + 2. A single onnx model bytes. + onnx_save_path: The onnx path to save the model. + + Returns: + The path to the onnx model. + Name of the onnx model. + """ + if isinstance(onnx_model_bytes, dict): + if not os.path.exists(onnx_save_path): + os.makedirs(onnx_save_path) + for key in onnx_model_bytes["onnx_model"]: + write_bytes(onnx_model_bytes["onnx_model"][key], os.path.join(onnx_save_path, key)) + onnx_model_file = [s for s in onnx_model_bytes["onnx_model"].keys() if s.endswith(".onnx")] + assert len(onnx_model_file) == 1, "There should be only one onnx model file" + model_name = onnx_model_file[0].split(".")[0] + return os.path.join(onnx_save_path, onnx_model_file[0]), model_name + write_bytes(onnx_model_bytes, os.path.join(onnx_save_path, "model.onnx")) + return os.path.join(onnx_save_path, "model.onnx"), "" + + +def get_onnx_bytes(*args, **kwargs): + """Return onnx bytes only. + + See ``get_onnx_bytes_and_metadata()`` for more info. + """ + onnx_bytes = get_onnx_bytes_and_metadata(*args, **kwargs)[0] + onnx_bytes_obj = OnnxBytes.from_bytes(onnx_bytes) + return onnx_bytes_obj.onnx_model[f"{onnx_bytes_obj.model_name}.onnx"] + + +def remove_node_training_mode(onnx_model: ModelProto, node_op_type: str) -> ModelProto: + """Remove training_mode attribute from selected node type. + + Args: + onnx_model: The onnx model. + node_op_type: The node type to remove training_mode attribute from. + + Returns: + The onnx model with the training_mode attribute removed. + """ + for node in onnx_model.graph.node: + if node.op_type == node_op_type: + for attribute in node.attribute: + if attribute.name == "training_mode": + if attribute.i == 1: + node.output.remove(node.output[1]) + node.output.remove(node.output[1]) + attribute.i = 0 + + return onnx_model diff --git a/modelopt/torch/distill/__init__.py b/modelopt/torch/distill/__init__.py new file mode 100644 index 0000000000..a09aa6b8ef --- /dev/null +++ b/modelopt/torch/distill/__init__.py @@ -0,0 +1,28 @@ +# 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. + +"""Distillation API subpackage for torch.""" + +from . import mode +from .config import * +from .distillation import * +from .distillation_model import * +from .loss_balancers import * +from .losses import * +from .registry import * + +# isort: off +# Import plugins last to avoid circular imports +from . import plugins diff --git a/modelopt/torch/distill/config.py b/modelopt/torch/distill/config.py new file mode 100644 index 0000000000..79c70b80ec --- /dev/null +++ b/modelopt/torch/distill/config.py @@ -0,0 +1,130 @@ +# 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. + +"""Configurations for distillation modes.""" + +import warnings +from typing import Any, Callable, Optional, Union + +import pydantic +import torch.nn as nn +from torch.nn.modules.loss import _Loss as Loss + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + +from .loss_balancers import DistillationLossBalancer + +__all__ = ["KDLossConfig"] + +TeacherModel = Union[type[nn.Module], tuple, Callable] +Criterion = Union[Loss, dict[tuple[str, str], Loss]] + + +class KDLossConfig(ModeloptBaseConfig): + """Configuration for the Knowledge-Distillation mode. + + This mode is used to distill knowledge from a teacher model to a student model. + """ + + # TODO: we should really think about a better to configure KDLossConfig + model_config = pydantic.ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + teacher_model: Optional[TeacherModel] = ModeloptField( + default=None, + title="Teacher model", + description=( + "The class or callable or tuple to initialize the teacher model using" + " :meth:`init_model_from_model_like" + " `. This cannot already be an" + " instance of nn.Module." + ), + ) + criterion: Optional[Criterion] = ModeloptField( + default=None, + title="Output-only or layer-wise loss criterion", + description=( + "If an instance of Loss class, a distillation loss will only be computed " + "between outputs of a student and teacher; if a dictionary in the format " + "{(student_layer_name, teacher_layer_name): loss_module}, a distillation loss will be " + "computed for each specified student-teacher pair of layers using the corresponding " + "``loss_module``." + ), + ) + loss_balancer: Optional[Any] = ModeloptField( + default=None, + title="Loss balancer", + description=( + "A balancer to reduce distillation and non-distillation losses into a single " + "value using some weighing scheme." + ), + ) + expose_minimal_state_dict: bool = ModeloptField( + default=True, + title="Expose student state dict only", + description=( + "Hide teacher model's state_dict in the returned wrapped model. This reduces the " + "checkpoint size by not re-storing the teacher unnecessarily again. " + ".. note: Set to False if using `FSDP `_" + ), + ) + + @pydantic.field_validator("criterion") + @classmethod + def format_criterion(cls, criterion: Optional[Criterion]) -> dict[tuple[str, str], Loss]: + """Ensure criterion is a mapping from layer names to loss (potentially entire module).""" + if not isinstance(criterion, dict): + # Output-only distillation. + criterion = {("", ""): criterion} + return criterion + + def model_dump(self, *args, **kwargs) -> dict[str, Any]: + """Dump the config to a dictionary but avoid serializing teacher model to dict. + + This avoids issues when the teacher is a tuple with callable and args. + If any of the args are Dataclasses, they are dumped as a dict and cannot be restored with their class. + """ + return {**super().model_dump(*args, **kwargs), "teacher_model": self.teacher_model} + + def _strict_validate(self) -> None: + """Strictly validate the configuration.""" + for k, v in self.items(): + if k == "loss_balancer": + pydantic.TypeAdapter( + Optional[DistillationLossBalancer], config=self.model_config + ).validate_python(v) + continue + assert v is not None, f"Missing required field: {k}." + + # Cannot have multiple loss layers without LossBalancer. + if self.loss_balancer is not None: + assert ( + len(tuple(self.loss_balancer.parameters())) == 0 + ), "Loss Balancer cannot have parameters." + else: + if len(self.criterion) == 0: # type: ignore[arg-type] + warnings.warn( + "No distillation loss criterion provided. Ignore if in pipeline-parallel setup." + ) + elif len(self.criterion) > 1: # type: ignore[arg-type] + raise ValueError( + "Cannot have multiple layer-loss pairs without a `DistillationLossBalancer`" + ) + + +class ExportStudentConfig(ModeloptBaseConfig): + """Configuration for the export_student mode. + + This mode is used to export a student model after distillation training/fine-tuning. + """ diff --git a/modelopt/torch/distill/distillation.py b/modelopt/torch/distill/distillation.py new file mode 100644 index 0000000000..d1888c084f --- /dev/null +++ b/modelopt/torch/distill/distillation.py @@ -0,0 +1,84 @@ +# 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. + +"""API for converting a model into a `modelopt.torch.distill.DistillationModel` to be used directly in training.""" + +import warnings + +import torch.nn as nn + +from modelopt.torch.opt.conversion import apply_mode +from modelopt.torch.opt.mode import ModeLike +from modelopt.torch.utils import unwrap_model + +from .distillation_model import DistillationModel +from .mode import DistillModeRegistry + +__all__ = ["convert", "export"] + + +def convert(model: nn.Module, mode: ModeLike) -> nn.Module: + """Main conversion function to turn a student model into a distillation-ready model. + + Args: + model: The base model to be used as the student. + mode: A (list of) string(s) or Mode(s) or a list of tuples containing the mode and its + config indicating the desired mode(s) (and configurations) for the convert + process. Modes set up the model for different algorithms for model optimization. The + following modes are available: + + * :class:`"kd_loss"`: The + ``model`` will be converted into meta-model encapsulating both teacher and student. + The mode's config is described in + :class:`KDLossConfig`. + + If the mode argument is specified as a dictionary, the keys should indicate the mode and + the values specify the per-mode configuration. + + Returns: + An instance of :class:`DistillationModel nn.Module: + """Export a distillation meta-model to the original student model. + + Args: + model: Model to be exported out of a distillation mode into the student only. + + Returns: + The inner student model. + + """ + # Unwrap a DP/DDP model. + model = unwrap_model( + model, + warn=True, + msg=( + f"Unwrapping a {type(model).__name__} model for export! Note that the export is" + " in-place and the model wrapper should be re-created after export since the wrapper" + " might not support changing parameters after initialization." + ), + ) + + # Check if the model is a distillation model. + if not isinstance(model, DistillationModel): + warnings.warn("Export of model failed: Model is not a `DistillationModel`") + return model + + # apply export mode + return apply_mode(model, "export_student", registry=DistillModeRegistry, init_state=False) diff --git a/modelopt/torch/distill/distillation_model.py b/modelopt/torch/distill/distillation_model.py new file mode 100644 index 0000000000..ce6c8142be --- /dev/null +++ b/modelopt/torch/distill/distillation_model.py @@ -0,0 +1,317 @@ +# 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. + +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +"""Meta-model wrapper to support knowledge-distillation learning.""" + +import inspect +import warnings +from contextlib import contextmanager +from typing import Any, Callable, Optional, Union + +import torch +import torch.nn as nn +from torch.nn.modules.loss import _Loss as Loss + +from modelopt.torch.opt.dynamic import DynamicModule + +from .loss_balancers import STUDENT_LOSS_KEY, DistillationLossBalancer + +__all__ = ["DistillationModel"] + + +class DistillationModel(DynamicModule): + """Class to encapsulate multiple teacher and student models as a single model.""" + + def _setup(self): + self._register_temp_attribute("_layers_to_loss", {}) + self._register_temp_attribute("_loss_balancer", None) + self._register_temp_attribute("_expose_minimal_state_dict", None) + self._register_temp_attribute("_teacher_model", None) + self._register_temp_attribute("_loss_modules", nn.ModuleList()) + self._register_temp_attribute("_only_teacher_fwd", False) + self._register_temp_attribute("_only_student_fwd", False) + + # HACK: set model's forward signature to match student class' original. + # Needed for HF `transformers.utils.find_labels` which relies on inspecting class signature. + sig_old = inspect.signature(self.original_cls.forward) + sig_new = inspect.signature(type(self).forward) + type(self).forward.__signature__ = sig_new.replace( # type: ignore[attr-defined] + parameters=tuple(sig_old.parameters.values()), + return_annotation=sig_old.return_annotation, + ) + + def modify( + self, + teacher_model: nn.Module, # To be frozen. + criterion: dict[ + tuple[ + str, # Student model layer whose output to capture. + str, # Teacher model layer whose output to capture. + ], + Loss, # Loss fn. + ], + loss_balancer: Optional[DistillationLossBalancer] = None, + expose_minimal_state_dict: bool = True, + ): + """Constructor. + + Args: + teacher_model: A teacher model which this class would encapsulate. + criterion: A dictionary mapping the tuple of student and teacher + model layer names to the loss function to apply to that layer pair. + loss_balancer: Instance of + :class:`DistillationLossBalancer ` + which reduces distillation and non-distillation losses into a single value using some weighing scheme. + expose_minimal_state_dict: If True, will hide teacher's state dict when calling ``state_dict`` on this + class. This allows avoiding to save the teacher state unnecessarily during checkpointing. + .. note: Set to False if using `FSDP `_ + """ + self._loss_balancer = loss_balancer + self._expose_minimal_state_dict = expose_minimal_state_dict + + # Assign loss to specified modules. + self._layers_to_loss = { + ( + self.get_submodule(student_layer_name), + teacher_model.get_submodule(teacher_layer_name), + ): loss_fn + for ( + student_layer_name, + teacher_layer_name, + ), loss_fn in criterion.items() + } + + # Register all child modules not automatically registered with assignment operator. + # This is done to ensure that the parameters of all the underlying nn.Modules + # in the DistillationModel appear to the caller when querying for parameters or + # child modules of DistillationModel object. These might be needed when attaching + # model params to optimizer, saving/restoring state of DistillationModel etc. + self._teacher_model = teacher_model + self._loss_modules = nn.ModuleList( + set(m for m in self._layers_to_loss.values() if len(list(m.parameters())) > 0) + ) + + # Disable grad for teacher + self._teacher_model.requires_grad_(False) + + # Register hooks for intermediate outputs from teacher models and the student model. + # HACK: For inexplicable reasons, sometimes a model will have hooks remain after + # `ato.restore()` so we check if they are present accidentally first. + for student_layer, teacher_layer in self._layers_to_loss: + setattr(student_layer, "_intermediate_output", None) + if output_capture_fwd_hook not in student_layer._forward_hooks.values(): + student_layer.register_forward_hook(output_capture_fwd_hook) + setattr(teacher_layer, "_intermediate_output", []) + if output_capture_fwd_hook not in teacher_layer._forward_hooks.values(): + teacher_layer.register_forward_hook(output_capture_fwd_hook) + + @property + def teacher_model(self) -> nn.ModuleList: + """Fetch the teacher model.""" + return self._teacher_model + + @property + def loss_modules(self) -> nn.ModuleList: + """Fetch the loss modules list.""" + return self._loss_modules + + @property + def loss_balancer(self) -> Optional[DistillationLossBalancer]: + """Fetch the loss balancer, if any.""" + return self._loss_balancer + + @contextmanager + def hide_teacher_model(self, enable=True): + """Context manager to temporarily hide teacher model from the model.""" + teacher_model = self._teacher_model + if enable: + self._teacher_model = nn.Module() + try: + yield + finally: + self._teacher_model = teacher_model + + @contextmanager + def hide_loss_modules(self, enable=True): + """Context manager to temporarily hide teacher model from the model.""" + loss_modules = self._loss_modules + if enable: + self._loss_modules = nn.ModuleList() + try: + yield + finally: + self._loss_modules = loss_modules + + @contextmanager + def only_teacher_forward(self, enable=True): + """Context manager to temporarily disable forward passes on the student model.""" + if enable: + self._only_teacher_fwd = True + try: + yield + finally: + self._only_teacher_fwd = False + + @contextmanager + def only_student_forward(self, enable=True): + """Context manager to temporarily disable forward passes on the student model.""" + if enable: + self._only_student_fwd = True + try: + yield + finally: + self._only_student_fwd = False + + def train(self, mode: bool = True): + """Override to prevent warnings of stored intermediate outputs in future forwards.""" + if self.training is True and mode is False: + # If switching from train to eval, clear outputs and stop forward hook + for student_layer, teacher_layer in self._layers_to_loss: + student_layer._intermediate_output = False + teacher_layer._intermediate_output = False + elif self.training is False and mode is True: + # If switching from eval to train, resume forward hook + for student_layer, teacher_layer in self._layers_to_loss: + student_layer._intermediate_output = None + teacher_layer._intermediate_output = [] + super().train(mode) + + def state_dict(self, *args, **kwargs) -> dict[str, Any]: + """Override to potentially return the state without teacher's.""" + with self.hide_teacher_model(enable=self._expose_minimal_state_dict): + return super().state_dict(*args, **kwargs) + + def load_state_dict(self, state_dict, *args, **kwargs) -> Any: + """Override to potentially load the state without teacher's or loss modules'.""" + # Don't expose teacher/loss modules when loading at first from a non-DistillationModel checkpoint. + hide_teacher = self._expose_minimal_state_dict or not any( + k.startswith("_teacher_model") for k in state_dict + ) + hide_losses = len(self._loss_modules) > 0 and not any( + k.startswith("_loss_modules") for k in state_dict + ) + with self.hide_teacher_model(enable=hide_teacher): + with self.hide_loss_modules(enable=hide_losses): + return super().load_state_dict(state_dict, *args, **kwargs) + + def forward(self, *args, **kwargs) -> Any: + """Implement forward pass. + + Args: + *args: Positional inputs to the student and teacher model. + **kwargs: Named inputs to the student and teacher model. + + Returns: + The student model's output. + """ + if not self._only_student_fwd: + # Call teacher model's forward pass for layer outputs to get computed. + # no_grad() context lets pytorch know not to save activations for + # teacher models in memory as there won't be any gradient updates applied + # to these layers. This consumes less memory than just freezing teacher model weights. + with torch.no_grad(): + # Calling `.train()` on this class inadvertently calls it on teacher too. + self._teacher_model.eval() + teacher_output = self._teacher_model(*args, **kwargs) + + if self._only_teacher_fwd: + # Special case for pipeline parallelism when it's desired to run student and teacher separately. + return teacher_output + + student_output = super().forward(*args, **kwargs) + + return student_output + + def compute_kd_loss( + self, + student_loss: Optional[torch.Tensor] = None, + loss_reduction_fn: Callable = None, + skip_balancer: bool = False, + ) -> Union[torch.Tensor, dict[str, torch.Tensor]]: + """Compute total loss for distillation backpropagation. + + Args: + student_loss: Original loss computed from the student's output. + loss_reduction_fn: Callable to be called on each loss tensor prior to balancing. Useful for + loss-masking situations where the callable changes arguments each iteration. + skip_balancer: Whether or not to use loss balancer to reduce the loss dict into a scalar. + + Returns: + If reduce is True, the scalar total loss weighted between ``student_loss`` and the distillation losses. + If reduce is False, a dict of student model output loss and layer-wise distillation losses. + """ + if not self.training: + raise AssertionError( + "`DistillationModel.compute_kd_loss()` should only be called in training mode." + ) + if self._loss_balancer is None: + assert student_loss is None, "Cannot pass in student loss without using Loss Balancer." + + loss_dict: dict[str, torch.Tensor] = {} + if student_loss is not None: + loss_dict[STUDENT_LOSS_KEY] = student_loss + + idx = 0 + for (student_layer, teacher_layer), loss_fn in self._layers_to_loss.items(): + out_s = student_layer._intermediate_output + out_t = teacher_layer._intermediate_output.pop(0) # can store multiple in special cases + student_layer._intermediate_output = None + + loss = loss_fn(out_s, out_t) # Student is pred, Teacher is target + if loss_reduction_fn is not None: + # Needed in cases where a loss mask is used on non-scalar loss-fn outputs, prior to + # reducing to a scalar loss value. + loss = loss_reduction_fn(loss) + loss_dict[f"{loss_fn.__class__.__name__}_{idx}"] = loss + idx += 1 + + if skip_balancer: + # Needed for special case if reduction needs to be done separately before balancing. + return loss_dict + + if self._loss_balancer is None: + assert len(loss_dict) == 1 # we ensure this in constructor + loss_total = next(iter(loss_dict.values())) + else: + loss_total = self._loss_balancer(loss_dict) + + return loss_total + + +def output_capture_fwd_hook(module: nn.Module, input: Any, output: Any): # pylint: disable=redefined-builtin # noqa + """A hook to capture layer output.""" + # NOTE: Defined externally to allow pickling. + + if module._intermediate_output is False: + # Special indicator to temporarily disable hook when not training mode. + return + elif isinstance(module._intermediate_output, list): + # Teacher + if len(module._intermediate_output) > 0: + warnings.warn( + f"Teacher's Module `{type(module).__name__}` already has an intermediate output stored." + " This is undesired behavior unless Pipeline Parallelism is in use." + ) + module._intermediate_output.append(output) + elif not getattr(module, "_only_teacher_fwd", False): # might be hooked on entire model fwd + # Student + if module._intermediate_output is not None: + warnings.warn( + f"Student's Module `{type(module).__name__}` already has an intermediate output stored." + " This is undesired behavior unless Gradient Checkpointing is in use." + ) + module._intermediate_output = output diff --git a/modelopt/torch/distill/loss_balancers.py b/modelopt/torch/distill/loss_balancers.py new file mode 100644 index 0000000000..efb5594421 --- /dev/null +++ b/modelopt/torch/distill/loss_balancers.py @@ -0,0 +1,139 @@ +# 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. + +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +"""Basic loss balancers for Distillation task.""" + +import abc +import warnings +from typing import Any, Callable, Union + +import torch +import torch.nn as nn + +__all__ = ["DistillationLossBalancer", "StaticLossBalancer"] + +STUDENT_LOSS_KEY = "student_loss" + + +class DistillationLossBalancer(nn.Module): + """Interface for loss balancers.""" + + def __init__(self): + """Constructor.""" + super().__init__() + self._student_loss_reduction_fn = None + + def set_student_loss_reduction_fn( + self, student_loss_reduction_fn: Callable[[Any], torch.Tensor] + ): + """Set student loss reduction function value. + + Needed in special case of loss-reducing the student loss prior to balancing. + """ + # NOTE: Cannot use setter decorator since ``nn.Module`` intercepts + # ``__setattr__`` for objects which are also instances of ``nn.Module``. + self._student_loss_reduction_fn = student_loss_reduction_fn + + @abc.abstractmethod + def forward(self, loss: dict[str, torch.Tensor]) -> torch.Tensor: + """Compute aggregate loss. + + Args: + loss: The loss dict to aggregate. + The keys will be the class name of the loss function applied to obtain the loss, + suffixed by `_{idx}` for uniqueness. And if a student loss is provided to + ``mtd.DistillationModel.compute_kd_loss`` then it will have the key + ``mtd.loss_balancers.STUDENT_LOSS_KEY``. For example, if the ``criterion`` argument + to ``mtd.convert`` is + ``{("mod1_s", "mod1_t"): torch.nn.MSELoss(), ("mod2_s", "mod2_t"): torch.nn.MSELoss()}`` + and the student_loss provided to ``mtd.DistillationModel.compute_kd_loss`` is not + None, then the loss dict here will look like + ``{"student_loss": torch.tensor(...), "MSELoss_0": torch.tensor(...), "MSELoss_1": torch.tensor(...)}``. + + Returns: + The total loss after balancing student and kd loss loss components. + """ + raise NotImplementedError + + +class StaticLossBalancer(DistillationLossBalancer): + """Static weights-based loss aggregation of KD losses.""" + + def __init__(self, kd_loss_weight: Union[float, list[float]] = 0.5): + """Constructor. + + Args: + kd_loss_weight: The static weight to be applied to balance the knowledge distillation + loss and original student loss. + If it is a float, it would be applied to the sum(KD losses). + If it is a list, the keys are the KD loss keys, in order specified to the + ``criterion`` argument, and the weight corresponding to each key is applied to + the corresponding loss value. + If the weights do not sum to 1.0, a ``student_loss`` should be passed into + ``mtd.DistillationModel.compute_kd_loss``, and the weight difference will be applied + to this loss value. + + Raises: + ValueError if kd_loss_weight is out of bounds. + """ + super().__init__() + if isinstance(kd_loss_weight, float): + kd_loss_weight = [kd_loss_weight] + + sum_kd_loss_weight = sum(kd_loss_weight) + if sum_kd_loss_weight < 0.0 or sum_kd_loss_weight > 1.0: + raise ValueError( + "The sum of values of kd_loss_weight should be [0., 1.]," + f" actual sum {sum_kd_loss_weight}" + ) + elif sum_kd_loss_weight < 1.0: + warnings.warn( + "`StaticLossBalancer` weights do not sum to 1.0." + " Argument `student_loss` should be passed into `DistillationModel.compute_kd_loss`" + ) + + self._kd_loss_weight = kd_loss_weight + + def forward(self, loss: dict[str, torch.Tensor]) -> torch.Tensor: + """Compute aggregate loss. + + Args: + loss: The loss dict to aggregate. + + Returns: + The total loss after balancing student and kd loss loss components. + """ + kd_loss_dict = loss.copy() + + if STUDENT_LOSS_KEY in loss: + output_loss = kd_loss_dict.pop(STUDENT_LOSS_KEY) + if self._student_loss_reduction_fn is not None: + output_loss = self._student_loss_reduction_fn(output_loss) + else: + output_loss = None + + assert len(kd_loss_dict) == len( + self._kd_loss_weight + ), "Number of `kd_loss_weight` does not correspond to number of kd losses computed." + + aggregate_loss = sum( + loss * weight for loss, weight in zip(kd_loss_dict.values(), self._kd_loss_weight) + ) + if output_loss is not None: + aggregate_loss += (1.0 - sum(self._kd_loss_weight)) * output_loss + + return aggregate_loss diff --git a/modelopt/torch/distill/losses.py b/modelopt/torch/distill/losses.py new file mode 100644 index 0000000000..3ebbdd7903 --- /dev/null +++ b/modelopt/torch/distill/losses.py @@ -0,0 +1,150 @@ +# 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. + +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +"""Different types of distillation losses.""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.modules.loss import _Loss as Loss + +__all__ = ["LogitsDistillationLoss", "MGDLoss"] + + +class LogitsDistillationLoss(Loss): + """KL-Divergence loss on output logits. + + This function implements the distillation loss found in the paper: https://arxiv.org/abs/1503.02531. + """ + + def __init__(self, temperature: float = 1.0, reduction: str = "batchmean"): + """Constructor. + + Args: + temperature: A value used to soften the logits_t and logits_s before computing loss on them. + reduction: How to reduce the final pointwise loss before returning. Pass ``"none"`` to + use your own reduction function afterwards, i.e. with loss masks. + """ + super().__init__() + self._temperature = temperature + self._reduction = reduction + + def forward(self, logits_s: torch.Tensor, logits_t: torch.Tensor) -> torch.Tensor: + """Compute KD loss on student and teacher logits. + + Args: + logits_s: Student's logits, treated as prediction. + logits_t: Teacher's logits, treated as label. + + .. note:: + + Assumes class logits dimension is last. + """ + soft_log_probs = F.log_softmax(logits_s / self._temperature, dim=-1) + soft_targets = F.softmax(logits_t / self._temperature, dim=-1) + + soft_log_probs = soft_log_probs.view(-1, soft_log_probs.size(-1)) + soft_targets = soft_targets.view(-1, soft_targets.size(-1)) + + kd_loss = F.kl_div(soft_log_probs, soft_targets.detach(), reduction=self._reduction) + + # Since the magnitudes of the gradients produced by the soft logits scale as 1/(T^2), + # multiplying them by T^2 ensures that the relative contributions of the logits + # remain roughly unchanged while experimenting with meta-parameters. + kd_loss *= self._temperature**2 + + return kd_loss + + +class MGDLoss(Loss): + """PyTorch version of Masked Generative Distillation. + + This function implements the distillation loss found in the paper: https://arxiv.org/abs/2205.01529. + """ + + def __init__( + self, + num_student_channels: int, + num_teacher_channels: int, + alpha_mgd: float = 1.0, + lambda_mgd: float = 0.65, + ): + """Constructor. + + Args: + num_student_channels: Number of channels in the student's feature map. + num_teacher_channels: Number of channels in the teacher's feature map. + alpha_mgd: Scalar final loss is multiplied by. Defaults to 1.0. + lambda_mgd: Masked ratio. Defaults to 0.65. + """ + super().__init__() + self._alpha_mgd = alpha_mgd + self._lambda_mgd = lambda_mgd + + if num_student_channels != num_teacher_channels: + self.align = nn.Conv2d( + num_student_channels, + num_teacher_channels, + kernel_size=1, + stride=1, + padding=0, + ) + else: + self.align = nn.Identity() + + self.generation = nn.Sequential( + nn.Conv2d( + num_teacher_channels, + num_teacher_channels, + kernel_size=3, + padding=1, + ), + nn.ReLU(inplace=True), + nn.Conv2d( + num_teacher_channels, + num_teacher_channels, + kernel_size=3, + padding=1, + ), + ) + + def _loss_fn(self, out_s: torch.Tensor, out_t: torch.Tensor): + n, _, h, w = out_t.shape + + mat = torch.rand((n, 1, h, w), device=out_s.device) + mat = torch.where(mat > 1 - self._lambda_mgd, 0, 1) + + masked_feats = torch.mul(out_s, mat) + new_feats = self.generation(masked_feats) + + kd_loss = F.mse_loss(new_feats, out_t) + + return kd_loss + + def forward(self, out_s: torch.Tensor, out_t: torch.Tensor): + """Forward function. + + Args: + out_s: Student's feature map (shape BxCxHxW). + out_t: Teacher's feature map (shape BxCxHxW). + """ + assert out_s.shape[-2:] == out_t.shape[-2:] + + out_s = self.align(out_s) + loss = self._loss_fn(out_s, out_t) * self._alpha_mgd + + return loss diff --git a/modelopt/torch/distill/mode.py b/modelopt/torch/distill/mode.py new file mode 100644 index 0000000000..26c4f380ae --- /dev/null +++ b/modelopt/torch/distill/mode.py @@ -0,0 +1,224 @@ +# 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. + +"""Module implementing and describing modes that can be used during the NAS convert process. + +Check out :meth:`mtn.convert ` to learn more about modes. +""" + +import warnings +from typing import Optional + +import torch.nn as nn +from torch.nn.modules.loss import _Loss as Loss + +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.conversion import ModeloptStateManager +from modelopt.torch.opt.mode import ( + ConvertEntrypoint, + ConvertReturnType, + MetadataDict, + RestoreEntrypoint, + UpdateEntrypoint, + _ModeDescriptor, + _ModeRegistryCls, +) +from modelopt.torch.utils import init_model_from_model_like, unwrap_model + +from .config import ExportStudentConfig, KDLossConfig +from .distillation_model import DistillationModel +from .registry import DistillationDMRegistry + +DistillModeRegistry = _ModeRegistryCls() + + +@DistillModeRegistry.register_mode +class KnowledgeDistillationModeDescriptor(_ModeDescriptor): + """Class to describe the Knowledge-Distillation mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "kd_loss" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return KDLossConfig + + @property + def next_modes(self) -> Optional[set[str]]: + """Modes that must immediately follow this mode.""" + return {"export_student"} + + @property + def export_mode(self) -> Optional[str]: + """The mode that corresponds to the export mode of this mode.""" + return "export_student" + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return _convert_for_kd + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return _restore_kd_model + + @property + def update_for_new_mode(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models state for adding new mode.""" + return _update_kd_state_before_new_mode + + +@DistillModeRegistry.register_mode +class ExportStudentModeDescriptor(_ModeDescriptor): + """Class to describe the specific Export mode to be used with Knowledge Distillation. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "export_student" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return ExportStudentConfig + + @property + def is_export_mode(self) -> bool: + """Specifies whether the mode is an export mode.""" + return True + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return _export_student + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return _restore_exported_student + + +def _convert_for_kd(model: nn.Module, config: KDLossConfig) -> ConvertReturnType: + """Function for converting a model to a distillation meta-model. + + This is the only utility needed to use the ``modelopt.torch.distill`` API directly. + + Args: + model: The base model to be used as the student. + config: A KDLossConfig instance defining the configuration options. + + Returns: + A ``DistillationModel`` encapsulating the teacher and students, to be used in place + of the original model. + Metadata dictionary containing the config arguments needed to recreate the model + from a checkpoint. + """ + # we need to do a "strict" validation here to ensure that all required fields are present. + # we don't do the strict validation by default since this is okay during restore or when other + # modes are added on top. + config._strict_validate() + + teacher = init_model_from_model_like(config.teacher_model) + student = model + + # hide teacher's state if it exists within teacher (will be reverted during export) + if ModeloptStateManager.is_converted(teacher): + teacher._modelopt_state_storage = [nn.Module()] + ModeloptStateManager.transfer_state_dict(teacher, teacher._modelopt_state_storage[0]) + warnings.warn( + "Distillation mode in progress. ModelOpt operations on teacher disabled until the" + " 'export_student' mode is used." + ) + + # initialize distillation model + original_cls = type(student) + if original_cls not in DistillationDMRegistry: + DistillationDMRegistry.register({original_cls: "student_class"})(DistillationModel) + # TODO (lucasl): look into ways to avoid registering every class manually + # (e.g. by just registering nn.Module and disable the "forward" check for the inherited class check + + distillation_model = DistillationDMRegistry.convert(student) + distillation_model.modify( + **{**config, "teacher_model": teacher} # overwrite with instantiated teacher + ) + + # no metadata, all specifed via config. + metadata = {} + + return distillation_model, metadata + + +def _restore_kd_model(model: nn.Module, config: KDLossConfig, metadata: MetadataDict) -> nn.Module: + """Function for restoring a previously convert model to a distillation meta-model.""" + # the metadata should be empty + assert not metadata, "No metadata expected!" + + return _convert_for_kd(model, config)[0] + + +def _update_kd_state_before_new_mode( + model: nn.Module, config: KDLossConfig, metadata: MetadataDict +) -> None: + """Function for updating the model's state before new mode.""" + config.teacher_model = nn.Module + config.criterion = Loss() + config.loss_balancer = None + + +def _export_student(model: nn.Module, config: ExportStudentConfig) -> ConvertReturnType: + """Export a ``DistillationModel`` to its inner student model, including modelopt state transfer.""" + if not isinstance(model, DistillationModel): + raise RuntimeError("`model` must be of type `mtd.DistillationModel`") + + # restore previously-renamed state in teacher just in case being used after export. + teacher = model.teacher_model + if hasattr(teacher, "_modelopt_state_storage"): + mod_storage = teacher._modelopt_state_storage[0] + delattr(teacher, "_modelopt_state_storage") + ModeloptStateManager.transfer_state_dict(mod_storage, teacher) + + # retrieve student model + student_model = model.export() + # just in case; This way all our APIs input and output unwrapped models. + student_model = unwrap_model( + student_model, + warn=True, + msg=( + f"The student model is wrapped into {type(student_model).__name__}. Unwrapping and" + " exporting it ..." + ), + ) + + return student_model, {} + + +def _restore_exported_student( + model: nn.Module, config: ExportStudentConfig, metadata: MetadataDict +) -> nn.Module: + """Function for restoring a previously exported distillation meta-model.""" + # no metadata is used by the mode + assert not metadata, "No metadata expected!" + + return _export_student(model, config)[0] diff --git a/modelopt/torch/distill/plugins/__init__.py b/modelopt/torch/distill/plugins/__init__.py new file mode 100644 index 0000000000..97b041e9b9 --- /dev/null +++ b/modelopt/torch/distill/plugins/__init__.py @@ -0,0 +1,21 @@ +# 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. + +"""Handles distillation plugins for third-party modules.""" + +from modelopt.torch.utils import import_plugin + +with import_plugin("megatron"): + from .megatron import * diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py new file mode 100644 index 0000000000..58b71586c8 --- /dev/null +++ b/modelopt/torch/distill/plugins/megatron.py @@ -0,0 +1,477 @@ +# 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. + +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +"""Distillation loss function(s).""" + +import logging +import types +from abc import ABCMeta +from typing import Any, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +import yaml +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.parallel_state import get_tensor_model_parallel_group +from megatron.core.tensor_parallel import gather_from_sequence_parallel_region +from megatron.core.transformer import MegatronModule, TransformerConfig +from torch import Tensor +from torch.nn.modules.loss import _Loss + +import modelopt.torch.distill as mtd + +logger = logging.getLogger(__name__) + + +def load_distillation_config( + config_path: Optional[str], student_cfg: TransformerConfig, teacher_cfg: TransformerConfig +) -> dict[str, Any]: + """Read the distillation yaml config file specified by ``args.export_kd_cfg``. + + Args: + config_path: Path to user-defined distillation settings yaml file. + If `None`, uses default logits-only distillation mode for GPT models. + student_cfg: Model config for student model. + teacher_cfg: Model config for teacher model. + + WARNING: Assumes intermediate hidden sizes are always that found in the model config's ``hidden_size`` attribute. + """ + if not config_path: + logger.warning("Distillation config not provided. Using default.") + cfg = { + "logit_layers": ["output_layer", "output_layer"], + "intermediate_layer_pairs": [], + "skip_lm_loss": True, + "kd_loss_scale": 1.0, + } + else: + with open(config_path) as f: + cfg = yaml.safe_load(f) + + intermediate_pairs: list[str] = cfg["intermediate_layer_pairs"] + logit_pair: list[str] = cfg["logit_layers"] + skip_lm_loss: bool = cfg["skip_lm_loss"] + loss_scale: float = cfg["kd_loss_scale"] + + criterion = {tuple(logit_pair): LogitsKLLoss(student_cfg, teacher_cfg)} + for layer_names in intermediate_pairs: + if torch.distributed.get_rank() == 0: + print( + "Distillation: Adding intermediate loss between" + f" `{layer_names[0]}` of student (hidden size {student_cfg.hidden_size}) and" + f" `{layer_names[1]}` of teacher (hidden size {teacher_cfg.hidden_size})." + ) + criterion[tuple(layer_names)] = HiddenStateCosineLoss(student_cfg, teacher_cfg) + + loss_balancer = LogitsAndIntermediatesLossBalancer( + kd_loss_scale=loss_scale, skip_original_loss=skip_lm_loss + ) + + cfg["criterion"] = criterion + cfg["loss_balancer"] = loss_balancer + + return cfg + + +######################################################## + + +class BaseLoss(_Loss, metaclass=ABCMeta): + """Abstract base class for Megatron distillation losses.""" + + def __init__( + self, + student_config: TransformerConfig, + teacher_config: TransformerConfig, + projection_layer: bool = False, + ): + """Constructor. + + Args: + student_config: Student's MCore transformer config. + teacher_config: Teacher's MCore transformer config. + projection_layer: If True, create a linear layer to project student tensor to teacher's hidden dim. + """ + super().__init__() + self._config = student_config + self._tensor_parallel = self._config.tensor_model_parallel_size > 1 + self._sequence_parallel = self._config.sequence_parallel + + if projection_layer: + self._projection = ProjectionLayer(student_config, teacher_config) + else: + self._projection = None + + def pre_forward(self, predictions: Tensor, targets: Tensor) -> tuple[Tensor, Tensor]: + """Performs projection of student tensor to match teacher's size if necessary.""" + if isinstance(predictions, tuple): + # `ColumnParallelLinear` returns bias too + predictions, targets = predictions[0], targets[0] + + if self._projection is not None: + predictions = self._projection(predictions) + targets = targets.detach() + + return predictions, targets + + def post_forward(self, loss: Tensor, tp_reduce: bool = False) -> Tensor: + """Reshapes tensor from [s, b] to [b, s] for upcoming loss masking.""" + loss = loss.transpose(0, 1).contiguous() + return (loss, tp_reduce) + + +class MSELoss(BaseLoss): + """Calculates Mean Squared Error loss between two tensors without reducing the sequence dim.""" + + def __init__(self, student_config: TransformerConfig, teacher_config: TransformerConfig): + """Constructor. + + Args: + student_config: Student's MCore transformer config. + teacher_config: Teacher's MCore transformer config. + """ + super().__init__(student_config, teacher_config) + + def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: + """Forward function. + + Args: + predictions: Student model tensors (size [s, b, h]) + targets: Teacher model tensors (size [s, b, h]) + + Returns: + MSE loss of tensors (size [b, s]) + """ + predictions, targets = self.pre_forward(predictions, targets) + + # TP irrelevant since MSE loss gradients are per-input element. + loss = F.mse_loss(predictions, targets, reduction="none") + loss = loss.sum(dim=-1) + + return self.post_forward(loss) + + +class HiddenStateCosineLoss(BaseLoss): + """Calculates Cosine loss between two tensors without reducing the sequence dim. + + The tensors are assumed to be intermediate activations, so extra restrictions are in place. + """ + + def __init__(self, student_config: TransformerConfig, teacher_config: TransformerConfig): + """Constructor. + + Args: + student_config: Student's MCore transformer config. + teacher_config: Teacher's MCore transformer config. + """ + super().__init__(student_config, teacher_config, projection_layer=True) + + if self._tensor_parallel and not self._sequence_parallel: + logger.warning( + "``HiddenStateCosineLoss`` only works with tensors with full hidden dim. Ensure the " + "tensor inputs meet this requirement or use `--sequence_parallel` if tensor parallel is enabled." + ) + + def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: + """Forward function. + + Args: + predictions: Student model tensors (size [s, b, h]) + targets: Teacher model tensors (size [s, b, h]) + + Returns: + Cosine loss of tensors (size [b, s]) + """ + predictions, targets = self.pre_forward(predictions, targets) + + loss = F.cosine_embedding_loss( + predictions.view(-1, predictions.size(-1)), + targets.view(-1, targets.size(-1)), + targets.new_ones(1), + reduction="none", + ) + loss = loss.view(*predictions.shape[:2]) + + if self._sequence_parallel: + # Can efficiently gather size [s, b] tensor now for loss-masking purposes. + # TODO(aanoosheh) Reconsider for memory savings by splitting loss mask instead. + loss = gather_from_sequence_parallel_region(loss) + + return self.post_forward(loss) + + +class LogitsKLLoss(BaseLoss): + """Calculates KL-Divergence loss between two logits tensors without reducing the sequence dim.""" + + def __init__( + self, + student_config: TransformerConfig, + teacher_config: TransformerConfig, + temperature: float = 1.0, + reverse: bool = False, + ): + """Constructor. + + Args: + student_config: Student's MCore transformer config. + teacher_config: Teacher's MCore transformer config. + temperature: Divide tensors by this value prior to calculating loss. + reverse: Whether to reverse the loss as KLD(teacher, student) instead of KLD(student, teacher) + """ + super().__init__(student_config, teacher_config) + self._temperature = temperature + self._reverse = reverse + + def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: + """Forward function. + + Args: + predictions: Student model tensors (size [s, b, h]) + targets: Teacher model tensors (size [s, b, h]) + + Returns: + KLD loss of tensors (size [b, s]) + """ + predictions, targets = self.pre_forward(predictions, targets) + + # Division by temp should happen prior to finding max for both student and teacher. + # Currently we don't use temperature in any of ours runs (temp=1.0) + output_teacher = targets.float() / self._temperature + output_student = predictions.float() / self._temperature + + # Compute local softmax, and the reweight to compute global softmax. + if self._tensor_parallel: + # Maximum value along vocab dimension across all GPUs. + teacher_logits_max, _ = torch.max(output_teacher, dim=-1) + torch.distributed.all_reduce( + teacher_logits_max, + op=torch.distributed.ReduceOp.MAX, + group=get_tensor_model_parallel_group(), + ) + output_teacher = output_teacher - teacher_logits_max.unsqueeze(dim=-1) + + denom_teacher = torch.sum(torch.exp(output_teacher), dim=-1) + # We can't use `gather_from_tensor_model_parallel_region` here since it discards + # gradients from other ranks - we need to all_reduce the gradients as well. + denom_teacher = all_reduce_autograd( + denom_teacher, group=get_tensor_model_parallel_group() + ) + + # Maximum value along vocab dimension across all GPUs. + student_logits_max, _ = torch.max(output_student, dim=-1) + torch.distributed.all_reduce( + student_logits_max, + op=torch.distributed.ReduceOp.MAX, + group=get_tensor_model_parallel_group(), + ) + output_student = output_student - student_logits_max.unsqueeze(dim=-1).detach() + + denom_student = torch.sum(torch.exp(output_student), dim=-1) + denom_student = all_reduce_autograd( + denom_student, group=get_tensor_model_parallel_group() + ) + + slen, bsz, sharded_vocab_size = output_student.shape + student_log_prob = output_student - torch.log(denom_student).view(slen, bsz, 1).expand( + slen, bsz, sharded_vocab_size + ) + teacher_log_prob = output_teacher - torch.log(denom_teacher).view(slen, bsz, 1).expand( + slen, bsz, sharded_vocab_size + ) + + if self._reverse: + loss = torch.sum( + F.kl_div(teacher_log_prob, student_log_prob, reduction="none", log_target=True), + dim=-1, + ) + else: + loss = torch.sum( + F.kl_div(student_log_prob, teacher_log_prob, reduction="none", log_target=True), + dim=-1, + ) + + else: + if self._reverse: + loss = torch.sum( + F.kl_div( + F.log_softmax(output_teacher, dim=-1), + F.softmax(output_student, dim=-1), + reduction="none", + ), + dim=-1, + ) + else: + loss = torch.sum( + F.kl_div( + F.log_softmax(output_student, dim=-1), + F.softmax(output_teacher, dim=-1), + reduction="none", + ), + dim=-1, + ) + + return self.post_forward(loss, tp_reduce=True) + + +######################################################## + + +class LogitsAndIntermediatesLossBalancer(mtd.DistillationLossBalancer): + """LossBalancer implementation for Logit and Intermediate losses. + + Dynamically weighs distillation and original losses to balance during training. + """ + + def __init__(self, kd_loss_scale: float = 1.0, skip_original_loss: bool = False): + """Constructor. + + Args: + kd_loss_scale: Multiply distillation losses by this before weighing. + (Not used when `skip_original_loss` is True.) + skip_original_loss: Used to signal whether the original loss should be used, regardless + of whether it was passed into ``mtd.DistillationModel.compute_kd_loss()`` or not. + """ + super().__init__() + self._kd_loss_scale = kd_loss_scale + self._skip_original_loss = skip_original_loss + + def forward(self, loss_dict: dict[str, Tensor]) -> Tensor: + """Forward function. + + Args: + loss_dict: All individual scalar losses, passed in during ``mtd.DistillationModel.compute_kd_loss()`` + + Returns: + Aggregate total scalar loss. + """ + original_loss = loss_dict.pop(mtd.loss_balancers.STUDENT_LOSS_KEY) + for _key, _loss in loss_dict.items(): + if _key.startswith(LogitsKLLoss.__name__): + logits_loss = _loss # should only be one + intermediate_loss = sum(loss_dict.values()) + + if intermediate_loss > 0: + dynamic_scale = logits_loss.item() / intermediate_loss.item() + intermediate_loss *= dynamic_scale + kd_loss_scale = self._kd_loss_scale / 2.0 + else: + kd_loss_scale = self._kd_loss_scale + + if self._skip_original_loss: + kd_loss = logits_loss + intermediate_loss + total_loss = kd_loss + else: + kd_loss = (logits_loss + intermediate_loss) * kd_loss_scale + dynamic_scale = original_loss.item() / kd_loss.item() + total_loss = original_loss + kd_loss * dynamic_scale + + return total_loss + + +######################################################## + + +class ProjectionLayer(MegatronModule): + """Module to project student layer activations to teacher's size.""" + + def __init__( + self, + student_config: TransformerConfig, + teacher_config: TransformerConfig, + ): + """Constructor. + + Args: + student_config: Student's MCore transformer config. + teacher_config: Teacher's MCore transformer config. + """ + super().__init__(config=student_config) + if student_config.hidden_size == teacher_config.hidden_size: + self._fit = nn.Identity() + else: + self._fit = nn.Linear(student_config.hidden_size, teacher_config.hidden_size) + self.apply(self._init_weights) + setattr(self._fit.weight, "sequence_parallel", self.config.sequence_parallel) + setattr(self._fit.bias, "sequence_parallel", self.config.sequence_parallel) + + def forward(self, student_tensor: Tensor): + """Forward function. + + Args: + student_tensor: Tensor to be fit to teacher size. + """ + return self._fit(student_tensor) + + def _init_weights(self, module): + """Initialize the weights.""" + if isinstance(module, (nn.Linear, nn.Embedding)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + module.weight.data.normal_(mean=0.0, std=0.01) + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + if isinstance(module, nn.Linear) and module.bias is not None: + module.bias.data.zero_() + + +class _AllReduce(torch.autograd.Function): + """Implementation from old PyTorch `torch.distributed.nn.parallel`.""" + + @staticmethod + def forward(ctx, op, group, tensor): + ctx.group, ctx.op = group, op + tensor = tensor.clone() + torch.distributed.all_reduce(tensor, op=op, group=group) + return tensor + + @staticmethod + def backward(ctx, grad_output): + return (None, None, _AllReduce.apply(ctx.op, ctx.group, grad_output)) + + +def all_reduce_autograd( + tensor, op=torch.distributed.ReduceOp.SUM, group=torch.distributed.group.WORLD +): + """AllReduce with autograd.""" + return _AllReduce.apply(op, group, tensor) + + +######################################################## + + +def adjust_distillation_model_for_mcore(model: mtd.DistillationModel, distill_cfg: dict[str, Any]): + """Extra modifcations to ``mtd.DistillationModel`` requried for Megatron-Core.""" + + # HACK: Hide teacher during `sharded_state_dict` method. + def _sharded_state_dict(self, *args, **kwargs) -> ShardedStateDict: + with self.hide_teacher_model(): + return self._sharded_state_dict(*args, **kwargs) + + model._sharded_state_dict = model.sharded_state_dict + model.sharded_state_dict = types.MethodType(_sharded_state_dict, model) + + # HACK: Skip `lm_loss` bypassing it when training if not needed for backprop. + def _compute_language_model_loss(self, labels, logits) -> Tensor: + if self.training: + return torch.zeros_like(labels) + return self._compute_language_model_loss(labels, logits) + + if distill_cfg["skip_lm_loss"]: + model._compute_language_model_loss = model.compute_language_model_loss + model.compute_language_model_loss = types.MethodType(_compute_language_model_loss, model) diff --git a/modelopt/torch/distill/registry.py b/modelopt/torch/distill/registry.py new file mode 100644 index 0000000000..905378cc7c --- /dev/null +++ b/modelopt/torch/distill/registry.py @@ -0,0 +1,23 @@ +# 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. + +"""A dynamic module registry that holds information about available DynamicModules for Distillation-related purposes.""" + +from modelopt.torch.opt.dynamic import _DMRegistryCls + +__all__ = ["DistillationDMRegistry"] + + +DistillationDMRegistry = _DMRegistryCls(prefix="Distill") # global instance for the registry diff --git a/modelopt/torch/export/__init__.py b/modelopt/torch/export/__init__.py new file mode 100644 index 0000000000..9a9c5bc544 --- /dev/null +++ b/modelopt/torch/export/__init__.py @@ -0,0 +1,21 @@ +# 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. + +"""Export package. So far it only supports selected nemo and huggingface LLMs.""" + +from .model_config_export import * +from .transformer_engine import * +from .unified_export_hf import * +from .unified_export_megatron import * diff --git a/modelopt/torch/export/distribute.py b/modelopt/torch/export/distribute.py new file mode 100644 index 0000000000..f8ea500796 --- /dev/null +++ b/modelopt/torch/export/distribute.py @@ -0,0 +1,303 @@ +# 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. + +"""torch.distribute utils.""" + +import json +from contextlib import contextmanager +from io import BytesIO +from multiprocessing.shared_memory import SharedMemory +from pathlib import Path +from typing import Any, Optional, Union + +import torch + +from modelopt.torch.utils import distributed as dist + +from .model_config_utils import ( + model_config_from_dict, + model_config_to_dict, + restore_model_config, + split_config_and_weights, +) + + +class NFSWorkspace: + """A shared workspace implementation using Network File Storage (NFS). + + NOTE: all read/write/modifition to the NFS dir do not involve any collective + communication nor barrier. It is users' responsibility to synchronize + all ranks (local and remove processes). + + This implementation uses `torch.save` and `torch.load` for serialization. + + Args: + workspace_path: the path to the NFS directory for postprocess cross rank communication. + If not provided, SharedMemory will be used instead. + """ + + def __init__(self, workspace_path: Optional[Union[Path, str]] = None): + """Create the NFS work dir and clean up existing existing state files.""" + self.path = Path("") if workspace_path is None else Path(workspace_path) + self._is_initialized = workspace_path is not None + self.rank = dist.rank() + if self.is_initialized: + if self.rank == 0: + self.path.mkdir(parents=True, exist_ok=True) + self.state_path = self._get_state_path(self.rank) + self._clean_up() + + @property + def is_initialized(self): + """Whether the workspace is intialized.""" + return self._is_initialized + + def write_configs_and_weights(self, config_json: dict[str, Any], weights: dict[str, Any]): + """All ranks write the state file to the shared NFS dir. + + Args: + config_json: model or module config in json + weights: module weights in torch's state_dict format + """ + if not self.is_initialized: + raise ValueError("NFSWorkspace is not initialized!") + self._clean_up() + torch.save({"config": config_json, "weight": weights}, self.state_path) + + def read_configs_and_weights_from_rank( + self, target_rank: int + ) -> tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]: + """All ranks read the target_rank state file. + + Args: + target_rank: the target rank + + Returns: + the model/module config and the weights + """ + if not self.is_initialized: + raise ValueError("NFSWorkspace is not initialized!") + state_path = self._get_state_path(target_rank) + if state_path.exists(): + state = torch.load(state_path, map_location="cpu") + return state["config"], state["weight"] + else: + return None, None + + def _get_state_path(self, target_rank: int) -> Path: + """Return the state file name of a particular rank. + + Args: + target_rank: the target rank + + Returns: + the state file path of the target rank + """ + if not self.is_initialized: + raise ValueError("NFSWorkspace is not initialized!") + return self.path.joinpath("rank_{}_state.pth".format(target_rank)) + + def _clean_up(self): + """Remove existing state files.""" + if not self.is_initialized: + raise ValueError("NFSWorkspace is not initialized!") + self.state_path.unlink(missing_ok=True) + + +@contextmanager +def get_tensors_parallel(tensor: torch.Tensor, ranks: list[int], group=None): + """Gathers the tensors across distributed processes using shm. + + Args: + tensor: the tensor that each rank want to pass to the first rank. + The tensors across the ranks need to have the same size. + ranks: the list of the ranks + group: the barrier sync group. + + Yields: + the first rank in the ranks has the full access of the tensors across all the ranks. + the other ranks returns an empty list + + The shm will be destroyed after consumption. + """ + assert tensor is not None + assert len(ranks) > 1 + local_rank = dist.rank() + shm_writer = None + shm_readers = [] + tensor = tensor.cpu() + + is_merged_rank = local_rank == ranks[0] + # Create shm and copy the tensor to the shm if not the merged rank. + # Assume each tensor need up to 2KB additional space for metadata. + if not is_merged_rank: + shm_writer = SharedMemory(name=f"rank_{local_rank}", create=True, size=tensor.nbytes + 2048) + torch.save(tensor, shm_writer._mmap) # type: ignore[attr-defined] + # All ranks wait for this to complete. + dist.barrier(group) + + tensors = [] + # The merged rank gather the tensor from the other ranks (including itself). + if is_merged_rank: + for rank in ranks: + if rank == ranks[0]: + tensors.append(tensor) + else: + shm = SharedMemory(name=f"rank_{rank}", create=False) + shared_tensor = torch.load(BytesIO(shm.buf)) + tensors.append(shared_tensor) + shm_readers.append(shm) + try: + # Send the tensor list to the consumer. + # The merged rank will get a valid tensor while the other ranks an empty tensor. + yield tensors + finally: + # Reader closes the shms. + if shm_readers: + for shm in shm_readers: + shm.close() + + # All ranks wait for the reader to close the shms. + dist.barrier(group) + + # Writer frees the shm resource. + if shm_writer is not None: + shm_writer.close() + shm_writer.unlink() + + +@contextmanager +def get_configs_parallel( + config, ranks: list[int], group, workspace_path: Optional[Union[Path, str]] = None +): + """Gathers the layer config across distributed processes using shm or NFS. + + Args: + config: the config (nullable) that each rank want to pass to the first rank. + ranks: the list of the ranks + group: the barrier sync group. + workspace_path: the path to the NFS directory for postprocess cross rank communication. + + Yields: + the first rank in the ranks has the full access of the configs across all the ranks. + the other ranks returns an empty list + + When workspace_path is provided, an NFSWorkspace object is created to perform communication + across ranks. Otherwise, `SharedMemory` is used for local multi-process communication. + The shm will be destroyed after consumption. + """ + assert len(ranks) > 1 + local_rank = dist.rank() + shm_writer = None + shm_readers = [] + nfs_workspace = NFSWorkspace(workspace_path) + + is_merged_rank = local_rank == ranks[0] + + def _get_weights_nbytes(weights_dict: dict[str, torch.Tensor]): + total_nbytes = 0 + for k, v in weights_dict.items(): + # Assume each tensor need up to 2KB additional space for metadata. + # In reality this should be much smaller. + total_nbytes = total_nbytes + len(k) + v.nbytes + 2048 + + return total_nbytes + + # Create shm and copy the serialized config to the shm if not the merged rank. + if not is_merged_rank: + if config is not None: + config_dict = model_config_to_dict(config) + # Add additional config type name to the dict so we can later pick the right config type. + config_dict["__name__"] = str(type(config).__name__) + weights = {} + split_config_and_weights(config_dict, weights) + + config_json = json.dumps(config_dict) + + if nfs_workspace.is_initialized: + # All ranks except for the master merge rank write to the NFS dir. + nfs_workspace.write_configs_and_weights(config_dict, weights) + else: + # SHM data structure: 8B json size, serialized json bytes and the weights dict. + shm_writer = SharedMemory( + name=f"rank_{local_rank}_config", + create=True, + size=(8 + len(config_json) + _get_weights_nbytes(weights)), + ) + + # Write json length to the shm + shm_writer.buf[:8] = len(config_json).to_bytes(8, "little") + + # Write json to the shm + shm_writer.buf[8 : len(config_json) + 8] = config_json.encode() + + # Write np tensors to the shm. + shm_writer._mmap.seek(len(config_json) + 8) # type: ignore[attr-defined] + torch.save(weights, shm_writer._mmap) # type: ignore[attr-defined] + else: + # If the config is None, we just store the empty 0. + shm_writer = SharedMemory( + name=f"rank_{local_rank}_config", + create=True, + size=8, + ) + + shm_writer.buf[:8] = int(0).to_bytes(8, "little") + + # All ranks wait for this to complete. + dist.barrier(group) + + configs = [] + if is_merged_rank: + for rank in ranks: + if rank == ranks[0]: + configs.append(config) + else: + if nfs_workspace.is_initialized: + # The master merge rank read other configs from the NFS dir. + config_dict, weights = nfs_workspace.read_configs_and_weights_from_rank(rank) + if config_dict is not None: + restore_model_config(config_dict, weights) + config = model_config_from_dict(config_dict) + configs.append(config) + else: + shm = SharedMemory(name=f"rank_{rank}_config", create=False) + len_json = int.from_bytes(shm.buf[:8], "little") + + if len_json != 0: + config_dict = json.loads(shm.buf[8 : 8 + len_json].tobytes().decode()) + weights = torch.load(BytesIO(shm.buf[8 + len_json :])) + restore_model_config(config_dict, weights) + config = model_config_from_dict(config_dict) + + configs.append(config) + shm_readers.append(shm) + try: + # Send the config list to the consumer. + # The merged rank will get a valid config list while the other ranks an empty list. + yield configs + finally: + # Reader closes the shms. + if shm_readers: + for shm in shm_readers: + shm.close() + + # All ranks wait for the reader to close the shms. + dist.barrier(group) + + # Writer frees the shm resource. + if shm_writer is not None: + shm_writer.close() + shm_writer.unlink() diff --git a/modelopt/torch/export/hf_config_map.py b/modelopt/torch/export/hf_config_map.py new file mode 100644 index 0000000000..28ab6ff3e5 --- /dev/null +++ b/modelopt/torch/export/hf_config_map.py @@ -0,0 +1,81 @@ +# 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. + +"""Define the config mapping between HF and modelopt.""" + +# the map is a list of tuples, and each tuple has two elements +# first element: a list of possible fields on HF +# second element: a name of the layer config field inside modelopt +HF_CONFIG_MAP = [ + (["n_head", "num_attention_heads", "n_heads", "num_heads"], "num_attention_heads"), + ( + [ + "num_key_value_heads", + "num_kv_heads", + "num_kv", + "n_head_kv", + "multi_query_group_num", + "num_query_groups", + "kv_n_heads", + ], + "num_kv_heads", + ), + ( + ["n_positions", "max_position_embeddings", "max_seq_len", "max_sequence_length"], + "max_position_embeddings", + ), + (["rotary_percentage", "rotary_percent", "rotary_pct", "partial_rotary_factor"], "rotary_pct"), + (["alibi"], "use_alibi"), + (["alibi_bias_max"], "alibi_bias_max"), + (["new_decoder_architecture"], "new_decoder_architecture"), # Falcon + (["apply_residual_connection_post_layernorm"], "apply_residual_connection_post_layernorm"), + (["use_cache"], "use_cache"), + (["rope_ratio"], "rope_ratio"), # Chatglm + (["parallel_attn"], "parallel_attention"), + (["seq_length"], "seq_length"), # Qwen + # For Qwen and phi3small rotary_emb_base + # CodeLlama using different rotary_base in comparison to LLaMA v1/v2 models + (["rotary_emb_base", "rope_theta", "rotary_base", "rope_embedding_base"], "rotary_base"), + (["original_max_position_embeddings"], "original_max_position_embeddings"), # Phi3 + (["partial_rotary_factor"], "partial_rotary_factor"), # Phi3 + (["mup_attn_multiplier"], "mup_attn_multiplier"), # Phi3-small + (["mup_embedding_multiplier"], "mup_embedding_multiplier"), # Phi3-small + (["mup_use_scaling"], "mup_use_scaling"), # Phi3-small + (["mup_width_multiplier"], "mup_width_multiplier"), # Phi3-small + (["blocksparse_block_size"], "blocksparse_block_size"), # Phi3-small + (["blocksparse_homo_head_pattern"], "blocksparse_homo_head_pattern"), # Phi3-small + (["blocksparse_num_local_blocks"], "blocksparse_num_local_blocks"), # Phi3-small + (["blocksparse_vert_stride"], "blocksparse_vertical_stride"), # Phi3-small + (["dense_attention_every_n_layers"], "dense_attention_every_n_layers"), # Phi3-small + (["gegelu_limit"], "gegelu_limit"), # Phi3-small + ( + ["num_local_experts", "moe_num_experts"], + "moe_num_experts", + ), # Mixture of Experts (Mixtral, DBRX) + (["num_experts_per_tok", "moe_top_k"], "moe_top_k"), # Mixture of Experts (Mixtral, DBRX) + (["model_type"], "qwen_type"), # qwen + (["lru_width"], "rnn_hidden_size"), # Recurrent Gemma + (["embeddings_scale_by_sqrt_dim"], "emb_scale_by_sqrt_dim"), # Recurrent Gemma + (["logits_soft_cap"], "logits_soft_cap"), # Recurrent Gemma + (["_block_types"], "layer_types"), # Recurrent Gemma + (["final_logit_softcapping"], "final_logit_softcapping"), # Gemma2 + (["attn_logit_softcapping"], "attn_logit_softcapping"), # Gemma2 + (["query_pre_attn_scalar"], "query_pre_attn_scalar"), # Gemma2 + (["clip_qkv"], "clip_qkv"), # DBRX + (["rope_scaling"], "rope_scaling"), # Llama 3.1 + (["cross_attention_layers"], "cross_attention_layers"), # Mllama + (["vision_output_dim"], "vision_output_dim"), # Mllama + (["router_jitter_noise"], "sparse_mixer_epsilon"), # Phi3.5-moe +] diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py new file mode 100755 index 0000000000..12571130f8 --- /dev/null +++ b/modelopt/torch/export/layer_utils.py @@ -0,0 +1,1570 @@ +# 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. + +"""Utils for model_config export. + +Some of the logics in this file are empirical and needs constant update if exceptions occur. +""" + +from typing import Optional +from warnings import warn + +import torch +import torch.nn as nn + +try: + from transformers.activations import ACT2FN +except Exception: + warn("Cannot find transformers package. Hugginface modules cannot be exported.") + pass + +from modelopt.torch.utils import distributed as dist +from modelopt.torch.utils import import_plugin + +from ..quantization.nn import SequentialQuantizer +from .hf_config_map import HF_CONFIG_MAP +from .mcore_config_map import MCORE_CONFIG_MAP +from .model_config import ( + LAYERNORM_DEFAULT, + LAYERNORM_RMS, + LINEAR_COLUMN, + LINEAR_GROUP, + LINEAR_ROW, + QUANTIZATION_FP8, + QUANTIZATION_NONE, + AttentionConfig, + ConvConfig, + DecoderLayerConfig, + EmbeddingConfig, + ExpertConfig, + LayernormConfig, + LinearActConfig, + LinearConfig, + MedusaHeadConfig, + MLPConfig, + MOEConfig, + QKVConfig, + RecurrentConfig, + RelativeAttentionTableConfig, + RgLruConfig, +) +from .model_config_utils import pad_weights +from .postprocess import view_as_float8_e4m3fn_if_needed, view_as_uint8_if_needed +from .quant_utils import ( + get_activation_scaling_factor, + get_kv_cache_dtype, + get_kv_cache_scaling_factor, + get_prequant_scaling_factor, + get_quantization_format, + get_weight_block_size, + get_weight_scaling_factor, + get_weight_scaling_factor_2, + preprocess_linear_fusion, +) + +has_mcore = False +with import_plugin("megatron", verbose=False): + from megatron.core.transformer.module import MegatronModule + + has_mcore = True + + +def get_experts_list(module: torch.nn.Module, model_type: str): + """Returns list of grouped experts by linear name for given module.""" + experts_list = [] + if "mixtralforcausallm" in model_type: + for linear_name in ["w1", "w2", "w3"]: + experts_list.append( + [ + _get_mixtral_expert(module.experts, i, linear_name) + for i in range(len(module.experts)) + ] + ) + else: + raise NotImplementedError(f" {model_type} not supported") + + return experts_list + + +def get_dtype(model): + """Returns the default dtype of the model.""" + for weight in model.parameters(): + if torch.is_floating_point(weight): + return weight.dtype + return None + + +def check_model_compatibility(module_list: list[nn.Module]) -> tuple[bool, bool, bool]: + """Returns whether the list of modules is compatible with the export logic. + + And if positional embedding and embedding layernorm exists. + + We assumes the model to be assembled with one or two embedding layers, + a ModuleList of transformer decoders, + and a final layernorm with optional embedding layernorm. + Otherwise it will not be supported. + """ + num_embeddings = 0 + num_module_list = 0 + num_layer_norm = 0 + for module in module_list: + if is_embedding(module): + num_embeddings += 1 + elif is_decoder_list(module): + num_module_list += 1 + elif is_layernorm(module): + num_layer_norm += 1 + + return ( + 1 <= num_embeddings + and num_embeddings <= 2 + and num_module_list == 1 + and 1 <= num_layer_norm + and num_layer_norm <= 2, + num_embeddings > 1, + num_layer_norm > 1, + ) + + +def get_transformer_layers(model: nn.Module) -> list[nn.Module]: + """Returns the root module of the transformer model.""" + if "Megatron" in type(model).__name__: + if hasattr(model, "model") and "GPTModel" in type(model.model).__name__: + # NEMO mcore models can be handled with the following branch. + model = model.model + + # NEMO non mcore models, we need to find the language_model module first. + children = [model] + language_model = None + while children and not language_model: + next_children = [] + for child in children: + if type(child).__name__ == "TransformerLanguageModel": + language_model = child + break + for m in child.children(): + next_children.append(m) + children = next_children + if language_model: + warn("Warning: this is an old NEMO checkpoint format and will be deprecated soon.") + layers = [m for m in language_model.embedding.children()] + [ + m for m in language_model.encoder.children() + ] + + if hasattr(language_model, "output_layer"): + layers.append(language_model.output_layer) + + return layers + + if "GPTModel" in type(model).__name__: + # mcore models + layers = [] + if hasattr(model, "embedding"): + layers = layers + [m for m in model.embedding.children()] + layers = layers + [m for m in model.decoder.children()] + if hasattr(model, "output_layer"): + layers.append(model.output_layer) + return layers + + if hasattr(model, "glm"): + model = model.glm + + if hasattr(model, "transformer"): + # This is a LMHead model + # Add lm_head to be processed along with transformer layers + modules = [] + for m in model.transformer.children(): + # QwenVL's visual encoder name as 'VisionTransformer' has no `layers`. + if ( + "Transformer" in type(m).__name__ + and hasattr(m, "layers") + and is_decoder_list(m.layers) + ): + modules.append(m.layers) + modules.append(m.final_layernorm) + else: + modules.append(m) + if hasattr(model, "lm_head"): + modules += [model.lm_head] + return modules + + if hasattr(model, "model"): + # LLAMA, InternLM2 + modules = [m for m in model.model.children()] + # LLAMA + if hasattr(model, "lm_head"): + modules += [model.lm_head] + # InternLM2 + elif hasattr(model, "output"): + modules += [model.output] + + return modules + + return [m for m in model.children()] + + +def is_linear(module: nn.Module) -> bool: + """Returns whether the module is a linear layer.""" + return any([k in type(module).__name__ for k in ["Linear", "Conv1D", "NormHead"]]) + + +def is_embedding(module: nn.Module) -> bool: + """Returns whether the module is an embedding layer.""" + module_type_name = type(module).__name__ + return ( + "Embedding" in module_type_name + and "Rotary" not in module_type_name + and "PhiImage" not in module_type_name + and "Phi3Image" not in module_type_name + ) + + +def build_embedding_config(module: nn.Module, normalization_constant: float = 1) -> EmbeddingConfig: + """Builds the embedding config from the module.""" + assert is_embedding(module) + + world_size = dist.size() + rank = dist.rank() + + # Special case for chatglm + if hasattr(module, "word_embeddings"): + module = module.word_embeddings + + weight = module.weight + normalized_weight = weight * normalization_constant + if "Parallel" in type(module).__name__: + local_weight = normalized_weight + else: + padded_weight = pad_weights(normalized_weight, dist.size()) + local_weight = torch.chunk(padded_weight, world_size, dim=0)[rank] + return EmbeddingConfig(weight=local_weight) + + +def is_layernorm(module: nn.Module) -> bool: + """Returns whether the module is a layernorm layer.""" + module_name = type(module).__name__ + return any(norm in module_name for norm in ["LayerNorm", "RMSNorm"]) + + +def build_layernorm_config(module: nn.Module) -> LayernormConfig: + """Builds the layernorm config from the module.""" + assert is_layernorm(module) + + layernorm_type = LAYERNORM_DEFAULT + if "RMS" in type(module).__name__: + layernorm_type = LAYERNORM_RMS + + weight = module.weight + + def _weights_plus_one(module): + if any( + name in type(module).__name__ + for name in ["LayerNorm1P", "GemmaRMSNorm", "Gemma2RMSNorm"] + ): + return True + + if hasattr(module, "zero_centered_gamma") and module.zero_centered_gamma: + return True + + return False + + if _weights_plus_one(module): + # megatron layernorm's weight needs to be updated. + weight = weight.float() + 1.0 + + config = LayernormConfig( + weight=weight, + bias=(module.bias if hasattr(module, "bias") and module.bias is not None else None), + layernorm_type=layernorm_type, + ) + + # TODO: handle the nemo llama eps config. + for eps_key in ["eps", "variance_epsilon"]: + if hasattr(module, eps_key): + config.eps = getattr(module, eps_key) + break + + return config + + +def is_decoder_list(module: nn.Module) -> bool: + """Returns whether the module is a decoder list.""" + return type(module) is nn.ModuleList + + +def is_attention(module: nn.Module) -> bool: + """Returns whether the module is an attention layer.""" + return "Attention" in type(module).__name__ + + +def is_mlp(module: nn.Module) -> bool: + """Returns whether the module is an MLP layer.""" + return any([key in type(module).__name__.upper() for key in ("MLP", "T5DENSE")]) + + +def is_moe(module: nn.Module) -> bool: + """Returns whether the module is an MOE layer.""" + return type(module).__name__ in [ + "MixtralSparseMoeBlock", + "ArcticMoE", + "DbrxFFN", + "MoELayer", + "PhiMoESparseMoeBlock", + "DeepseekMoE", + ] + + +def is_quantlinear(module: nn.Module) -> bool: + """Returns whether the module is a quantized linear layer.""" + return "QuantLinear" in type(module).__name__ + + +def dup_kv_weight(v: torch.Tensor, head_size: int, num_head: int, tp_size: int) -> torch.Tensor: + """Repeat kv heads if tp_size > num_kv_heads.""" + assert tp_size % num_head == 0 + reps = tp_size // num_head + v = v.view(-1, head_size, v.size(-1)) + v = v.repeat_interleave(reps, dim=0) + v = v.view(-1, v.size(-1)) + return v.contiguous() + + +def build_qkv( + qkv_modules: list[nn.Module], + model_metadata_config, + ext_config: DecoderLayerConfig = None, + tp_size: int = 1, +) -> QKVConfig: + """Converts the qkv modules to the config.""" + config = QKVConfig() + q_bias = None + k_bias = None + v_bias = None + + block_size = get_weight_block_size(qkv_modules[0]) + + num_heads = ext_config.num_attention_heads + training_tp = model_metadata_config["training_tensor_parallel"] + if len(qkv_modules) == 1: + # QKV layers combined as a single module, e.g. gpt + qkv_module = qkv_modules[0] + assert ext_config is not None, "ext_config is None" + num_kv_heads = ext_config.num_kv_heads + + if "ColumnParallelLinear" in type(qkv_module).__name__: + # For NEMO model, num_kv_heads/num_attention_heads is the first dimension of QKV + model_metadata_config["head_is_first_dim"] = True + + qkv_weight = qkv_module.weight + if type(qkv_module).__name__ == "Conv1D": + if not hasattr(qkv_module, "input_quantizer") and not hasattr( + qkv_module, "output_quantizer" + ): + # For unquantized nn.Conv1D, the weights are transposed compared with the nn.Linear + qkv_weight = qkv_weight.T + + # Handle the case that num_kv_heads/num_attention_heads is the first dimension of QKV. + # This logic covers MQA and GQA as well. + keep_channel_order = not model_metadata_config.get("head_is_first_dim", False) + + q_weight, k_weight, v_weight = _split_fused_qkv_weight_and_scaling( + qkv_weight, + num_heads, + num_kv_heads, + training_tp, + keep_channel_order, + ) + qkv_activation_scaling_factor = get_activation_scaling_factor(qkv_module) + q_activation_scaling_factor = qkv_activation_scaling_factor + k_activation_scaling_factor = qkv_activation_scaling_factor + v_activation_scaling_factor = qkv_activation_scaling_factor + + qkv_weight_scaling_factor = get_weight_scaling_factor(qkv_module) + + if qkv_weight_scaling_factor is not None and qkv_weight_scaling_factor.numel() != 1: + # SQ and AWQ case + ( + q_weight_scaling_factor, + k_weight_scaling_factor, + v_weight_scaling_factor, + ) = _split_fused_qkv_weight_and_scaling( + qkv_weight_scaling_factor, + num_heads, + num_kv_heads, + training_tp, + keep_channel_order, + ) + else: + q_weight_scaling_factor = qkv_weight_scaling_factor + k_weight_scaling_factor = qkv_weight_scaling_factor + v_weight_scaling_factor = qkv_weight_scaling_factor + + # bias + if qkv_module.bias is not None: + q_bias, k_bias, v_bias = _split_fused_qkv_weight_and_scaling( + qkv_module.bias, + num_heads, + num_kv_heads, + training_tp, + keep_channel_order, + ) + + q_weight_scaling_factor_2 = k_weight_scaling_factor_2 = v_weight_scaling_factor_2 = ( + get_weight_scaling_factor_2(qkv_module) + ) + + q_prequant_scaling_factor = k_prequant_scaling_factor = v_prequant_scaling_factor = ( + get_prequant_scaling_factor(qkv_module) + ) + + q_quantization = k_quantization = v_quantization = get_quantization_format(qkv_module) + + elif len(qkv_modules) == 3: + preprocess_linear_fusion(qkv_modules) + # Separate QKV layers + q_weight = qkv_modules[0].weight + q_activation_scaling_factor = get_activation_scaling_factor(qkv_modules[0]) + q_weight_scaling_factor = get_weight_scaling_factor(qkv_modules[0]) + q_quantization = get_quantization_format(qkv_modules[0]) + k_weight = qkv_modules[1].weight + k_activation_scaling_factor = get_activation_scaling_factor(qkv_modules[1]) + k_weight_scaling_factor = get_weight_scaling_factor(qkv_modules[1]) + k_quantization = get_quantization_format(qkv_modules[1]) + v_weight = qkv_modules[2].weight + v_activation_scaling_factor = get_activation_scaling_factor(qkv_modules[2]) + v_weight_scaling_factor = get_weight_scaling_factor(qkv_modules[2]) + v_quantization = get_quantization_format(qkv_modules[2]) + + q_weight_scaling_factor_2 = get_weight_scaling_factor_2(qkv_modules[0]) + k_weight_scaling_factor_2 = get_weight_scaling_factor_2(qkv_modules[1]) + v_weight_scaling_factor_2 = get_weight_scaling_factor_2(qkv_modules[2]) + + q_prequant_scaling_factor = get_prequant_scaling_factor(qkv_modules[0]) + k_prequant_scaling_factor = get_prequant_scaling_factor(qkv_modules[1]) + v_prequant_scaling_factor = get_prequant_scaling_factor(qkv_modules[2]) + + if hasattr(qkv_modules[0], "bias"): + q_bias = qkv_modules[0].bias + + if hasattr(qkv_modules[1], "bias"): + k_bias = qkv_modules[1].bias + + if hasattr(qkv_modules[2], "bias"): + v_bias = qkv_modules[2].bias + + else: + raise NotImplementedError(f"QKV modules format {qkv_modules} not supported") + + # derive num_kv_heads + head_size = q_weight.size(0) // num_heads + num_kv_heads = ext_config.num_kv_heads or k_weight.size(0) // head_size + if tp_size > num_kv_heads: + if any( + map( + lambda t: t is not None and t.numel() > 1, + [ + k_activation_scaling_factor, + k_weight_scaling_factor_2, + k_bias, + v_activation_scaling_factor, + v_weight_scaling_factor_2, + v_bias, + ], + ) + ): + # TODO(oargov): handle cases with biases / scales + raise NotImplementedError( + "Duplicating KV heads for KV with non-scalar scales and/or biases is not supported" + ) + + # duplicte kv heads as needed + k_weight = dup_kv_weight(k_weight, head_size, num_kv_heads, tp_size) + v_weight = dup_kv_weight(v_weight, head_size, num_kv_heads, tp_size) + if k_weight_scaling_factor is not None and k_weight_scaling_factor.numel() > 1: + if len(k_weight_scaling_factor.shape) == 1: + raise NotImplementedError( + "Duplicating KV heads per-channel scales is not supported" + ) + k_weight_scaling_factor = dup_kv_weight( + k_weight_scaling_factor, head_size, num_kv_heads, tp_size + ) + if v_weight_scaling_factor is not None and v_weight_scaling_factor.numel() > 1: + if len(v_weight_scaling_factor.shape) == 1: + raise NotImplementedError( + "Duplicating KV heads per-channel scales is not supported" + ) + v_weight_scaling_factor = dup_kv_weight( + v_weight_scaling_factor, head_size, num_kv_heads, tp_size + ) + + config.q = LinearConfig(linear_type=LINEAR_COLUMN) + config.q.weight = q_weight + config.q.bias = q_bias if q_bias is not None else None + config.q.activation_scaling_factor = q_activation_scaling_factor + config.q.weights_scaling_factor = q_weight_scaling_factor + config.q.weights_scaling_factor_2 = q_weight_scaling_factor_2 + config.q.prequant_scaling_factor = q_prequant_scaling_factor + config.q.awq_block_size = block_size + config.q.quantization = q_quantization + + config.k = LinearConfig(linear_type=LINEAR_COLUMN) + config.k.weight = k_weight + config.k.bias = k_bias if k_bias is not None else None + config.k.activation_scaling_factor = k_activation_scaling_factor + config.k.weights_scaling_factor = k_weight_scaling_factor + config.k.weights_scaling_factor_2 = k_weight_scaling_factor_2 + config.k.prequant_scaling_factor = k_prequant_scaling_factor + config.k.awq_block_size = block_size + config.k.quantization = k_quantization + + config.v = LinearConfig(linear_type=LINEAR_COLUMN) + config.v.weight = v_weight + config.v.bias = v_bias if v_bias is not None else None + config.v.activation_scaling_factor = v_activation_scaling_factor + config.v.weights_scaling_factor = v_weight_scaling_factor + config.v.weights_scaling_factor_2 = v_weight_scaling_factor_2 + config.v.prequant_scaling_factor = v_prequant_scaling_factor + config.v.awq_block_size = block_size + config.v.quantization = v_quantization + + if not ext_config.attention_head_size: + ext_config.attention_head_size = config.q.weight.shape[0] * training_tp // num_heads + + return config + + +def build_linear_config(module: nn.Module, linear_type: str) -> LinearConfig: + """Builds the linear config for the module.""" + if has_mcore and not isinstance(module, MegatronModule): + # Check only for HF model, not Mcore model + assert is_linear(module) + + torch_weight = module.weight + + if "NormHead" in type(module).__name__: + torch_weight = torch.nn.functional.normalize(torch_weight) + elif "Conv1D" in type(module).__name__ and not ( + hasattr(module, "input_quantizer") or hasattr(module, "output_quantizer") + ): + # Transpose Conv1D weights to linear unless it has been transposed by the quantization. + torch_weight = torch_weight.T + + weight = torch_weight + + config = LinearConfig(linear_type=linear_type) + config.weight = weight + + if hasattr(module, "bias") and module.bias is not None: + config.bias = module.bias + + config.activation_scaling_factor = get_activation_scaling_factor(module) + config.weights_scaling_factor = get_weight_scaling_factor(module) + config.weights_scaling_factor_2 = get_weight_scaling_factor_2(module) + config.prequant_scaling_factor = get_prequant_scaling_factor(module) + config.awq_block_size = get_weight_block_size(module) + config.quantization = get_quantization_format(module) + return config + + +def build_fused_linear_config(modules: list[nn.Module], linear_type: str) -> LinearConfig: + """Returns a fused linear config from a list of modules.""" + assert linear_type == LINEAR_COLUMN, "Only support column fuse" + + preprocess_linear_fusion(modules) + + config = build_linear_config(modules[0], linear_type=linear_type) + config.weight = torch.cat([module.weight for module in modules], dim=0) + + if config.weights_scaling_factor.numel() != 1: + config.weights_scaling_factor = torch.cat( + [get_weight_scaling_factor(module) for module in modules], dim=0 + ) + + return config + + +def build_attention_config( + module: nn.Module, + model_metadata_config, + ext_config: DecoderLayerConfig = None, + tp_size: int = 1, +) -> AttentionConfig: + """Builds the attention config from the module.""" + assert is_attention(module) + + config = AttentionConfig() + if hasattr(module, "rotary_dim"): + config.rotary_dim = module.rotary_dim + if hasattr(module, "clip_qkv"): + config.clip_qkv = module.clip_qkv + + qkv_modules = [] + q = None + k = None + v = None + for name, layer in module.named_children(): + if is_linear(layer): + if _is_qkv(name): + qkv_modules.append(layer) + elif "q" in name: + q = layer + elif "k" in name: + k = layer + elif "v" in name: + v = layer + else: + # The dense layer + config.dense = build_linear_config(layer, LINEAR_ROW) + elif is_layernorm(layer): + if "q" in name.lower(): + config.q_layernorm = build_layernorm_config(layer) + elif "k" in name.lower(): + config.k_layernorm = build_layernorm_config(layer) + else: + raise NotImplementedError(f"{name}: {layer} not recognized") + elif "model_type" in model_metadata_config: + if model_metadata_config["model_type"] == "t5": + config.rel_attn_table = RelativeAttentionTableConfig(weight=layer.weight.T) + + if not qkv_modules: + assert q + assert k + assert v + qkv_modules = [q, k, v] + + config.qkv = build_qkv(qkv_modules, model_metadata_config, ext_config, tp_size=tp_size) + + config.kv_cache_scaling_factor = get_kv_cache_scaling_factor(qkv_modules) + if config.kv_cache_scaling_factor is not None: + config.kv_cache_dtype = get_kv_cache_dtype(qkv_modules) + + return config + + +def _is_qkv(name) -> bool: + return all([k in name for k in ["q", "k", "v"]]) or "W_pack" in name or "c_attn" in name + + +def _get_hidden_act(act_func) -> str: + """Returns the name of the hidden activation functon based on ACT2FN.""" + if isinstance(act_func, str): + return act_func + + # Falcon activation, "nn.GELU" is equivalent to "gelu" in ACT2FN + if isinstance(act_func, nn.GELU): + return "gelu" + + if hasattr(act_func, "func") and act_func.func == nn.functional.gelu: + return "gelu" + + for name, func in ACT2FN.items(): + # TRT LLM uses "squared-relu" activation keyword. + if name == "relu2": + name = "squared-relu" + if isinstance(func, tuple): + if isinstance(act_func, func[0]): + return name + elif isinstance(act_func, func): + return name + + return act_func.__name__ + + +def build_mlp_config( + module: nn.Module, + decoder_type, + hidden_act: Optional[str] = None, + merge_gate_fc: bool = False, +) -> MLPConfig: + """Builds the MLP config for the module.""" + assert is_mlp(module) + + config = MLPConfig() + + def _split_gate_from_fc(decoder_type, module, fc_name, fc_layer): + if ( + "ColumnParallelLinear" in type(fc_layer).__name__ + and hasattr(module.config, "gated_linear_unit") + and module.config.gated_linear_unit + ): + return True + + if decoder_type != "gpt": + return False + + if "dense_h_to_4h" in fc_name and "dense_h_to_4h_2" not in fc_name: + return True + + return False + + # TODO: We may want to refactor these keywords/model mapping + fc_keywords = set( + [ + "c_fc", # gpt2 + "fc_in", # gptj + "gate_proj", # llama, baichuan, recurrentgemma, deepseek + "dense_h_to_4h", # falcon, chatglm, bloom + "linear_fc1", + "w2", # qwen + "fc1", # phi, gemma + "gate_up_proj", # phi + "wi_0", # t5 + "wi", # t5 + "c_fc_0", # exaone + ] + ) + proj_keywords = set( + [ + "c_proj", # gpt2, qwen, exaone + "fc_out", # gptj + "dense_4h_to_h", # falcon, chatglm, bloom + "4h_to_h", + "down_proj", # llama, baichuan, mpt, phi, recurrentgemma, nemotron, deepseek + "linear_fc2", + "proj", + "fc2", # phi, gemma + "wo", # t5 + ] + ) + gate_keywords = set( + [ + "up_proj", # llama, baichuan, recurrentgemma, deepseek + "dense_h_to_4h_2", + "w1", # qwen + "wi_1", # t5 + "c_fc_1", # exaone + ] + ) + + # Arctic (llama-based MoE, decoder_type is "llama") has MLP keyword conflicts with Qwen + # Arctic's residual MLP use w1 for fc, w2 for proj, w3 for gate + if type(module).__name__ in ["ArcticMLP", "InternLM2MLP"]: + fc_keywords.discard("w2") + gate_keywords.discard("w1") + fc_keywords.add("w1") + proj_keywords.add("w2") + gate_keywords.add("w3") + + if decoder_type == "mpt": + fc_keywords.add("up_proj") + gate_keywords.discard("up_proj") + + if type(module).__name__ in [ + "TLGv4MLP", + "Phi3SmallMLP", + "NemotronMLP", + ]: # for TLGv4ForCausalLM + fc_keywords.add("up_proj") + gate_keywords.discard("up_proj") + + fc_linear: nn.Module = None + gate_linear: nn.Module = None + proj_linear: nn.Module = None + for name, layer in module.named_children(): + if is_linear(layer): + if any([keyword == name for keyword in fc_keywords]): + fc_linear = layer + elif any([keyword == name for keyword in gate_keywords]): + gate_linear = layer + elif any([keyword == name for keyword in proj_keywords]): + proj_linear = layer + + if merge_gate_fc: + assert gate_linear is not None and fc_linear is not None + preprocess_linear_fusion([fc_linear, gate_linear]) + + if fc_linear is not None: + weight_quantizer = None + if hasattr(fc_linear, "weight_quantizer"): + weight_quantizer = fc_linear.weight_quantizer + if isinstance(weight_quantizer, SequentialQuantizer): + weight_quantizer = weight_quantizer[0] + + # swap fused fc and gate + if decoder_type in ["chatglm", "phi3"]: + weights = torch.chunk(fc_linear.weight, 2, dim=0) + weights = (weights[1], weights[0]) + fc_linear.weight.data = torch.cat(weights, dim=0) + + if ( + weight_quantizer is not None + and weight_quantizer.is_enabled + and weight_quantizer.amax.numel() != 1 + ): + amax_chunks = torch.chunk(weight_quantizer.amax, 2, dim=0) + weight_quantizer.amax = torch.cat([amax_chunks[1], amax_chunks[0]], dim=0) + + split_gate = _split_gate_from_fc(decoder_type, module, name, fc_linear) + if merge_gate_fc: + config.fc = build_fused_linear_config([gate_linear, fc_linear], LINEAR_COLUMN) + gate_linear = None + + elif split_gate: + # We have to split the gate from the fc + weights = torch.chunk(fc_linear.weight, 2, dim=0) + weight_scaling_factor = get_weight_scaling_factor(fc_linear) + + weight_scaling_factors = None + + if weight_scaling_factor is not None and weight_scaling_factor.numel() != 1: + # for Int8 SQ case, we split the weight scaling factor into two parts. + weight_scaling_factors = torch.chunk(weight_scaling_factor, 2, dim=0) + + config.fc = build_linear_config(fc_linear) + config.gate = build_linear_config(fc_linear) + config.fc.weight = weights[0] + config.gate.weight = weights[1] + if weight_scaling_factors is not None: + config.fc.weights_scaling_factor = weight_scaling_factors[0] + config.gate.weights_scaling_factor = weight_scaling_factors[1] + else: + config.fc = build_linear_config(fc_linear, LINEAR_COLUMN) + + if proj_linear is not None: + config.proj = build_linear_config(proj_linear, LINEAR_ROW) + + if gate_linear is not None: + config.gate = build_linear_config(gate_linear, LINEAR_COLUMN) + + assert config.proj is not None and config.fc is not None, "proj or fc can not be found" + + # Override hidden_act based on decoder_type + if decoder_type in ["bloom", "glm"]: + hidden_act = "gelu" + if decoder_type == "phi3": + hidden_act = "swiglu" + + if hidden_act is None: + if hasattr(module, "activation"): + hidden_act = module.activation + elif hasattr(module, "activation_func"): + # MCore activation_func can be swiglu (gated silu) or squared_relu. + hidden_act = module.activation_func.__name__.replace("_", "-") + if hidden_act in ["glu", "silu"]: + hidden_act = "swiglu" if decoder_type == "gpt" else "silu" + else: + for act in ["act", "act_fn", "activation_fn"]: + if hasattr(module, act): + hidden_act = _get_hidden_act(getattr(module, act)).split("_")[0] + break + + if hidden_act is None: + raise NotImplementedError(f"{module} not supported.") + + config.hidden_act = hidden_act + return config + + +def _get_mixtral_expert(experts: nn.Module, export_id: int, linear_name: str): + # Mixtral experts layout is: + # experts[0]: + # w1 + # w2 + # w3 + # experts[1]: + # w1 + # w2 + # w3 + # ... + return getattr(experts[export_id], linear_name) + + +def _get_dbrx_expert(experts: nn.Module, export_id: int, linear_name: str): + # DBRX experts layout is: + # experts: + # w1[0] + # w1[1] + # ... + # w2[0] + # w2[1] + # ... + # v1[0] + # v1[1] + # ... + return getattr(experts, linear_name)[export_id] + + +def _build_stacked_linear(experts: nn.Module, module_name, linear_type, num_experts, expert_getter): + config = LinearConfig(linear_type=linear_type) + + # weights + config.weight = torch.stack( + [expert_getter(experts, i, module_name).weight for i in range(num_experts)] + ) + + # bias + first_module = expert_getter(experts, 0, module_name) + if hasattr(first_module, "bias") and first_module.bias is not None: + raise ValueError("Unexpected bias tensors inside MOE modules.") + + # scaling factors + def get_stacked_scaling_factors(experts, get_function, module_name): + expert_0_scaling_factor = get_function(expert_getter(experts, 0, module_name)) + + if expert_0_scaling_factor is None: + return None + + dtype = expert_0_scaling_factor.dtype + scaling_factors = [ + get_function(expert_getter(experts, i, module_name)) for i in range(num_experts) + ] + + if dtype == torch.float8_e4m3fn: + scaling_factors = [sf.view(torch.uint8) for sf in scaling_factors] + return torch.stack(scaling_factors).view(torch.float8_e4m3fn) + + return torch.stack(scaling_factors) + + config.activation_scaling_factor = get_stacked_scaling_factors( + experts, get_activation_scaling_factor, module_name + ) + # The moe plugin only supports a single activation scaling factor for all experts + if config.activation_scaling_factor is not None: + config.activation_scaling_factor = config.activation_scaling_factor.max().unsqueeze(0) + config.weights_scaling_factor = get_stacked_scaling_factors( + experts, get_weight_scaling_factor, module_name + ) + config.weights_scaling_factor_2 = get_stacked_scaling_factors( + experts, get_weight_scaling_factor_2, module_name + ) + config.prequant_scaling_factor = get_stacked_scaling_factors( + experts, get_prequant_scaling_factor, module_name + ) + config.awq_block_size = get_weight_block_size(expert_getter(experts, 0, module_name)) + config.quantization = get_quantization_format(experts) + + return config + + +def build_stacked_experts( + experts: nn.Module, + linear_names: list[str], + num_experts, + expert_getter, +): + """Builds the experts_weight_1 and experts_weight_2 configs for the experts.""" + # Resmooth all experts + preprocess_linear_fusion( + [expert_getter(experts, i, linear_names[2]) for i in range(num_experts)], + resmooth_only=True, + ) + + w1w3_names = [linear_names[0]] + if len(linear_names) == 3: + w1w3_names.append(linear_names[2]) + + preprocess_linear_fusion( + [ + expert_getter(experts, i, module_name) + for i in range(num_experts) + for module_name in w1w3_names + ], + resmooth_only=True, + ) + + # Pre-fuse W1 and W3 + if len(linear_names) == 3: + for i in range(num_experts): + preprocess_linear_fusion( + [ + expert_getter(experts, i, module_name) + for module_name in [linear_names[0], linear_names[2]] + ] + ) + + experts_weight_1 = _build_stacked_linear( + experts, linear_names[0], LINEAR_COLUMN, num_experts, expert_getter + ) + experts_weight_2 = _build_stacked_linear( + experts, linear_names[1], LINEAR_ROW, num_experts, expert_getter + ) + if len(linear_names) > 2: + # Only for HF model, as Mcore model only has two fc layers in MoE + experts_weight_3 = _build_stacked_linear( + experts, linear_names[2], LINEAR_COLUMN, num_experts, expert_getter + ) + + # Concat w1 and w3 into w1 + experts_weight_1.weight = torch.concat( + [experts_weight_3.weight, experts_weight_1.weight], dim=-2 + ) + + if experts_weight_1.weights_scaling_factor is not None: + if experts_weight_1.weights_scaling_factor.numel() != num_experts: + experts_weight_1.weights_scaling_factor = view_as_float8_e4m3fn_if_needed( + torch.concat( + [ + view_as_uint8_if_needed(experts_weight_3.weights_scaling_factor), + view_as_uint8_if_needed(experts_weight_1.weights_scaling_factor), + ], + dim=-2, + ) + ) + else: + assert torch.equal( + experts_weight_1.weights_scaling_factor, + experts_weight_3.weights_scaling_factor, + ) + + if experts_weight_1.activation_scaling_factor is not None: + assert torch.equal( + experts_weight_1.activation_scaling_factor, + experts_weight_3.activation_scaling_factor, + ) + + if experts_weight_1.weights_scaling_factor_2 is not None: + assert torch.equal( + experts_weight_1.weights_scaling_factor_2, + experts_weight_3.weights_scaling_factor_2, + ) + + # Explicitly move weight to CPU to reduce GPU memory requirement. + experts_weight_1.weight = experts_weight_1.weight.cpu() + experts_weight_2.weight = experts_weight_2.weight.cpu() + return experts_weight_1, experts_weight_2 + + +def build_moe_config(module: nn.Module, decoder_type) -> MOEConfig: + """Builds the MOE config for the module.""" + assert is_moe(module) + assert decoder_type in ["llama", "dbrx", "phi3", "deepseek"] + + config = MOEConfig() + + # Router: TRT-LLM uses fp32 for router to keep precision + if decoder_type in ["llama", "phi3"]: + if has_mcore and isinstance(module, MegatronModule): + # For Mcore model + config.router = build_linear_config(module.router, LINEAR_ROW) + else: + # For huggingface model + config.router = build_linear_config(module.gate, LINEAR_ROW) + elif decoder_type == "dbrx": + config.router = build_linear_config(module.router.layer, LINEAR_ROW) + elif decoder_type == "deepseek": + # Build a linear on the fly. + router_linear = nn.Linear(module.gate.gating_dim, module.gate.n_routed_experts, bias=None) + router_linear.weight = module.gate.weight + config.router = build_linear_config(router_linear, LINEAR_ROW) + # Add shared config + preprocess_linear_fusion([module.shared_experts.gate_proj, module.shared_experts.up_proj]) + config.shared_expert = build_mlp_config( + module.shared_experts, decoder_type, merge_gate_fc=True + ) + else: + raise NotImplementedError(f"{decoder_type} not supported") + + config.router.weight = config.router.weight.type(torch.float) + + # Experts + experts = ExpertConfig() + if decoder_type in ["llama", "phi3"]: + if has_mcore and isinstance(module, MegatronModule): + # For Mcore model + experts.fc, experts.proj = build_stacked_experts( + module.experts.local_experts, + ["linear_fc1", "linear_fc2"], + len(module.experts.local_experts), + _get_mixtral_expert, + ) + # For Mcore model, experts.fc.weight needs to be flipped along axis = 1 + mid_point = experts.fc.weight.shape[1] // 2 + experts.fc.weight = torch.cat( + [ + experts.fc.weight[:, mid_point:, :], + experts.fc.weight[:, :mid_point, :], + ], + dim=1, + ) + else: + # For Huggingface model + experts.fc, experts.proj = build_stacked_experts( + module.experts, + ["w1", "w2", "w3"], + len(module.experts), + _get_mixtral_expert, + ) + elif decoder_type == "dbrx": + experts.fc, experts.proj = build_stacked_experts( + module.experts.mlp, + ["w1_linear", "w2_linear", "v1_linear"], + len(module.experts.mlp.w1_linear), + _get_dbrx_expert, + ) + elif decoder_type == "deepseek": + experts.fc, experts.proj = build_stacked_experts( + module.experts, + ["gate_proj", "down_proj", "up_proj"], + len(module.experts), + _get_mixtral_expert, + ) + else: + raise NotImplementedError(f"{decoder_type} not supported") + + config.experts = experts + + # activation for mixtral and dbrx + config.hidden_act = "swiglu" + + return config + + +def build_conv_config(module: nn.Module) -> ConvConfig: + """Builds the conv config for this module.""" + return ConvConfig( + weight=module.weight.unsqueeze(dim=-1), + bias=module.bias, + ) + + +def is_recurrent(module: nn.Module) -> bool: + """Returns whether the module is a recurrent layer.""" + module_name = type(module).__name__ + return "RecurrentBlock" in module_name + + +def build_recurrent_config(module: nn.Module): + """Builds the recurrent config for this module.""" + assert is_recurrent(module) + + config = RecurrentConfig() + config.linear_y = build_linear_config(module.linear_y, linear_type=LINEAR_COLUMN) + # Separate bias from linear_y to y_bias as guided by TRT LLM. + config.y_bias = config.linear_y.bias + config.linear_y.bias = None + + config.linear_x = build_linear_config(module.linear_x, linear_type=LINEAR_COLUMN) + config.linear_out = build_linear_config(module.linear_out, linear_type=LINEAR_ROW) + config.conv1d = build_conv_config(module.conv_1d) + config.rg_lru = RgLruConfig() + + config.rg_lru.input_gate = LinearConfig( + linear_type=LINEAR_GROUP, + weight=module.rg_lru.input_gate_weight, + bias=module.rg_lru.input_gate_bias, + quantization=get_quantization_format(module), + ) + + config.rg_lru.recurrent_gate = LinearConfig( + linear_type=LINEAR_GROUP, + weight=module.rg_lru.recurrent_gate_weight, + bias=module.rg_lru.recurrent_gate_bias, + quantization=get_quantization_format(module), + ) + + config.rg_lru.recurrent_param = module.rg_lru.recurrent_param + return config + + +def _set_layer_config_from_metaconfig(layer_config, metaconfig): + for keys, name in HF_CONFIG_MAP: + for k in keys: + if k in metaconfig: + setattr(layer_config, name, metaconfig[k]) + for keys, name in MCORE_CONFIG_MAP: + for k in keys: + if k in metaconfig: + setattr(layer_config, name, metaconfig[k]) + + # 2048 is the default TRT LLM max_position_embeddings + if layer_config.max_position_embeddings == 0: + layer_config.max_position_embeddings = 2048 + + for k in ["_name_or_path"]: + if k in metaconfig: + model_path_name = metaconfig[k].replace("-", "_") + chatglm_version = "" + + for n in ["glm_4", "chatglm3", "chatglm2", "chatglm"]: + if n in model_path_name: + chatglm_version = n.replace("_", "") # glm_4 -> glm + break + + layer_config.chatglm_version = chatglm_version + + # For Falcon variants do not provide new_decoder_architecture, but we can infer it from the model_type + if "model_type" in metaconfig: + if metaconfig["model_type"] == "RefinedWeb": + # Case 1. Falcon-40B / Falcon-40B-instruct + # https://huggingface.co/tiiuae/falcon-40b/blob/main/layer_config.json + layer_config.new_decoder_architecture = True + elif metaconfig["model_type"] == "RefinedWebModel": + # Case 2. Falcon-7B / Falcon-7B-instruct + # https://huggingface.co/tiiuae/falcon-7b/blob/main/layer_config.json + layer_config.new_decoder_architecture = False + + # For Falcon variants, they might not specify the number of kv heads with MQA models, e.g., 7b + if ( + not layer_config.new_decoder_architecture + and "multi_query" in metaconfig + and metaconfig["multi_query"] + ): + layer_config.num_kv_heads = 1 + + # For Phi3 + for k in ["rope_scaling"]: + if k in metaconfig and metaconfig[k] is not None: + for k2 in ["short_factor"]: + if k2 in metaconfig[k]: + layer_config.longrope_scaling_short_factors = metaconfig[k][k2] + for k2 in ["long_factor"]: + if k2 in metaconfig[k]: + layer_config.longrope_scaling_long_factors = metaconfig[k][k2] + # For Phi3 small + for k2 in ["short_mscale"]: + if k2 in metaconfig[k]: + layer_config.longrope_short_mscale = metaconfig[k][k2] + for k2 in ["long_mscale"]: + if k2 in metaconfig[k]: + layer_config.longrope_long_mscale = metaconfig[k][k2] + + +def build_decoder_config( + module: nn.Module, + model_metadata_config, + decoder_type: str, + tp_size: int = 1, +) -> DecoderLayerConfig: + """Builds the full decoder config from the module.""" + quantization = get_quantization_format(module) + config = DecoderLayerConfig(decoder_type=decoder_type) + # Supporting different attention layer config in MCoreGPTModel. If per layer config + # exists, override the global config. + if hasattr(module, "self_attention"): + if hasattr(module.self_attention, "config"): + if hasattr(module.self_attention.config, "num_attention_heads"): + config.num_attention_heads = module.self_attention.config.num_attention_heads + if hasattr(module.self_attention.config, "kv_channels"): + config.attention_head_size = module.self_attention.config.kv_channels + if hasattr(module.self_attention.config, "num_query_groups"): + config.num_kv_heads = module.self_attention.config.num_query_groups + + # Set all config fields in modelopt from HF config + _set_layer_config_from_metaconfig(config, model_metadata_config) + + if type(module).__name__ == "DbrxBlock": + # Flatten DBRX attention and ffn + module_layers = {} + module_layers.update(dict(getattr(module, "norm_attn_norm").named_children())) + module_layers.update({"ffn": module.ffn}) + elif decoder_type in ["t5"]: + # Combine two modules (T5LayerSelfAttention, T5LayerFF) / three modules + # ((T5LayerSelfAttention, T5LayerCrossAttention, T5LayerFF)) of T5 model + # (depending on whether it's encoder / decoder) into one decoder layer + combined_module = nn.ModuleList() + for sub_module in module: + for layer in sub_module.children(): + combined_module.append(layer) + module_layers = dict(combined_module.named_children()) + else: + module_layers = dict(module.named_children()) + if decoder_type in ["exaone"]: + module_layers.update({"attn": module_layers["attn"].attention}) + + for name, layer in module_layers.items(): + # We assume input_layernorm should be before the post_layernorm in decoder block layout, + # and residual_layernorm could be after post_layernorm + if is_layernorm(layer): + layernorm_config = build_layernorm_config(layer) + # Special attributes naming for T5 model (Encoder-Decoder model) + if decoder_type in ["t5"]: + _update_encoder_decoder_layernorm_config( + model_metadata_config, config, layernorm_config + ) + # For all decoder only models + else: + if name in ["ln_mlp"]: + config.mlp_layernorm = layernorm_config + elif config.decoder_type == "gemma2" and "post_attention_layernorm" in name: + config.post_layernorm = layernorm_config + elif config.decoder_type == "gemma2" and "pre_feedforward_layernorm" in name: + config.pre_feedforward_layernorm = layernorm_config + elif config.decoder_type == "gemma2" and "post_feedforward_layernorm" in name: + config.post_feedforward_layernorm = layernorm_config + elif config.input_layernorm is None: + config.input_layernorm = layernorm_config + elif config.post_layernorm is None: + config.post_layernorm = layernorm_config + else: + assert model_metadata_config[ + "parallel_attn_mlp_res" + ], "Unexpected layernorm in a layer" + config.residual_layernorm = layernorm_config + + elif is_attention(layer): + # For models where a linear may replace the attention/MLP module (e.g. Deci models) + if is_linear(layer): + config.attn_replacing_linear = build_linear_config(layer.linear_attn, "column") + else: + if decoder_type in ["bloom", "falcon", "phi3small", "internlm"]: + model_metadata_config["head_is_first_dim"] = True + attention_config = build_attention_config( + layer, model_metadata_config, config, tp_size=tp_size + ) + # For T5 decoder + if decoder_type in ["t5"] and model_metadata_config["enc_dec"] == "dec": + # We assume self_attention should be before the cross_attention in decoder block layout + if config.self_attention is None: + config.self_attention = attention_config + else: + config.cross_attention = attention_config + elif decoder_type == "mllama": + if "cross" in type(layer).__name__.lower(): + config.cross_attention = attention_config + else: + config.self_attention = attention_config + else: + config.attention = attention_config + + elif is_moe(layer): + if quantization not in [QUANTIZATION_NONE, QUANTIZATION_FP8]: + warn(f"TensorRT-LLM may not support MOE quantization {quantization}") + config.mlp = build_moe_config(layer, decoder_type) + + elif is_recurrent(layer): + config.recurrent = build_recurrent_config(layer) + + # We assume MoE layer should be before the residual MLP layer in decoder block layout, + # so MLP layer after a MoE layer is treated as a residual MLP layer + elif is_mlp(layer): + if config.mlp is None: + # For models where a linear may replace the attention/MLP module (e.g. Deci models) + if is_linear(layer): + config.mlp_replacing_linear = build_linear_config(layer.linear_mlp, "column") + else: + config.mlp = build_mlp_config( + layer, + decoder_type, + hidden_act=model_metadata_config.get("hidden_act", None), + ) + else: + assert model_metadata_config["parallel_attn_mlp_res"], "Unexpected mlp in a layer" + config.residual_mlp = build_mlp_config( + layer, + decoder_type, + hidden_act=model_metadata_config.get("hidden_act", None), + ) + + config.gate_attn = getattr(module, "cross_attn_attn_gate", None) + config.gate_ffwd = getattr(module, "cross_attn_mlp_gate", None) + + config = _move_input_layernorm_for_noop_attention(config) + return config + + +def build_medusa_heads_config( + model: Optional[nn.Module], +) -> Optional[list[MedusaHeadConfig]]: + """Build a list of MedusaHeadConfig if exists. + + Following TensorRT-LLM's Medusa implementation, all Medusa heads (num_medusa_heads) should be + placed inside a 'torch.nn.ModuleList' with attribute name 'medsua_heads'. A Medusa head composes + an additional 'lm_head' (vocab_size, hidden_size) and a list (num_medusa_layers) of Medusa layer + (LinearActConfig). The only supported hidden_act for the layer is 'silu'. All Linear layers are + column-parallel. + """ + # medusa_heads: Optional[nn.Module] = None + + def get_medusa_heads(model: nn.Module) -> Optional[nn.Module]: + """Return the MedusaHead is exists.""" + # MCore GPTModel impl + if hasattr(model, "medusa_heads"): + return model.medusa_heads + return None + + medusa_heads = get_medusa_heads(model) + + if medusa_heads is None: + return None + + configs = [] + for medusa_head in medusa_heads: + config = MedusaHeadConfig() + layer_configs = [] + if isinstance(medusa_head, torch.nn.Sequential): + # In the HF MedusaModel, nn.Sequencetial is used. + config.lm_head = build_linear_config(medusa_head[-1], LINEAR_COLUMN) + for layer in medusa_head[0:-1]: + layer_config = LinearActConfig() + layer_config.linear = build_linear_config(layer.linear, LINEAR_COLUMN) + # NOTE: only silu is supported now. + layer_config.hidden_act = "silu" + layer_configs.append(layer_config) + + else: + # In the Megatron, we have MedusaHead define and lm_head is a submodule. + config.lm_head = build_linear_config(medusa_head.lm_head, LINEAR_COLUMN) + for layer in medusa_head.medusa_layers: + layer_config = LinearActConfig() + layer_config.linear = build_linear_config(layer.linear, LINEAR_COLUMN) + # NOTE: only silu is supported now. + layer_config.hidden_act = "silu" + layer_configs.append(layer_config) + config.medusa_layers = layer_configs + configs.append(config) + return configs + + +def _split_fused_qkv_weight_and_scaling( + weight: torch.Tensor, + num_heads: int, + num_kv_heads: Optional[int] = None, + training_tp: int = 1, + keep_channel_order: bool = False, +): + """Reorder the qkv weight for spliting QKV weights. + + 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 contigous. + 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. + if keep_channel_order + HF : [q1, ..., qh, k1, ..., kh, v1, ..., vh] + TRT-LLM: [q1, ..., qh, k1, ..., kh, v1, ..., vh] + """ + # 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 + weight_dim = weight.dim() + assert 1 <= weight_dim and weight_dim <= 2 + + qkv_in = weight.shape[-1] if weight_dim > 1 else 1 + + num_kv_heads = num_kv_heads if num_kv_heads else num_heads + assert ( + num_heads % num_kv_heads == 0 + ), f"num_heads({num_heads}) must be divisible by 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 + num_kv_heads_single_tp = max(num_kv_heads // training_tp, 1) + size_per_head = weight.shape[0] // num_kv_heads_single_tp // num_group_heads + + if keep_channel_order: + kv_size = num_kv_heads * size_per_head + q_w = weight[: -2 * kv_size, ...] + k_w = weight[-2 * kv_size : -1 * kv_size, ...] + v_w = weight[-1 * kv_size :, ...] + else: + # Split Q/K/V weights + weight = weight.reshape(num_kv_heads_single_tp, num_group_heads, size_per_head, qkv_in) + q_w = weight[:, :-2, ...] # (nKV, num_heads // nKV, size_per_head, qkv_in) + k_w = weight[:, -2:-1, ...] # (nKV, 1, size_per_head, qkv_in) + v_w = weight[:, -1:, ...] # (nKV, 1, size_per_head, qkv_in) + + q_w = q_w.reshape(-1, qkv_in) + k_w = k_w.reshape(-1, qkv_in) + v_w = v_w.reshape(-1, qkv_in) + + if weight_dim == 1: + q_w = q_w.reshape(-1) + k_w = k_w.reshape(-1) + v_w = v_w.reshape(-1) + + return q_w, k_w, v_w + + +def _update_encoder_decoder_layernorm_config(model_metadata_config, config, layernorm_config): + # For encoder of Encoder-Decoder model + if model_metadata_config["enc_dec"] == "enc": + if config.attention_layernorm is None: + config.attention_layernorm = layernorm_config + else: + config.mlp_layernorm = layernorm_config + # For decoder of Encoder-Decoder model + else: + if config.self_attention_layernorm is None: + config.self_attention_layernorm = layernorm_config + elif config.cross_attention_layernorm is None: + config.cross_attention_layernorm = layernorm_config + else: + config.mlp_layernorm = layernorm_config + + +def _move_input_layernorm_for_noop_attention( + complete_decoder_config: DecoderLayerConfig, +) -> DecoderLayerConfig: + """For models with decoder blocks that skip the attention or recurrent, the only layer norm is AFTER the MLP. + + We build DecoderLayerConfigs with the assumption that there are both norms. This function fixes + this assumption and moved the input_layernorm to be a post_layernorm if necessary. + """ + if all( + config is None + for config in [ + complete_decoder_config.attention, + complete_decoder_config.self_attention, + complete_decoder_config.cross_attention, + complete_decoder_config.attn_replacing_linear, + complete_decoder_config.recurrent, + ] + ): + assert ( + complete_decoder_config.post_layernorm is None + ), "Should not have 2 layer norms with no attention" + complete_decoder_config.post_layernorm = complete_decoder_config.input_layernorm + complete_decoder_config.input_layernorm = None + + return complete_decoder_config + + +def update_experts_avg_prequant_scale(experts: nn.Module): + """Registers experts_pre_quant_scale attribute of each expert with average pre_quant_scale amongst experts.""" + """In NVFP4_AWQ and INT4_AWQ all the experts share prequant_scaling_factor. """ + experts_linear_names = get_experts_linear_names(experts) + if "mixtral" in type(experts).__name__.lower(): + get_func = _get_mixtral_expert + num_experts = len(experts.experts) + experts = experts.experts + elif "dbrx" in type(experts).__name__.lower(): + get_func = _get_dbrx_expert + num_experts = len(getattr(experts.experts.mlp, experts_linear_names[0])) + experts = experts.experts.mlp + else: + raise NotImplementedError("MoE model not supported") + + for linear_name in experts_linear_names: + # Check if pre_quant_scale exists + assert hasattr( + get_func(experts, 0, linear_name).input_quantizer, "pre_quant_scale" + ), "Layer does not have attribute pre_quant_scale" + experts_avg_pre_quant_scale = torch.mean( + torch.stack( + [ + get_func(experts, i, linear_name).input_quantizer.pre_quant_scale + for i in range(num_experts) + ] + ), + dim=0, + ) + # Register a new experts_pre_quant_scale attribute of each expert with average pre_quant_scale amongst experts + for i in range(num_experts): + get_func( + experts, i, linear_name + ).input_quantizer.experts_avg_pre_quant_scale = experts_avg_pre_quant_scale.clone() + + +def get_experts_linear_names(model: torch.nn.Module): + """Returns linear layer names based on decoder type for MoE models.""" + if "mixtral" in type(model).__name__.lower(): + return ["w1", "w2", "w3"] + elif "dbrx" in type(model).__name__.lower(): + return ["w1_linear", "w2_linear", "v1_linear"] + else: + raise NotImplementedError("MoE model not supported") diff --git a/modelopt/torch/export/mcore_config_map.py b/modelopt/torch/export/mcore_config_map.py new file mode 100644 index 0000000000..24c86bdf2b --- /dev/null +++ b/modelopt/torch/export/mcore_config_map.py @@ -0,0 +1,24 @@ +# 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. + +"""Define the config mapping between Mcore and modelopt.""" + +# the map is a list of tuples, and each tuple has two elements +# first element: a list of possible fields on Mcore +# second element: a name of the layer config field inside modelopt +MCORE_CONFIG_MAP = [ + (["moe_router_topk"], "moe_top_k"), + (["num_moe_experts"], "moe_num_experts"), +] diff --git a/modelopt/torch/export/mcore_hf_export_map.py b/modelopt/torch/export/mcore_hf_export_map.py new file mode 100644 index 0000000000..b3b7eb5dc1 --- /dev/null +++ b/modelopt/torch/export/mcore_hf_export_map.py @@ -0,0 +1,163 @@ +# 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. + + +"""Custom mapping from Megatron Core models to their Hugging Face counter part.""" + +import copy +from typing import Any, Dict + + +class CustomModuleMapping: + """A custom module mapping from Megatron Core to its HF counter part.""" + + def __init__( + self, func_name: str = "", target_name_or_prefix: str = "", func_kwargs: Dict[str, Any] = {} + ): + """Create a custom module mapping.""" + self.func_name = func_name + self.target_name_or_prefix = target_name_or_prefix + self.func_kwargs = func_kwargs + + +gpt_model_export: Dict[str, CustomModuleMapping] = { + # MCore GPTModel + "word_embeddings": CustomModuleMapping("name_remapping", "embedding.word_embeddings."), + "input_layernorm": CustomModuleMapping("name_remapping", "decoder.layers.{}.input_layernorm."), + "linear_qkv": CustomModuleMapping( + "name_remapping", + "decoder.layers.{}.self_attention.linear_qkv.", + {"skip_output_scale": False}, + ), + "linear_proj": CustomModuleMapping( + "name_remapping", "decoder.layers.{}.self_attention.linear_proj." + ), + "pre_mlp_layernorm": CustomModuleMapping( + "name_remapping", "decoder.layers.{}.pre_mlp_layernorm." + ), + "linear_fc1": CustomModuleMapping("name_remapping", "decoder.layers.{}.mlp.linear_fc1."), + "linear_fc2": CustomModuleMapping("name_remapping", "decoder.layers.{}.mlp.linear_fc2."), + "final_layernorm": CustomModuleMapping("name_remapping", "decoder.final_layernorm."), + "output_layer": CustomModuleMapping("name_remapping", "output_layer."), + # ModelOpt MedusaGPTModel support + "medusa_heads.lm_head": CustomModuleMapping("name_remapping", "medusa_heads.{}.1."), + "medusa_heads.medusa_layers.linear": CustomModuleMapping( + "name_remapping", "medusa_heads.{}.{}.linear." + ), + # ModelOpt EagleGPTModel support + "eagle_module.fc": CustomModuleMapping("name_remapping", "eagle_module.fc."), + "eagle_module.input_layernorm": CustomModuleMapping( + "name_remapping", "eagle_module.decoder.layers.{}.input_layernorm." + ), + "eagle_module.linear_qkv": CustomModuleMapping( + "name_remapping", + "eagle_module.decoder.layers.{}.self_attention.linear_qkv.", + {"skip_output_scale": False}, + ), + "eagle_module.linear_proj": CustomModuleMapping( + "name_remapping", "eagle_module.decoder.layers.{}.self_attention.linear_proj." + ), + "eagle_module.pre_mlp_layernorm": CustomModuleMapping( + "name_remapping", "eagle_module.decoder.layers.{}.pre_mlp_layernorm." + ), + "eagle_module.linear_fc1": CustomModuleMapping( + "name_remapping", "eagle_module.decoder.layers.{}.mlp.linear_fc1." + ), + "eagle_module.linear_fc2": CustomModuleMapping( + "name_remapping", "eagle_module.decoder.layers.{}.mlp.linear_fc2." + ), +} + + +llama_causal_lm_export: Dict[str, CustomModuleMapping] = { + "word_embeddings": CustomModuleMapping("name_remapping", "model.embed_tokens."), + "input_layernorm": CustomModuleMapping("name_remapping", "model.layers.{}.input_layernorm."), + "linear_qkv": CustomModuleMapping( + "qkv_slicing", + "model.layers.{}.self_attn.", + ), + "linear_proj": CustomModuleMapping("name_remapping", "model.layers.{}.self_attn.o_proj."), + "pre_mlp_layernorm": CustomModuleMapping( + "name_remapping", "model.layers.{}.post_attention_layernorm." + ), + "linear_fc1": CustomModuleMapping("gated_mlp_slicing", "model.layers.{}.mlp."), + "linear_fc2": CustomModuleMapping("name_remapping", "model.layers.{}.mlp.down_proj."), + "final_layernorm": CustomModuleMapping("name_remapping", "model.norm."), + "output_layer": CustomModuleMapping("name_remapping", "lm_head."), + # MedusaForCausalLM support + "medusa_heads.lm_head": CustomModuleMapping("name_remapping", "medusa_heads.{}.1."), + "medusa_heads.medusa_layers.linear": CustomModuleMapping( + "name_remapping", "medusa_heads.{}.{}.linear." + ), + # EagleForCausalLM support + "eagle_module.fc": CustomModuleMapping("name_remapping", "eagle_module.fc."), + "eagle_module.input_layernorm": CustomModuleMapping( + "name_remapping", "eagle_module.layers.{}.input_layernorm." + ), + "eagle_module.linear_qkv": CustomModuleMapping( + "qkv_slicing", + "eagle_module.layers.{}.self_attn.", + ), + "eagle_module.linear_proj": CustomModuleMapping( + "name_remapping", "eagle_module.layers.{}.self_attn.o_proj." + ), + "eagle_module.pre_mlp_layernorm": CustomModuleMapping( + "name_remapping", "eagle_module.layers.{}.post_attention_layernorm." + ), + "eagle_module.linear_fc1": CustomModuleMapping( + "gated_mlp_slicing", "eagle_module.layers.{}.mlp." + ), + "eagle_module.linear_fc2": CustomModuleMapping( + "name_remapping", "eagle_module.layers.{}.mlp.down_proj." + ), +} + +# Example on adding a new CausalLM. +nemotron_causal_lm_export_delta: Dict[str, CustomModuleMapping] = { + # NemotronForCausalLM is using square-relu where no gated handle is needed. + "linear_fc1": CustomModuleMapping("name_remapping", "model.layers.{}.mlp.up_proj."), + # EagleForCausalLM support + "eagle_module.linear_fc1": CustomModuleMapping( + "name_remapping", "eagle_module.layers.{}.mlp.up_proj." + ), +} +# Copy the Llama rule book and overwrite the delta +nemotron_causal_lm_export = copy.deepcopy(llama_causal_lm_export) +nemotron_causal_lm_export.update(nemotron_causal_lm_export_delta) + +all_mcore_hf_export_mapping: Dict[str, Any] = { + "GPTModel": gpt_model_export, + "LlamaForCausalLM": llama_causal_lm_export, + "NemotronForCausalLM": nemotron_causal_lm_export, +} + + +llama_causal_lm_import: Dict[str, CustomModuleMapping] = { + "word_embeddings": CustomModuleMapping("name_remapping", "model.embed_tokens."), + "input_layernorm": CustomModuleMapping("name_remapping", "model.layers.{}.input_layernorm."), + "linear_qkv": CustomModuleMapping("qkv_merging", "model.layers.{}.self_attn."), + "linear_proj": CustomModuleMapping("name_remapping", "model.layers.{}.self_attn.o_proj."), + "pre_mlp_layernorm": CustomModuleMapping( + "name_remapping", "model.layers.{}.post_attention_layernorm." + ), + "linear_fc1": CustomModuleMapping("gated_mlp_merging", "model.layers.{}.mlp."), + "linear_fc2": CustomModuleMapping("name_remapping", "model.layers.{}.mlp.down_proj."), + "final_layernorm": CustomModuleMapping("name_remapping", "model.norm."), + "output_layer": CustomModuleMapping("name_remapping", "lm_head."), +} + +all_mcore_hf_import_mapping: Dict[str, Any] = { + "LlamaForCausalLM": llama_causal_lm_import, +} diff --git a/modelopt/torch/export/model_config.py b/modelopt/torch/export/model_config.py new file mode 100755 index 0000000000..4c0f946548 --- /dev/null +++ b/modelopt/torch/export/model_config.py @@ -0,0 +1,597 @@ +# 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. + +"""This module defines the model_config format. + +This format can be converted from huggingface, nemo or modelopt-quantized model. +And we will build tensorrt_llm engine from the context saved with this format. +""" + +import math +from dataclasses import dataclass, field +from typing import Optional, Union + +import torch + +from modelopt.torch.quantization.qtensor import NVFP4QTensor + +QUANTIZATION_NONE = None +QUANTIZATION_FP8 = "fp8" +QUANTIZATION_INT8_SQ = "int8_sq" +QUANTIZATION_INT4_AWQ = "int4_awq" +QUANTIZATION_W4A8_AWQ = "w4a8_awq" +QUANTIZATION_NVFP4 = "nvfp4" +QUANTIZATION_NVFP4_AWQ = "nvfp4_awq" + +KV_CACHE_FP8 = "FP8" +KV_CACHE_INT8 = "INT8" + +LINEAR_COLUMN = "column" +LINEAR_ROW = "row" +LINEAR_GROUP = "group" + +# These need to be synced with torch.export.tensorrt_llm_type +LAYERNORM_DEFAULT = "LayerNorm" +LAYERNORM_RMS = "RmsNorm" + + +@dataclass +class EmbeddingConfig: + """The embedding layer config.""" + + weight: torch.Tensor = None + + @property + def local_vocab_size(self): + """Infers the vocab_size from the embedding layer weights shape.""" + return self.weight.shape[0] + + @property + def hidden_size(self): + """Infers the hidden_size from the embedding layer weights shape.""" + return self.weight.shape[1] + + +@dataclass +class LayernormConfig: + """The layernorm layer config.""" + + weight: torch.Tensor = None + bias: torch.Tensor = None + layernorm_type: str = LAYERNORM_DEFAULT + eps: float = 1e-5 + + +@dataclass +class LinearConfig: + """The linear layer config.""" + + quantization: Optional[str] = QUANTIZATION_NONE + linear_type: str = LINEAR_COLUMN + weight: torch.Tensor = None + bias: torch.Tensor = None + activation_scaling_factor: torch.Tensor = None + weights_scaling_factor: torch.Tensor = None + + # For methods like W4A8 AWQ, we have two quantizers for weights + # For W4A8, the first quantizer is for INT4 quantization and the second quantizer is for FP8 quantization + # `weight_scaling_factor_2` is the scaling factor the the second FP8 quantizer + weights_scaling_factor_2: torch.Tensor = None + + prequant_scaling_factor: torch.Tensor = None + awq_block_size: int = 0 + + def __del__(self): + del self.weight + del self.bias + del self.activation_scaling_factor + del self.weights_scaling_factor + del self.weights_scaling_factor_2 + del self.prequant_scaling_factor + + +@dataclass +class LinearActConfig: + """The linear + activation layer config.""" + + linear: LinearConfig = None + hidden_act: str = "" + + +@dataclass +class ConvConfig: + """The Conv layer config.""" + + quantization: Optional[str] = QUANTIZATION_NONE + weight: torch.Tensor = None + bias: torch.Tensor = None + + +@dataclass +class QKVConfig: + """The QKV layer config.""" + + q: LinearConfig = None + k: LinearConfig = None + v: LinearConfig = None + + @property + def weight(self): + """The generated linear layer weight. + + The Q, K, V weights are concat together to fit the TensorRT-LLM QKV linear layer. + """ + return torch.cat((self.q.weight, self.k.weight, self.v.weight)) + + @property + def bias(self): + """The generated linear layer bias. + + The Q, K, V bias are concat together to fit the TensorRT-LLM QKV linear layer. + """ + if self.q.bias is None: + assert ( + self.k.bias is None and self.v.bias is None + ), "K and V should have valid bias as Q" + return None + return torch.cat((self.q.bias, self.k.bias, self.v.bias)) + + @property + def activation_scaling_factor(self): + """Returns the merged activation_scaling_factor across Q, K and V. + + The max of the Q, K, V activation scaling factors is returned. + """ + if ( + self.q.activation_scaling_factor is None + or self.k.activation_scaling_factor is None + or self.v.activation_scaling_factor is None + ): + return None + + return ( + torch.stack( + [ + self.q.activation_scaling_factor, + self.k.activation_scaling_factor, + self.v.activation_scaling_factor, + ] + ) + .max(dim=0) + .values + ) + + @property + def weights_scaling_factor(self): + """Returns the merged weights_scaling_factor across Q, K and V. + + If the quantization is FP8, the max of the Q, K, V weight scaling factors is returned. + If the quanitzation is INT8_SQ, the concat value is returned. + """ + if ( + self.q.weights_scaling_factor is None + or self.k.weights_scaling_factor is None + or self.v.weights_scaling_factor is None + ): + return None + + if self.q.weights_scaling_factor.numel() != 1: + # for NVFP4, Int4 AWQ and Int8 SQ case, we concatenate the + # q_weight_scaling_factor, k_weight_scaling_factor, v_weight_scaling_factor + if self.q.weights_scaling_factor.dtype == torch.float8_e4m3fn: + # For NVFP4, we recompute the weights_scaling_factor using max weights_scaling_factor2 + qkv_weights_scaling_factor, _ = NVFP4QTensor.get_weights_scaling_factor( + self.weight, self.awq_block_size, self.weights_scaling_factor_2 + ) + else: + qkv_weights_scaling_factor = torch.cat( + ( + self.q.weights_scaling_factor, + self.k.weights_scaling_factor, + self.v.weights_scaling_factor, + ) + ) + else: + # for FP8 set qkv_weight_scaling_factor to the max of + # q_weight_scaling_factor, k_weight_scaling_factor, v_weight_scaling_factor + qkv_weights_scaling_factor = ( + torch.stack( + [ + self.q.weights_scaling_factor, + self.k.weights_scaling_factor, + self.v.weights_scaling_factor, + ], + ) + .max(dim=0) + .values + ) + return qkv_weights_scaling_factor + + @property + def weights_scaling_factor_2(self): + """Returns the merged weights_scaling_factor_2 across Q, K and V. + + weight_scaling_factor_2 is needed for W4A8 AWQ. + """ + if ( + self.q.weights_scaling_factor_2 is None + or self.k.weights_scaling_factor_2 is None + or self.v.weights_scaling_factor_2 is None + ): + return None + + # For W4A8 AWQ, weight_scaling_factor_2 corresponds to the per-tensor FP8 quantization. + # Hence weight_scaling_factor_2 should be a scalar. + assert self.q.weights_scaling_factor_2.numel() == 1 + + # set qkv_weight_scaling_factor_2 to the max of q,k,v weight_scaling_factor_2 + qkv_weights_scaling_factor_2 = ( + torch.stack( + [ + self.q.weights_scaling_factor_2, + self.k.weights_scaling_factor_2, + self.v.weights_scaling_factor_2, + ] + ) + .max(dim=0) + .values + ) + + return qkv_weights_scaling_factor_2 + + @property + def prequant_scaling_factor(self): + """Returns the merged prequant_scaling_factor across Q, K and V. + + Prequant scaling factors for Q, K, V should be the same. So just return one of them. + """ + if ( + self.q.prequant_scaling_factor is None + or self.k.prequant_scaling_factor is None + or self.v.prequant_scaling_factor is None + ): + return None + + assert torch.equal( + self.q.prequant_scaling_factor, self.k.prequant_scaling_factor + ) and torch.equal( + self.k.prequant_scaling_factor, self.v.prequant_scaling_factor + ), "Prequant scaling factors of Q, K and V should be the same" + return self.q.prequant_scaling_factor + + @property + def awq_block_size(self): + """Returns the awq_block_size of this QKV layer.""" + assert ( + self.q.awq_block_size == self.k.awq_block_size == self.v.awq_block_size + ), "awq_block_size of QKV should be the same." + return self.q.awq_block_size + + +@dataclass +class RelativeAttentionTableConfig: + """Relative attention table config. For splitting purpose.""" + + weight: torch.Tensor = None + + +@dataclass +class AttentionConfig: + """The attention layer config.""" + + # QKV can either be stored as splitted (for easier postprocessing) + # or merged (for TRT LLM export) + qkv: Union[QKVConfig, LinearConfig] = None + dense: LinearConfig = None + kv_cache_scaling_factor: torch.Tensor = None + kv_cache_dtype: Optional[str] = None + + rotary_dim: int = -math.inf + # MPT variants + clip_qkv: float = None + # T5 + rel_attn_table: RelativeAttentionTableConfig = None + + # Mllama + q_layernorm: LayernormConfig = None + k_layernorm: LayernormConfig = None + + +@dataclass +class MLPConfig: + """The MLP layer config.""" + + fc: LinearConfig = None + gate: LinearConfig = None + proj: LinearConfig = None + hidden_act: str = "" + + +@dataclass +class ExpertConfig: + """The Expert config.""" + + # Aligning the naming convesion with TRT-LLM + fc: LinearConfig = None # stacked experts for concatenated w3 and w1 + proj: LinearConfig = None # stacked experts for w2 + + +@dataclass +class RgLruConfig: + """The RG LRU from recurrentgemma.""" + + recurrent_param: torch.Tensor = None + input_gate: LinearConfig = None + recurrent_gate: LinearConfig = None + + +@dataclass +class RecurrentConfig: + """The RecurrentBlock from recurrentgemma.""" + + linear_y: LinearConfig = None + y_bias: torch.Tensor = None + linear_x: LinearConfig = None + linear_out: LinearConfig = None + conv1d: ConvConfig = None + rg_lru: RgLruConfig = None + + +@dataclass +class MOEConfig: + """The Mixture of Expert layer config.""" + + router: LinearConfig = None + experts: ExpertConfig = None + shared_expert: MLPConfig = None # Deepseek MOE + hidden_act: str = "" + + @property + def fc(self): + """Return the fc module from experts.""" + return self.experts.fc + + +@dataclass +class DecoderLayerConfig: + """The decoder layer config.""" + + decoder_type: str = "" + input_layernorm: LayernormConfig = None + mlp_layernorm: LayernormConfig = None + attention: AttentionConfig = None + recurrent: RecurrentConfig = None # recurrent gemma + post_layernorm: LayernormConfig = None + pre_feedforward_layernorm: LayernormConfig = None # gemma 2 + post_feedforward_layernorm: LayernormConfig = None # gemma 2 + mlp: Union[MLPConfig, MOEConfig] = None + + num_attention_heads: int = 0 + # Supporting different attention_head_size per layer. + attention_head_size: int = None + + num_kv_heads: int = 0 + max_position_embeddings: int = 0 + rotary_pct: float = 1.0 + + # Falcon and Baichuan variants + use_alibi: bool = False + new_decoder_architecture: bool = False + parallel_attention: bool = False + + # chatglm variants + apply_residual_connection_post_layernorm: bool = False + use_cache: bool = True + chatglm_version: str = "" + rope_ratio: float = 1.0 + + # Qwen config + seq_length: int = 0 + qwen_type: str = "" + + # Qwen and CodeLlama + rotary_base: int = 0 + + # Phi + partial_rotary_factor: float = 0 + + # Phi3 + original_max_position_embeddings: int = 0 + longrope_scaling_short_factors: list[float] = None + longrope_scaling_long_factors: list[float] = None + + # Phi3 small + mup_attn_multiplier: float = 0 + mup_embedding_multiplier: float = 0 + mup_use_scaling: float = 0 + mup_width_multiplier: float = 0 + blocksparse_block_size: int = 0 + blocksparse_homo_head_pattern: bool = False + blocksparse_num_local_blocks: int = 0 + blocksparse_vertical_stride: int = 0 + dense_attention_every_n_layers: int = 0 + gegelu_limit: float = 0 + longrope_short_mscale: float = 0 + longrope_long_mscale: float = 0 + + # Mixture of Experts + moe_num_experts: int = 0 + moe_top_k: int = 0 + moe_tp_mode: int = 0 + moe_renorm_mode: int = 0 + + # MPT + alibi_bias_max: int = 0 + + # Arctic variants + residual_layernorm: LayernormConfig = None + residual_mlp: MLPConfig = None + + # Recurrent Gemma + rnn_hidden_size: int = 0 + logits_soft_cap: float = 0 + emb_scale_by_sqrt_dim: bool = False + layer_types: list[str] = field(default_factory=list) + + # Deci models + attn_replacing_linear: LinearConfig = None + mlp_replacing_linear: LinearConfig = None + block_config: dict = None + + # Gemma2 + final_logit_softcapping: float = 0 + attn_logit_softcapping: float = 0 + query_pre_attn_scalar: float = 0 + + # DBRX + clip_qkv: int = 0 + + # T5, Mllama + cross_attention: AttentionConfig = None + cross_attention_layernorm: LayernormConfig = None + self_attention: AttentionConfig = None + self_attention_layernorm: LayernormConfig = None + attention_layernorm: LayernormConfig = None + rel_attn_max_distance: int = 0 + rel_attn_num_buckets: int = 0 + + # Llama 3.1 + rope_scaling: dict = None + + # Mllama + cross_attention_layers: dict = None + vision_output_dim: int = 0 + gate_ffwd: torch.Tensor = None + gate_attn: torch.Tensor = None + + # Phi3.5 MOE + sparse_mixer_epsilon: float = 0 + + @property + def hidden_size(self): + """Returns the hidden size of the transformer model.""" + if isinstance(self.mlp, MOEConfig): + # fc.weight for MOE is stacked + if self.mlp.fc.quantization in [QUANTIZATION_NVFP4, QUANTIZATION_NVFP4_AWQ]: + return self.mlp.fc.weight.shape[-1] * 2 + return self.mlp.fc.weight.shape[-1] + else: + k = self.mlp.fc.weight.shape[1] + if self.mlp.fc.quantization in [QUANTIZATION_NVFP4, QUANTIZATION_NVFP4_AWQ]: + return k * 2 + return k + + @property + def ffn_hidden_size_local(self): + """Returns the ffn hidden size of the transformer model.""" + fc = self.mlp.fc + if isinstance(self.mlp, MOEConfig): + # fc in MoE merge fc and gate + k = fc.weight.shape[1] // 2 + else: + k = fc.weight.shape[0] + if fc.quantization not in [QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ]: + return k + return k * 2 + + +@dataclass +class MedusaHeadConfig: + """The decoder layer config.""" + + medusa_layers: list[LinearActConfig] = None + lm_head: LinearConfig = None + + +@dataclass +class ModelConfig: + """The full LLM model config that includes the full information needed for tensorrt_llm engine building. + + This class includes all the fields that tensorrt_llm supports, but not all of the fields are required. + pipeline_parallel > 1 is only supported for TensorRT-LLM checkpoint. + """ + + # Global metadata + architecture: str = "" + quantization: str = QUANTIZATION_NONE + dtype: str = "float16" + vocab_size: int = 0 + + # Parallel metadata + rank: int = 0 + tensor_parallel: int = 1 + pipeline_parallel: int = 1 + + # Model structure and weights + vocab_embedding: EmbeddingConfig = None + position_embedding: EmbeddingConfig = None + block_embedding: EmbeddingConfig = None + ln_embed: LayernormConfig = None + layers: list[DecoderLayerConfig] = field(default_factory=list) + + ln_f: LayernormConfig = None + + lm_head: LinearConfig = None + share_embedding_table: bool = False + + medusa_heads: list[MedusaHeadConfig] = None + num_medusa_heads: int = 0 + num_medusa_layers: int = 0 + + # To differentiate encoder and decoder of Encoder-Deocder model + enc_dec: str = "" + + # For decoder of Encoder-Decoder model that needs encoder information + encoder_hidden_size: int = 0 + encoder_num_heads: int = 0 + encoder_head_size: int = 0 + + @property + def vocab_size_padded(self): + """Returns the padded vocab_size of the model rounds to the tensor_parallel.""" + + def _pad_vocab_size(vocab_size, tp_size): + return int(math.ceil(vocab_size / tp_size) * tp_size) + + return _pad_vocab_size(self.vocab_size, self.tensor_parallel) + + @property + def hidden_size(self): + """Returns the hidden_size of the model.""" + return self.layers[0].hidden_size + + @property + def max_position_embeddings(self): + """Returns the max_position_embedding of the model.""" + return self.layers[0].max_position_embeddings + + @property + def num_attention_heads(self): + """Returns the num_attention_heads of the model.""" + return self.layers[0].num_attention_heads + + @property + def num_kv_heads(self): + """Returns the num_key_value_heads of the model.""" + return ( + self.layers[0].num_kv_heads + if self.layers[0].num_kv_heads is not None and self.layers[0].num_kv_heads > 0 + else self.num_attention_heads + ) + + @property + def hidden_act(self): + """Returns the hidden_act of the model.""" + return self.layers[0].mlp.hidden_act diff --git a/modelopt/torch/export/model_config_export.py b/modelopt/torch/export/model_config_export.py new file mode 100644 index 0000000000..686750a1f2 --- /dev/null +++ b/modelopt/torch/export/model_config_export.py @@ -0,0 +1,549 @@ +# 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. + +"""Code that export optimized models to the TensorRT-LLM checkpoint.""" + +import copy +import json +import math +import tempfile +from dataclasses import asdict +from pathlib import Path +from typing import Any, Iterator, Optional, Union +from warnings import warn + +import torch +import torch.nn as nn +from safetensors.torch import save_file + +from modelopt.torch.utils import distributed as dist +from modelopt.torch.utils import import_plugin + +from .layer_utils import ( + build_decoder_config, + build_embedding_config, + build_layernorm_config, + build_linear_config, + build_medusa_heads_config, + check_model_compatibility, + get_dtype, + get_transformer_layers, + is_decoder_list, + is_embedding, + is_layernorm, + is_linear, +) +from .model_config import QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ, ModelConfig +from .model_config_utils import ( + merge_qkv, + model_config_to_dict, + naive_quantization, + pack_linear_weights, + split_config_and_weights, +) +from .postprocess import ( + check_weight_shape_valid, + pad_embedding_lm_head, + postprocess_model_config, + postprocess_tensors, + update_lm_head_quantization, +) +from .quant_utils import get_quantization_format, process_layer_quant_config +from .tensorrt_llm_utils import ( + convert_to_tensorrt_llm_config, + is_tensorrt_llm_0_8_or_9, + prepare_enc_dec_decoder_layer, + prepare_enc_dec_export_dir, +) + +has_mcore = False +with import_plugin("megatron", verbose=False): + from megatron.core.models.gpt import GPTModel as MCoreGPTModel + from megatron.core.transformer.module import MegatronModule + + has_mcore = True + +__all__ = ["export_tensorrt_llm_checkpoint", "torch_to_tensorrt_llm_checkpoint"] + + +def torch_to_tensorrt_llm_checkpoint( + model: nn.Module, + decoder_type: str, + dtype: Optional[torch.dtype] = None, + inference_tensor_parallel: int = 0, + inference_pipeline_parallel: int = 1, + naive_fp8_quantization: bool = False, + workspace_path: Optional[Union[Path, str]] = None, +) -> Iterator[tuple[dict[str, Any], dict[str, torch.Tensor], dict[str, Any]]]: + """Converts the torch model to the TensorRT-LLM checkpoint per GPU rank. + + TensorRT-LLM checkpoint is the LLM model format that can be used by the TensorRT-LLM build API. + for the engine building process. + https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/architecture/checkpoint.md + + Args: + model: the torch model. + decoder_type: the type of the decoder, e.g. gpt, gptj, llama. + dtype: the weights data type to export the unquantized layers or the default model data type if None. + inference_tensor_parallel: The target inference time tensor parallel. + We will merge or split the calibration tensor parallelism to inference. + Default is 0, meaning using the calibration without manual config merge or split. + inference_pipeline_parallel: The target inference time pipeline parallel. + We will merge or split the calibration pipeline parallelism to inference. + Default is 1, meaning no pipeline parallelism. + naive_fp8_quantization: Quantize the model naively to FP8 without calibration. + All scaling factors are set to 1. + workspace_path: the path to the NFS directory for postprocess cross rank communication. + + Yields: + A tuple of + tensorrt_llm_config: A dict that maps to the ``PretrainedConfig`` in TensorRT-LLM. + https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/models/modeling_utils.py + weights: A dict that stores all model weights and scaling factors for each rank. + per_layer_quantization: A dict that contains layer-wise quantization information for all quantized layers + for mixed_precision, empty dictionary otherwise. + """ + if dtype is None: + dtype = get_dtype(model) + + if dtype not in [torch.float16, torch.bfloat16]: + warn( + f"dtype {dtype} not fully compatible with TensorRT-LLM optimizations, Default to float16." + ) + dtype = torch.float16 + + architecture = "" + hf_config = None + if has_mcore and isinstance(model, MegatronModule): + if not isinstance(model, MCoreGPTModel): + raise ValueError("Only megatron.core.models.gpt.GPTModel is supported!") + # MCoreGPTModel.config has type TransformerConfig + # + # We choose to deepcopy here since TransformerConfig deserialization is sensitive to + # additional attributes. + model_metadata_config = copy.deepcopy(model.config.__dict__) + vocab_size = model.vocab_size + model_metadata_config["max_position_embeddings"] = model.max_position_embeddings + model_metadata_config["rotary_percent"] = model.rotary_percent + model_metadata_config["rotary_base"] = getattr(model, "rotary_base", 10000) + rope_scaling = getattr(model, "rotary_scaling", False) + if rope_scaling: + model_metadata_config["rope_scaling"] = {"type": "llama3"} + elif hasattr(model, "config"): + # Huggingface models + model_metadata_config = model.config.__dict__ + vocab_size = model.config.vocab_size + hf_config = model.config + architecture = model.config.architectures[0] + + # For Baichuan 13B, we check if alibi is used with the alibi_mask property. + if hasattr(model, "model") and hasattr(model.model, "alibi_mask"): + model_metadata_config["alibi"] = True + + # For MPT, DBRX + for config_key in ["attn_config", "ffn_config"]: + config_value = model_metadata_config.get(config_key, None) + if config_value: + model_metadata_config.update( + config_value if isinstance(config_value, dict) else config_value.to_dict() + ) + + elif hasattr(model, "cfg"): + # NeMo MegatronGPTModel + model_metadata_config = dict(model.cfg) + vocab_size = model.tokenizer.vocab_size + else: + raise ValueError("Cannot find valid model metadata config in model") + + if "multi_query_group_num" in model_metadata_config.keys(): + if model_metadata_config["multi_query_group_num"] % inference_tensor_parallel != 0: + raise ValueError( + "Cannot divide {} kv_heads into {} gpus".format( + model_metadata_config["multi_query_group_num"], + inference_tensor_parallel, + ) + ) + + training_pipeline_parallel = model_metadata_config.get("pipeline_model_parallel_size", 1) + training_tensor_parallel = dist.size() // training_pipeline_parallel + model_metadata_config["training_pipeline_parallel"] = training_pipeline_parallel + model_metadata_config["training_tensor_parallel"] = training_tensor_parallel + + if "make_vocab_size_divisible_by" in model_metadata_config: + # For some nemo models, the vocab_size is pre-padded. + # We calculate the pre-padded vocab_size with this config: make_vocab_size_divisible_by. + make_vocab_size_divisible_by = model_metadata_config["make_vocab_size_divisible_by"] + make_vocab_size_divisible_by_with_tp = ( + make_vocab_size_divisible_by * training_tensor_parallel + ) + vocab_size = int( + math.ceil(vocab_size / make_vocab_size_divisible_by_with_tp) + * make_vocab_size_divisible_by_with_tp + ) + print( + f"the new vocab_size is updated: {vocab_size}, make_vocab_size_divisible_by" + f" {make_vocab_size_divisible_by}, training_tensor_parallel" + f" {training_tensor_parallel}." + ) + + models = [model] + if decoder_type in ["t5"]: + model_lm_head = model.lm_head + # For T5 model with encoder and decoder, we process the checkpoint separately. + models = [model.encoder, model.decoder] + + for model in models: + transformer_layers = get_transformer_layers(model) + if training_pipeline_parallel == 1: + compatible, has_position_embedding, has_embedding_layernorm = check_model_compatibility( + transformer_layers + ) + else: + # For Megatron models with more than one PP, + # we skip the compatibility check as not all ranks have the full model. + # For Megatron Core GPTModel, both nemotron and llama do not have position embedding + # nor embedding layernorm. + compatible = len(transformer_layers) > 0 + has_position_embedding = False + has_embedding_layernorm = False + assert compatible, "The model is not supported" + + config = ModelConfig( + architecture=architecture, + dtype=str(dtype).split(".")[1], + rank=dist.rank(), + tensor_parallel=training_tensor_parallel, + pipeline_parallel=training_pipeline_parallel, + vocab_size=vocab_size, + ) + + # For Encoder-Decoder Model like T5 + if decoder_type in ["t5"]: + if model.is_decoder is False: + config.enc_dec = "enc" + model_metadata_config["enc_dec"] = "enc" + else: + config.enc_dec = "dec" + model_metadata_config["enc_dec"] = "dec" + + # GLM has a different handling of word_embeddings. + if decoder_type == "glm": + model = model.glm + transformer_layers.insert(0, model.word_embeddings) + + # Build the full model_config dict layer by layer. + for module in transformer_layers: + if is_embedding(module): + if config.vocab_embedding is None: + # We assume the first embedding in the list the vocab_embedding. + + normalization_constant = 1 + # Normalize vocab embedding for gemma. + if ( + decoder_type == "gemma" and is_tensorrt_llm_0_8_or_9() + ) or decoder_type == "recurrentgemma": + normalization_constant = model_metadata_config["hidden_size"] ** 0.5 + + config.vocab_embedding = build_embedding_config( + module, normalization_constant=normalization_constant + ) + elif has_position_embedding and config.position_embedding is None: + config.position_embedding = build_embedding_config(module) + elif decoder_type == "glm": + config.block_embedding = build_embedding_config(module) + elif is_decoder_list(module): + layers = [] + for idx, layer in enumerate(module.children()): + # Special process due to T5 model structure's specialty + if decoder_type in ["t5"]: + layer = layer.layer + layer_config = build_decoder_config( + layer, + model_metadata_config, + decoder_type, + tp_size=inference_tensor_parallel, + ) + if decoder_type in ["deci"]: + block_config = model_metadata_config["block_configs"][idx] + layer_config.block_config = asdict(block_config) + # Special process for each decoder layer of Encoder-Decoder Model + if decoder_type in ["t5"]: + prepare_enc_dec_decoder_layer( + layer_config, + model.config, + model_metadata_config["enc_dec"], + layers, + ) + layers.append(layer_config) + config.layers = layers + elif is_layernorm(module): + if has_embedding_layernorm and config.ln_embed is None: + # Assume embedding_layernorm is placed before the ln_f. + config.ln_embed = build_layernorm_config(module) + else: + config.ln_f = build_layernorm_config(module) + elif is_linear(module): + if model_metadata_config.get("share_embeddings_and_output_weights", False): + # NeMo/MCore models with shared embeddings - for example Gemma - + # the model head weight is None so we just skip processing + config.share_embedding_table = True + continue + # TRT LLM forces the embedding table to be shared for the following models. + force_share_embedding_table = decoder_type in ["gemma", "gemma2"] + if force_share_embedding_table and torch.equal( + module.weight, config.vocab_embedding.weight + ): + config.share_embedding_table = True + else: + # This will update lm_head quantization config according to constraints from TRT-LLM + update_lm_head_quantization(config, module, inference_pipeline_parallel) + config.lm_head = build_linear_config(module, "column") + + # For decoder of Encoder-Decoder model, it needs some encoder information + if decoder_type in ["t5"]: + if model_metadata_config["enc_dec"] == "dec": + config.encoder_hidden_size = models[0].config.d_model + config.encoder_head_size = models[0].config.d_kv + config.encoder_num_heads = models[0].config.num_heads + + # For the training time PP, not all ranks will have the lm_head layer. + if config.lm_head is None: + if decoder_type in ["t5"]: + config.share_embedding_table = False + if model_metadata_config["enc_dec"] == "dec": + config.lm_head = build_linear_config(model_lm_head, "column") + elif training_pipeline_parallel == 1: + # Models that share weights for lm_head and vocab_embedding + assert decoder_type in [ + "mpt", + "gpt2", + "gemma", + "gemma2", + "glm", + ], f"lm_head not available for decoder {decoder_type}" + config.share_embedding_table = True + + # Handle Medusa Heads + # TODO (chenhany): post-processing is not implemented yet + config.medusa_heads = build_medusa_heads_config(model) + + config.quantization = get_quantization_format(model) + + if config.quantization in [QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ]: + if config.vocab_size % 64 != 0: + # TODO: Check if this works for Mixtral + assert ( + training_tensor_parallel == 1 + ), "We do not support padding for training time TP" + print("Padding vocab_embedding and lm_head for AWQ weights export") + pad_embedding_lm_head(config) + if hf_config is not None: + hf_config.vocab_size = config.vocab_size + + check_weight_shape_valid( + config, + inference_tensor_parallel, + training_tensor_parallel, + ) + + # If inference_tensor_parallel or inference_pipeline_parallel is different from world_size, + # we try to merge or split the model configs based on the rank selected. + if ( + inference_tensor_parallel > 0 + or inference_pipeline_parallel > 0 + or training_pipeline_parallel > 1 + ): + model_configs = postprocess_model_config( + config, + inference_tensor_parallel, + inference_pipeline_parallel, + training_pipeline_parallel=training_pipeline_parallel, + workspace_path=workspace_path, + ) + else: + model_configs = [config] + + for model_config in model_configs: + assert model_config.rank >= 0, "Invalid model_config, postprocess_model_config fails." + + if not model_config.quantization and naive_fp8_quantization: + naive_quantization(model_config) + + merge_qkv(model_config) + pack_linear_weights(model_config) + + weights = {} + # Layer config dict holds quantization format of each layer. + # It also holds awq_block_size information for applicable layers. + layer_config_dict = {} + model_config_dict = model_config_to_dict(model_config) + # We split the weights from model_config and save them separately as two files. + split_config_and_weights(model_config_dict, weights, "transformer", layer_config_dict) + # Process per layer quantization config dict + per_layer_quantization = process_layer_quant_config(layer_config_dict) + # We only export the json once across ranks as all jsons should be the same except for the rank. + tensorrt_llm_config = convert_to_tensorrt_llm_config( + model_config, weights.keys(), hf_config=hf_config + ) + + # Postprocess the tensors in the model_config. + # Exporting the safetensors also allows the tensor to be a view. + postprocess_tensors( + weights, + dtype=dtype, + force_cpu=True, + force_contiguous=True, + force_non_view=False, + ) + + yield tensorrt_llm_config, weights, per_layer_quantization + + +def export_tensorrt_llm_checkpoint( + model: nn.Module, + decoder_type: str, + dtype: Optional[torch.dtype] = None, + export_dir: Union[Path, str] = tempfile.gettempdir(), + inference_tensor_parallel: int = 0, + inference_pipeline_parallel: int = 1, + naive_fp8_quantization: bool = False, + use_nfs_workspace: bool = False, +): + """Exports the torch model to the TensorRT-LLM checkpoint and save to the export_dir. + + Args: + model: the torch model. + decoder_type: the type of the decoder, e.g. gpt, gptj, llama. + dtype: the weights data type to export the unquantized layers or the default model data type if None. + export_dir: the target export path. + inference_tensor_parallel: The target inference time tensor parallel. + We will merge or split the calibration tensor parallelism to inference. + Default is 0, meaning using the calibration without manual config merge or split. + inference_pipeline_parallel: The target inference time pipeline parallel. + We will merge or split the calibration pipeline parallelism to inference. + Default is 1, meaning no pipeline parallelism. + inference_pipeline_parallel: The target inference time pipeline parallel. + naive_fp8_quantization: Quantize the model naively to FP8 without calibration. + All scaling factors are set to 1. + use_nfs_workspace: if True, the an NFS workspace will be created under the export_dir and + used as a shared memory for cross process/node communication. + + For tensorrt_llm deployment, save the representation under ``export_dir``. + We will save the model_config as two files: + + * ``.json``: The nested dict that maps to the ``PretrainedConfig`` in TensorRT-LLM. + https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/models/modeling_utils.py. + * ``.safetensors``: The file for the list of weights as safetensors. Unique for each rank. + """ + export_dir = Path(export_dir) + export_root = export_dir + export_dir.mkdir(parents=True, exist_ok=True) + # Create a NFS workspace under the export folder which is also assumed to be NFS. + workspace_path = None + if use_nfs_workspace: + workspace_path = export_dir.joinpath("workspace") + workspace_path.mkdir(parents=True, exist_ok=True) + + try: + exclude_modules = set() + for ( + tensorrt_llm_config, + weights, + per_layer_quantization, + ) in torch_to_tensorrt_llm_checkpoint( + model=model, + decoder_type=decoder_type, + dtype=dtype, + inference_tensor_parallel=inference_tensor_parallel, + inference_pipeline_parallel=inference_pipeline_parallel, + naive_fp8_quantization=naive_fp8_quantization, + workspace_path=workspace_path, + ): + rank = tensorrt_llm_config["rank"] + world_size = tensorrt_llm_config["mapping"]["world_size"] + if tensorrt_llm_config["quantization"]: + excluded_modules_rank = tensorrt_llm_config["quantization"].get( + "exclude_modules", {} + ) + if excluded_modules_rank: + exclude_modules.update(excluded_modules_rank) + # For T5 model + if decoder_type in ["t5"]: + export_dir = prepare_enc_dec_export_dir(tensorrt_llm_config, export_root) + export_dir = Path(export_dir) + export_dir.mkdir(parents=True, exist_ok=True) + + if rank == world_size - 1: + # We only export the json once across ranks as all jsons should be the same except for the rank. + # If auto_quant is used, save per layer quantization information in quant_cfg.json + if per_layer_quantization: + # Update auto quant related information for quantization.json export + per_layer_quantization["kv_cache_quant_algo"] = tensorrt_llm_config[ + "quantization" + ]["kv_cache_quant_algo"] + + # Update auto quant related information for config.json export + # We remove group_size, has_zero_point, exclude_modules and pre_quant_scale information from config + tensorrt_llm_config["quantization"] = { + k: per_layer_quantization[k] for k in ("quant_algo", "kv_cache_quant_algo") + } + + with open(export_dir / "quant_cfg.json", "w") as f: + json.dump(per_layer_quantization, f, indent=4) + else: + # Excluded modules information is only included in non auto_quant case + tensorrt_llm_config["quantization"]["exclude_modules"] = list(exclude_modules) + + with open(export_dir / "config.json", "w") as f: + json.dump(tensorrt_llm_config, f, indent=4) + + # Hacky implementation for T5 for now + new_weights = {} + if decoder_type == "t5": + for key in weights.keys(): + if key == "transformer.vocab_embedding.weight": + new_key = "embedding.vocab_embedding.weight" + elif key.startswith("transformer.layers"): + # For encoder + if tensorrt_llm_config["architecture"] == "EncoderModel": + new_key = key.replace("transformer.layers", "encoder_layers") + # For decoder + else: + new_key = key.replace("transformer.layers", "decoder_layers") + elif key == "transformer.ln_f.weight": + new_key = "final_layernorm.weight" + elif key == "lm_head.weight": + new_key = key + if key.endswith("rel_attn_table.weight"): + new_key = new_key.replace("rel_attn_table.weight", "rel_attn_table") + if key.find(".0.") > 0: + # stores additional the relative attention for pre-compute feature in TRTLLM + new_weights["rel_attn_table"] = weights[key].clone() + new_weights[new_key] = weights[key] + weights = new_weights + # End of hacky implementation for T5 for now + + weights_path = export_dir / f"rank{rank}.safetensors" + save_file(weights, weights_path) + + except Exception as e: + fallback_model_path = export_dir / f"modelopt_model.{dist.rank()}.pth" + torch.save(model.state_dict(), fallback_model_path) + warn( + "Cannot export model to the model_config. The modelopt-optimized model state_dict" + f" (including the quantization factors) is saved to {fallback_model_path} using" + " torch.save for further inspection." + ) + raise e diff --git a/modelopt/torch/export/model_config_utils.py b/modelopt/torch/export/model_config_utils.py new file mode 100755 index 0000000000..5b02ec9b6d --- /dev/null +++ b/modelopt/torch/export/model_config_utils.py @@ -0,0 +1,401 @@ +# 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. + +"""Common utils for the ModelConfig.""" + +import dataclasses +import math +from typing import Optional, Union, get_args, get_origin + +import numpy as np +import torch + +from .model_config import ( + KV_CACHE_FP8, + QUANTIZATION_FP8, + QUANTIZATION_INT4_AWQ, + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + DecoderLayerConfig, + LayernormConfig, + LinearConfig, + MLPConfig, + ModelConfig, + MOEConfig, + QKVConfig, +) +from .quant_utils import to_quantized_weight + +# numpy doesn't know bfloat16, define abstract binary type instead +np_bfloat16 = np.dtype("V2", metadata={"dtype": "bfloat16"}) + + +def _numpy_to_torch(x): + """Convert numpy array to torch tensor.""" + if isinstance(x, torch.Tensor): + return x + + if x.dtype != np_bfloat16: + return torch.tensor(x) + return torch.tensor(x.view(np.int16)).view(torch.bfloat16) + + +def model_config_to_dict(model_config: ModelConfig) -> dict: + """Converts the instance to a python dict.""" + assert model_config is not None, "model_config is None" + return dataclasses.asdict(model_config) + + +def split_config_and_weights( + config, + weights: dict[str, torch.tensor], + prefix: str = "transformer", + layer_config_dict: dict = {}, +): + """Util function to split the weights or any torch.Tensor in nested config to weights. + + A weight id starts with transformers or lm_head will also be generated to link the original key to the weights dict. + The weights in the weights dict are contiguous. + + layer_config_dict: A dictionary containing layerwise quantization format information and awq_block_size information + when relevant. It is used to export quantization.json for auto_quant checkpoint. + """ + if isinstance(config, dict): + for k, v in config.items(): + if k == "lm_head" and "medusa_heads" not in prefix: + # lm_head is not part of the transformer. + array_key = k + elif k == "experts": + # Omit the 'experts' key that is not in the model name + array_key = prefix + elif k == "medusa_heads": + # medusa_heads is not part of the transformer + array_key = k + else: + array_key = f"{prefix}.{k}" + + # Construct per_layer quantization dictionary, with block size information + if array_key != "transformer.quantization" and ( + "quantization" in array_key or "awq_block_size" in array_key + ): + layer_config_dict[array_key] = v + + if isinstance(v, torch.Tensor): + weights[array_key] = v + config[k] = f"{array_key}" + else: + split_config_and_weights(v, weights, array_key, layer_config_dict) + elif isinstance(config, list): + for i, v in enumerate(config): + array_key = f"{prefix}.{i}" + if isinstance(v, torch.Tensor): + weights[array_key] = v + config[i] = f"{array_key}" + else: + split_config_and_weights(v, weights, array_key, layer_config_dict) + + +def _unified_weights_key(k: str) -> str: + """Try to unify the weights dict key between old npz and the new safetensors format.""" + prefixes = ["transformer.", "_np:"] + for prefix in prefixes: + if k.startswith(prefix): + k = k[len(prefix) :] + + k = k.replace("final_layernorm", "ln_f") + + return k.replace(":", ".") + + +def _restore_model_config(model_config, weights: dict[str, Union[np.ndarray, torch.Tensor]]): + def _is_tensor_key(k): + return isinstance(k, str) and _unified_weights_key(k) in weights + + if isinstance(model_config, dict): + for k, v in model_config.items(): + if _is_tensor_key(v): + model_config[k] = _numpy_to_torch(weights[_unified_weights_key(v)]) + else: + _restore_model_config(v, weights) + if isinstance(model_config, list): + for i, v in enumerate(model_config): + if _is_tensor_key(v): + model_config[i] = _numpy_to_torch(weights[_unified_weights_key(v)]) + else: + _restore_model_config(v, weights) + + +def restore_model_config(model_config, weights: dict[str, Union[np.ndarray, torch.Tensor]]): + """Recursively restores the model_config from json and loads np.ndarray or torch.Tensor weights from weights.""" + unified_key_weights = {} + for k, v in weights.items(): + unified_key_weights[_unified_weights_key(k)] = v + + _restore_model_config(model_config, unified_key_weights) + + +def _from_dict(class_type, data): + """Helper function to load the data as a class_type. class_type must be a dataclass.""" + if data is None: + return None + + # class_type of quantization is Optional[str], which is catergrorized as Union + if class_type != Optional[str] and get_origin(class_type) == Union: + # Handle QKV + if all([key in data for key in ["q", "k", "v"]]): + # splitted qkv case + class_type = QKVConfig + elif all([key in data for key in ["router", "experts"]]): + # moe + class_type = MOEConfig + elif all([key in data for key in ["fc", "gate", "proj"]]): + # mlp + class_type = MLPConfig + else: + # merged qkv case + assert "linear_type" in data, f"{data} is not a valid LinearConfig" + class_type = LinearConfig + + if dataclasses.is_dataclass(class_type): + fieldtypes = {f.name: f.type for f in dataclasses.fields(class_type)} + fields_map = {} + for k, v in data.items(): + if k in fieldtypes: + # We only handle keys available in the fields. + # Deprecated fields in the checkpoint will be ignored. + fields_map[k] = _from_dict(fieldtypes[k], v) + return class_type(**fields_map) + elif get_origin(class_type) is list and dataclasses.is_dataclass(get_args(class_type)[0]): + list_value = [] + for child in data: + child_class_type = get_args(class_type)[0] + list_value.append(_from_dict(child_class_type, child)) + return list_value + else: + return data + + +def model_config_from_dict(d: dict) -> ModelConfig: + """Load a dict to a `ModelConfig` instance.""" + config_type = ModelConfig + + config_type_map = {} + for t in [ModelConfig, DecoderLayerConfig, LayernormConfig, LinearConfig]: + config_type_map[t.__name__] = t + + if "__name__" in d: + config_name = d.pop("__name__") + try: + config_type = config_type_map[config_name] + except Exception as e: + raise NotImplementedError(f"{config_name} not supported") from e + + return _from_dict(config_type, d) + + +def pad_weights(weights, tp_size): + """Returns the padded weights to tp_size.""" + assert len(weights.shape) > 1 + + def _pad_size(original_size, tp_size): + return int(math.ceil(original_size / tp_size) * tp_size) + + original_size = weights.shape[0] + padded_size = _pad_size(original_size, tp_size) + + if original_size != padded_size: + pad_width = padded_size - original_size + return torch.nn.functional.pad(weights, (0, 0, 0, pad_width), "constant", value=0) + return weights + + +def merge_qkv(model_config): + """Merges the qkv fields in model_config from QKVConfig to a single LinearConfig.""" + for decoder_config in model_config.layers: + for attention_key in ["attention", "self_attention", "cross_attention"]: + attention = getattr(decoder_config, attention_key, None) + if attention and isinstance(attention.qkv, QKVConfig): + splitted_qkv = attention.qkv + attention.qkv = LinearConfig() + attention.qkv.weight = splitted_qkv.weight + attention.qkv.bias = splitted_qkv.bias + attention.qkv.activation_scaling_factor = splitted_qkv.activation_scaling_factor + attention.qkv.weights_scaling_factor = splitted_qkv.weights_scaling_factor + attention.qkv.weights_scaling_factor_2 = splitted_qkv.weights_scaling_factor_2 + attention.qkv.prequant_scaling_factor = splitted_qkv.prequant_scaling_factor + attention.qkv.awq_block_size = splitted_qkv.awq_block_size + # Assert q,k,v have same quantization formats before merging + assert ( + splitted_qkv.q.quantization + == splitted_qkv.k.quantization + == splitted_qkv.v.quantization + ), "Quantization formats of q,k,v must be the same." + attention.qkv.quantization = splitted_qkv.q.quantization + + # Collect GPU memory from the deleted tensors + del splitted_qkv + + +def pack_linear_weights(model_config: ModelConfig): + """Packs the quantized linear weights in the model_config to the quantized format.""" + + def _linear_layer_to_quantized_weight(linear_layers): + for linear_layer in linear_layers: + if isinstance(linear_layer, LinearConfig): + # Check if quantization of the layer is None to support auto_quant + if ( + linear_layer.weights_scaling_factor is not None + and linear_layer.quantization is not None + ): + # Quantize on CPU if we are short of GPU memory. + # Using 2x of the tensor size as a threshold. + if linear_layer.weight.is_cuda: + free_mem, _ = torch.cuda.mem_get_info(linear_layer.weight.device) + if ( + free_mem + < 2 + * linear_layer.weight.element_size() + * linear_layer.weight.nelement() + ): + linear_layer.weight = linear_layer.weight.cpu() + + # Save the quantize layer weights to cpu and save gpu memory. + linear_layer.weight = to_quantized_weight( + linear_layer.weight, + linear_layer.weights_scaling_factor, + linear_layer.quantization, + linear_layer.weights_scaling_factor_2, + linear_layer.awq_block_size, + ).cpu() + + linear_layer.weights_scaling_factor = linear_layer.weights_scaling_factor.cpu() + + if not model_config.quantization: + return + + attention_key_list = ["attention", "self_attention", "cross_attention"] + for decoder_config in model_config.layers: + linear_layers = [] + if any([hasattr(decoder_config, attention_key) for attention_key in attention_key_list]): + for attention_key in attention_key_list: + attention = getattr(decoder_config, attention_key, None) + if attention: + linear_layers += [ + attention.qkv, + attention.dense, + ] + if decoder_config.recurrent: + linear_layers = [ + decoder_config.recurrent.linear_y, + decoder_config.recurrent.linear_x, + decoder_config.recurrent.linear_out, + ] + + if isinstance(decoder_config.mlp, MOEConfig): + if model_config.quantization not in [ + QUANTIZATION_FP8, + QUANTIZATION_INT4_AWQ, + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + ]: + raise NotImplementedError( + f"MOE quantization for {model_config.quantization} is not supported yet." + ) + else: + linear_layers.append(decoder_config.mlp.experts.fc) + linear_layers.append(decoder_config.mlp.experts.proj) + elif decoder_config.mlp is not None: + linear_layers.append(decoder_config.mlp.fc) + linear_layers.append(decoder_config.mlp.proj) + linear_layers.append(decoder_config.mlp.gate) + + _linear_layer_to_quantized_weight(linear_layers) + + if model_config.medusa_heads is not None: + linear_layers = [] + + for head in model_config.medusa_heads: + linear_layers.append(head.lm_head) + for layer in head.medusa_layers: + linear_layers.append(layer.linear) + + _linear_layer_to_quantized_weight(linear_layers) + + # lm_head can be quantized by AutoQuant + if model_config.lm_head is not None: + _linear_layer_to_quantized_weight([model_config.lm_head]) + + +def naive_quantization(config: ModelConfig): + """Generates a constant scaling factor (1) with target quantization. + + This is for debugging and performance measurement only. + """ + if isinstance(config.layers[0], DecoderLayerConfig): + assert ( + config.layers[0].decoder_type != "mixtral" + ), "Mixtral native quantization is not supported." + + config.quantization = QUANTIZATION_FP8 + default_scaling_factor = torch.tensor([1], dtype=torch.float32) + + for layer in config.layers: + if layer.attention: + linear_layers = [ + layer.attention.dense, + ] + + if isinstance(layer.attention.qkv, QKVConfig): + linear_layers += [ + layer.attention.qkv.q, + layer.attention.qkv.k, + layer.attention.qkv.v, + ] + elif isinstance(layer.attention.qkv, LinearConfig): + linear_layers += [layer.attention.qkv] + + elif layer.recurrent: + linear_layers = [ + layer.recurrent.linear_x, + layer.recurrent.linear_y, + layer.recurrent.linear_out, + ] + else: + linear_layers = [] + + moe_layers = [] + if isinstance(layer.mlp, MOEConfig): + moe_layers.append(layer.mlp.experts.fc) + moe_layers.append(layer.mlp.experts.proj) + else: + linear_layers.append(layer.mlp.fc) + linear_layers.append(layer.mlp.proj) + linear_layers.append(layer.mlp.gate) + + for linear_layer in linear_layers: + if linear_layer: + linear_layer.activation_scaling_factor = default_scaling_factor + linear_layer.weights_scaling_factor = default_scaling_factor + + for moe_layer in moe_layers: + if moe_layer: + num_expert = moe_layer.weight.shape[0] + # The moe plugin only supports a single activation scaling factor for all experts + moe_layer.activation_scaling_factor = default_scaling_factor + moe_layer.weights_scaling_factor = default_scaling_factor.expand(num_expert, 1) + + if layer.attention: + layer.attention.kv_cache_dtype = KV_CACHE_FP8 + layer.attention.kv_cache_scaling_factor = torch.tensor([1.0], dtype=torch.float) diff --git a/modelopt/torch/export/postprocess.py b/modelopt/torch/export/postprocess.py new file mode 100644 index 0000000000..edf7517cd5 --- /dev/null +++ b/modelopt/torch/export/postprocess.py @@ -0,0 +1,817 @@ +# 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. + +"""Utils to load and process model_config.""" + +import copy +from dataclasses import fields, is_dataclass +from pathlib import Path +from typing import Optional, Union +from warnings import warn + +import torch + +from modelopt.torch.quantization.model_calib import disable_pre_quant_scale_and_resmooth +from modelopt.torch.quantization.nn import SequentialQuantizer +from modelopt.torch.quantization.nn.modules.quant_linear import QuantLinear +from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.utils import distributed as dist + +from .distribute import get_configs_parallel, get_tensors_parallel +from .model_config import ( + LINEAR_COLUMN, + LINEAR_GROUP, + LINEAR_ROW, + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + ConvConfig, + EmbeddingConfig, + ExpertConfig, + LinearConfig, + ModelConfig, + MOEConfig, + RelativeAttentionTableConfig, +) +from .model_config_utils import pad_weights +from .quant_utils import get_scaling_factor_from_weight, resmooth_and_get_scale + + +def _same_tensor(tensors: list[torch.Tensor]): + return all(t is None for t in tensors) or all(torch.equal(tensors[0], t) for t in tensors[1:]) + + +# Define helper functions +def view_as_uint8_if_needed(tensor): + """Views fp8 tensor as uint8 tensor for processing using torch operations.""" + is_float8_e4m3fn = tensor.dtype == torch.float8_e4m3fn + return tensor.view(torch.uint8) if is_float8_e4m3fn else tensor + + +def view_as_float8_e4m3fn_if_needed(tensor): + """Views uint8 tensor as fp8 tensor after processing using torch operations.""" + is_uint8 = tensor.dtype == torch.uint8 + return tensor.view(torch.float8_e4m3fn) if is_uint8 else tensor + + +def _split_model_config_for_tp(merged_config, split_factor): + """This method splits the tensor fields for linear config so the config can be used with more GPUs. + + The implementation is recursive. + """ + configs = [copy.copy(merged_config) for _ in range(split_factor)] + + if isinstance(merged_config, ConvConfig): + raise NotImplementedError("Do not support ConvConfig TP merge or split") + + elif isinstance(merged_config, EmbeddingConfig): + weights = torch.chunk(pad_weights(merged_config.weight, split_factor), split_factor, dim=0) + for i, config in enumerate(configs): + config.weight = weights[i] + + elif isinstance(merged_config, MOEConfig): + split_expert_configs = _split_model_config_for_tp( + merged_config.experts, + split_factor, + ) + # TP for rounter of MoE is skipped for better performance + # See https://github.com/NVIDIA/TensorRT-LLM/pull/1091 for details + for i in range(split_factor): + configs[i].experts = split_expert_configs[i] + configs[i].router = merged_config.router + + elif isinstance(merged_config, ExpertConfig): + assert merged_config.proj.linear_type != LINEAR_COLUMN # row + assert merged_config.fc.linear_type == LINEAR_COLUMN # column + configs = [copy.deepcopy(merged_config) for _ in range(split_factor)] + weights = torch.chunk( + pad_weights(merged_config.proj.weight, split_factor), split_factor, dim=2 + ) + for i, config in enumerate(configs): + config.proj.weight = weights[i] + + # Split w3 and w1 weight tensors separately and then concat for each TP + merged_w3, merged_w1 = torch.chunk(merged_config.fc.weight, 2, dim=1) + merged_w3_weights = torch.chunk(merged_w3, split_factor, dim=1) + merged_w1_weights = torch.chunk(merged_w1, split_factor, dim=1) + for i, config in enumerate(configs): + config.fc.weight = torch.concat([merged_w3_weights[i], merged_w1_weights[i]], dim=1) + + # TODO: check if this works with int8_sq + # Split scaling factors for int4_awq + if ( + config.fc.weights_scaling_factor is not None + and config.fc.weights_scaling_factor.dim() == 3 + ): # column linear + merged_scaling_factor_w3, merged_scaling_factor_w1 = torch.chunk( + view_as_uint8_if_needed(merged_config.fc.weights_scaling_factor), 2, dim=1 + ) + scaling_factor_w3 = torch.chunk(merged_scaling_factor_w3, split_factor, dim=1) + scaling_factor_w1 = torch.chunk(merged_scaling_factor_w1, split_factor, dim=1) + for i, config in enumerate(configs): + config.fc.weights_scaling_factor = view_as_float8_e4m3fn_if_needed( + torch.concat([scaling_factor_w3[i], scaling_factor_w1[i]], dim=1) + ) + + if ( + config.proj.weights_scaling_factor is not None + and config.proj.weights_scaling_factor.dim() == 3 + ): # row linear + scaling_factor_proj = torch.chunk( + merged_config.proj.weights_scaling_factor, split_factor, dim=2 + ) + + for i, config in enumerate(configs): + config.proj.weights_scaling_factor = scaling_factor_proj[i] + + # For NVFP4_AWQ and INT4_AWQ we split pre_quant_scale across tp + if merged_config.proj.prequant_scaling_factor is not None: + prequant_scaling_factors = torch.chunk( + merged_config.proj.prequant_scaling_factor, split_factor, dim=-1 + ) + for i, config in enumerate(configs): + config.proj.prequant_scaling_factor = prequant_scaling_factors[i] + + elif isinstance(merged_config, LinearConfig): + # The scaling factors split rule is summarized as below: + + # S: all ranks should have the same scaling factor. + # P0: Split the scaling factors on dim 0. + # P1: Split the scaling factors on dim 1. + # NA: Not valid / present + + # ws: weight scaling factor + # as: activation scaling factor + # ps: prequant scaling factor + + # C: Colum Linear + # R: Row Linear + + # F: FP8 + # I8: INT8 SQ + # I4: INT4 AWQ + + # Split Rules: + # ws as ps + # FC S S NA + # FR S S NA + # I8C P0 S S + # I8R S S P0 + # I4C P0 S S + # I4R P1 S P0 + + # For INT4 AWQ reference implemention: please check examples/llama/weight.py in the tekit repo + # For normal linear layers, we split the column linear on the dim 0 and row on the dim 1 + assert ( + merged_config.linear_type != LINEAR_GROUP + ), "Do not support group linear TP merge or split" + + split_axis = 0 if merged_config.linear_type == LINEAR_COLUMN else 1 + if merged_config.linear_type == LINEAR_COLUMN: + merged_config.weight = pad_weights(merged_config.weight, split_factor) + weights = torch.chunk(merged_config.weight, split_factor, dim=split_axis) + for i, config in enumerate(configs): + config.weight = weights[i] + + # Only split the bias for column linear. + if merged_config.linear_type == LINEAR_COLUMN and merged_config.bias is not None: + biases = torch.chunk(merged_config.bias, split_factor, dim=0) + for i, config in enumerate(configs): + config.bias = biases[i] + + if merged_config.linear_type == LINEAR_COLUMN: + if ( + merged_config.weights_scaling_factor is not None + and merged_config.weights_scaling_factor.numel() != 1 + ): + # INT4 AWQ and INT8 sq all linear cases + scaling_factors = torch.chunk( + view_as_uint8_if_needed(merged_config.weights_scaling_factor), + split_factor, + dim=0, + ) + + for i, config in enumerate(configs): + config.weights_scaling_factor = view_as_float8_e4m3fn_if_needed( + scaling_factors[i] + ) + + else: + if ( + merged_config.weights_scaling_factor is not None + and merged_config.awq_block_size != 0 + ): + scaling_factors = torch.chunk( + view_as_uint8_if_needed(merged_config.weights_scaling_factor), + split_factor, + dim=1, + ) + + for i, config in enumerate(configs): + config.weights_scaling_factor = view_as_float8_e4m3fn_if_needed( + scaling_factors[i] + ) + + if merged_config.prequant_scaling_factor is not None: + prequant_scaling_factors = torch.chunk( + merged_config.prequant_scaling_factor, split_factor, dim=0 + ) + for i, config in enumerate(configs): + config.prequant_scaling_factor = prequant_scaling_factors[i] + + elif isinstance(merged_config, RelativeAttentionTableConfig): + weights = torch.chunk(merged_config.weight, split_factor, dim=0) + for i, config in enumerate(configs): + config.weight = weights[i] + + elif is_dataclass(merged_config): + for field in fields(merged_config): + field_configs = _split_model_config_for_tp( + getattr(merged_config, field.name), + split_factor, + ) + for r in range(split_factor): + setattr(configs[r], field.name, field_configs[r]) + elif isinstance(merged_config, list): + for i in range(len(merged_config)): + field_configs = _split_model_config_for_tp( + merged_config[i], + split_factor, + ) + for r in range(split_factor): + configs[r][i] = field_configs[r] + + return configs + + +def _split_model_config_for_pp(merged_config, split_factor): + """This method splits ModelConfig for inference pipeline parallel.""" + num_layers = len(merged_config.layers) + assert num_layers % split_factor == 0 + layers_per_pp = num_layers // split_factor + + configs = [copy.copy(merged_config) for _ in range(split_factor)] + for i, config in enumerate(configs): + if i > 0: + config.vocab_embedding = None + config.position_embedding = None + config.ln_embed = None + if i < split_factor - 1: + config.ln_f = None + config.lm_head = None + config.layers = config.layers[i * layers_per_pp : (i + 1) * layers_per_pp] + + return configs + + +def _merge_model_configs_to_first_tp(config, ranks: list[int], group=None): + """This method merges the tensor fields for linear config so the config can be used with fewer GPUs. + + The implementation is recursive. + """ + if isinstance(config, ConvConfig): + raise NotImplementedError("Do not support ConvConfig TP merge or split") + + elif isinstance(config, EmbeddingConfig): + assert config.weight is not None + with get_tensors_parallel(config.weight, ranks, group) as weights: + if weights: + config.weight = torch.cat(weights, dim=0) + + elif isinstance(config, LinearConfig): + # The scaling factors merge rule is summarized as below: + + # S: all ranks should have the same scaling factor. + # M: Pick elementwise max among the ranks. Merged shape same as single rank. + # C0: Concat the scaling factors on dim 0. Merged shape == tensor_parallel * original shape. + # C1: Concat the scaling factors on dim 1. Merged shape == original shape * tensor_parallel. + # NA: Not valid / present + + # ws: weight scaling factor + # as: activation scaling factor + # ps: prequant scaling factor + + # C: Colum Linear + # R: Row Linear + + # F: FP8 + # I8: INT8 SQ + # I4: INT4 AWQ + + # Merge Rules: + # ws as ps + # FC M M NA + # FR M M NA + # I8C C0 M S + # I8R M M C0 + # I4C C0 M S + # I4R C1 M C0 + + assert config.linear_type != LINEAR_GROUP, "Do not support group linear TP merge or split" + + # Handling constants + for field_name in [ + "activation_scaling_factor", + "weights_scaling_factor", + "weights_scaling_factor_2", + ]: + field_value = getattr(config, field_name) + if field_value is not None and field_value.numel() == 1: + with get_tensors_parallel(field_value, ranks, group) as scaling_factors: + if scaling_factors: + # Scaling factor is a scalar. + setattr( + config, + field_name, + torch.stack(scaling_factors).max(dim=0).values, + ) + + # We merge column linear on the dim 0 and row on the dim 1 + merge_axis = 0 if config.linear_type == LINEAR_COLUMN else 1 + + assert config.weight is not None + with get_tensors_parallel(config.weight, ranks, group) as weights: + if weights: + config.weight = torch.cat(weights, dim=merge_axis) + + # Only cat the bias for column linear. + if config.linear_type == LINEAR_COLUMN and config.bias is not None: + with get_tensors_parallel(config.bias, ranks, group) as biases: + if biases: + config.bias = torch.cat(biases, dim=0) + + if config.linear_type == LINEAR_COLUMN: + if ( + config.weights_scaling_factor is not None + and config.weights_scaling_factor.numel() != 1 + ): + # INT8 sq + with get_tensors_parallel( + config.weights_scaling_factor, ranks, group + ) as w_scaling_factors: + if w_scaling_factors: + config.weights_scaling_factor = torch.cat(w_scaling_factors, dim=0) + + if config.prequant_scaling_factor is not None: + with get_tensors_parallel( + config.prequant_scaling_factor, ranks, group + ) as p_scaling_factors: + if p_scaling_factors: + # INT4 AWQ, desmooth and de-smooth and re-smooth across all ranks + if config.awq_block_size != 0: + ( + config.weight, + config.weights_scaling_factor, + config.prequant_scaling_factor, + ) = resmooth_and_get_scale( + merged_weights=config.weight, + pre_quant_scales=p_scaling_factors, + ranks=len(ranks), + group_size=config.awq_block_size, + quantization=config.quantization, + ) + if config.quantization == QUANTIZATION_NVFP4_AWQ: + # We have to update weight_scaling_factor and weight_scaling_factor_2 + config.weights_scaling_factor, config.weights_scaling_factor_2 = ( + NVFP4QTensor.get_weights_scaling_factor( + config.weight, config.awq_block_size + ) + ) + else: + assert _same_tensor( + p_scaling_factors + ), f"Failed to merge config {config} with others" + else: + if config.weights_scaling_factor is not None: + with get_tensors_parallel( + config.weights_scaling_factor, ranks, group + ) as w_scaling_factors: + if w_scaling_factors: + if config.awq_block_size != 0: + # INT4 AWQ + if w_scaling_factors[0].ndim == 2: + scaling_factors_total_size = 0 + for _, w_scaling_factor in enumerate(w_scaling_factors): + scaling_factors_total_size += w_scaling_factor.numel() + if scaling_factors_total_size != config.weight.numel(): + # The weights from each rank are padded to a multiple of group_size in this case. + # We need to merge the weights and recalculate the scaling factors. + config.weights_scaling_factor = get_scaling_factor_from_weight( + config.weight, config.awq_block_size + ) + if config.quantization in [ + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + ]: + ( + config.weights_scaling_factor, + config.weights_scaling_factor_2, + ) = NVFP4QTensor.get_weights_scaling_factor( + config.weight, config.awq_block_size + ) + else: + config.weights_scaling_factor = torch.cat( + w_scaling_factors, dim=1 + ) + else: + raise NotImplementedError( + "Unexpected dimensions for scaling factors." + ) + else: + # INT8 SQ + config.weights_scaling_factor = ( + torch.stack(w_scaling_factors).max(dim=0).values + ) + + if config.prequant_scaling_factor is not None: + with get_tensors_parallel( + config.prequant_scaling_factor, ranks, group + ) as p_scaling_factors: + if p_scaling_factors: + config.prequant_scaling_factor = torch.cat(p_scaling_factors, dim=0) + + elif is_dataclass(config): + for field in fields(config): + _merge_model_configs_to_first_tp(getattr(config, field.name), ranks, group) + elif isinstance(config, list): + for i in range(len(config)): + _merge_model_configs_to_first_tp(config[i], ranks, group) + + +def _model_model_configs_to_first_pp( + model_config: ModelConfig, + ranks: list[int], + workspace_path: Optional[Union[Path, str]] = None, +): + """Merges the mode_config from each rank to the first pp rank. + + Args: + model_config: model or module config + ranks: a list of ranks in the same pipeline parallel group + workspace_path: the path to the NFS directory for postprocess cross rank communication. + """ + # TODO: There is an NCCL error if we try group sync based on the pp_ranks. + # So we just ask all groups to sync together for now. + group = None + + # Merge decoder layers. + decoder_layers = [] + for layer in model_config.layers: + with get_configs_parallel( + layer, ranks, group, workspace_path=workspace_path + ) as layer_configs: + layer_configs_copy = [] + if layer_configs: + layer_configs_copy.append(layer_configs[0]) + for config in layer_configs[1:]: + if config: + # Have to copy the config from the other pps as the shm will be releases after + layer_configs_copy.append(copy.deepcopy(config)) + decoder_layers.append(layer_configs_copy) + + # If we take a 80-layer TP8/PP4 export, then locally each pp rank has 20 layers. + # If we list the global layer id (0-base) in decoder.layers, then + # + # decoder_layers = [ + # [decoder.layers.0 (owned by rank 0-7), 20 (rank 8-15), 40 (rank 16-23), 60 (rank 24-31)], + # [decoder.layers.1 (owned by rank 0-7), 21 (rank 8-15), 41 (rank 16-23), 61 (rank 24-31)], + # ... + # [decoder.layers.19 (owned by rank 0-7), 39 (rank 8-15), 59 (rank 16-23), 79 (rank 24-31)], + # ] + # + # To merge into PP1, we need to gather and shuffle back to the order of + # + # shuffled_config_layers = [ + # decoder.layers.0 (owned by rank 0-7), + # decoder.layers.1 (owned by rank 0-7), + # ... + # decoder.layers.19 (owned by rank 0-7), + # decoder.layers.20 (owned by rank 0-7), + # decoder.layers.21 (owned by rank 0-7), + # ... + # decoder.layers.78 (owned by rank 0-7), + # decoder.layers.79 (owned by rank 0-7), + # ] + model_config.layers = [ + layer for shuffled_layers in zip(*decoder_layers) for layer in shuffled_layers + ] + + # Get the ln_f from the last PP rank + with get_configs_parallel( + model_config.ln_f, ranks, group, workspace_path=workspace_path + ) as configs: + if configs and configs[-1] is not None: + model_config.ln_f = configs[-1] + + # Get the lm_head from the last PP rank + with get_configs_parallel( + model_config.lm_head, ranks, group, workspace_path=workspace_path + ) as configs: + if configs and configs[-1] is not None: + model_config.lm_head = configs[-1] + + +def postprocess_model_config( + model_config, + inference_tensor_parallel: int = 1, + inference_pipeline_parallel: int = 1, + training_pipeline_parallel: int = 1, + workspace_path: Optional[Union[Path, str]] = None, +) -> list[ModelConfig]: + """Postprocesses the model configs with trained tensor parallel to target inference tensor parallel. + + If the training_pipeline_parallel > 1, the model configs across PP will be merged to one. + + Returns: + The processed model config as a list. + For the merging case: + The merged rank will return the merged model_config as an single item list. + The other ranks will return an empty list as we no longer export them. + For the split case: + The splitted model config list is returned. + """ + rank = dist.rank() + + # We assume the ranks ardistributed in training as [PP size, TP size]. + training_tensor_parallel = dist.size() // training_pipeline_parallel + tp_rank = rank % training_tensor_parallel + pp_rank = rank // training_tensor_parallel + + print(f"current rank: {rank}, tp rank: {tp_rank}, pp rank: {pp_rank}") + dist.barrier() + + # Merge PP ranks to the first + if training_pipeline_parallel > 1: + # The pp_ranks for the same tp is [tp_rank, tp_rank + tp, tp_rank + tp * 2, ...] + pp_ranks = torch.arange(tp_rank, dist.size(), training_tensor_parallel, dtype=int).tolist() + + print(f"PP: Current rank {rank}, merge to {pp_ranks[0]}. Merge group {pp_ranks}") + dist.barrier() + _model_model_configs_to_first_pp( + model_config, + pp_ranks, + workspace_path=workspace_path, + ) + + # Returns the empty model_config on other PP ranks. + if pp_rank != 0: + model_config.rank = -1 + return [] + + # Now we are with PP == 1 + model_config.pipeline_parallel = 1 + model_config.rank = tp_rank + + tp_world_size = dist.size() // training_pipeline_parallel + if inference_tensor_parallel < tp_world_size: + # Merge the model_configs to target inference tensor parallel. + assert ( + tp_world_size % inference_tensor_parallel == 0 + ), f"Cannot merge {tp_world_size} configs to {inference_tensor_parallel}" + + num_configs_per_group = tp_world_size // inference_tensor_parallel + local_tp_group_id = tp_rank // num_configs_per_group + tp_ranks = list( + range( + local_tp_group_id * num_configs_per_group, + (local_tp_group_id + 1) * num_configs_per_group, + ) + ) + + print(f"TP: Current rank {rank}, merge to {tp_ranks[0]}. Merge group {tp_ranks}.") + # We sync on all TP ranks (and pp_rank = 0) + group = dist.get_group(list(range(training_tensor_parallel))) + _merge_model_configs_to_first_tp(model_config, tp_ranks, group) + model_config.tensor_parallel = inference_tensor_parallel + if rank == tp_ranks[0]: + model_config.rank = local_tp_group_id + splitted_model_configs = [model_config] + else: + # Mark this config to be invalid and return it as invalid. + model_config.rank = -1 + return [] + + elif inference_tensor_parallel > tp_world_size: + assert ( + tp_world_size == 1 + ), "We only support splitting a single model config to multiple GPUs" + split_factor = inference_tensor_parallel // tp_world_size + splitted_model_configs = _split_model_config_for_tp( + model_config, + split_factor, + ) + for i, config in enumerate(splitted_model_configs): + config.rank = i + config.tensor_parallel = inference_tensor_parallel + + else: + splitted_model_configs = [model_config] + + if inference_pipeline_parallel > 1: + splitted_model_configs_tp_pp = [] + for i, model_config_tp in enumerate(splitted_model_configs): + splitted_model_configs_pp = _split_model_config_for_pp( + model_config_tp, inference_pipeline_parallel + ) + for j, config in enumerate(splitted_model_configs_pp): + config.rank = i + j * inference_tensor_parallel + config.tensor_parallel = inference_tensor_parallel + config.pipeline_parallel = inference_pipeline_parallel + splitted_model_configs_tp_pp.extend(splitted_model_configs_pp) + return splitted_model_configs_tp_pp + else: + return splitted_model_configs + + +def pad_embedding_lm_head(model_config: ModelConfig, padding_factor: int = 64): + """Pad lm_head and embedding as multiples of 64 for AWQ quantization.""" + vocab_size = model_config.vocab_size + # Pad the lm_head and vocab_embedding only if the lm_head is quantized with AWQ. + if vocab_size % padding_factor == 0: + return + + pad_vocab_size = int((vocab_size + padding_factor - 1) / padding_factor) * padding_factor + model_config.vocab_size = pad_vocab_size + + if hasattr(model_config, "vocab_embedding"): + embedding_config = model_config.vocab_embedding + original_weight = embedding_config.weight + pad_size = (0, 0, 0, pad_vocab_size - original_weight.shape[0]) + embedding_config.weight = torch.nn.functional.pad( + original_weight, pad_size, mode="constant", value=0 + ) + + if hasattr(model_config, "lm_head"): + lm_head_config = model_config.lm_head + original_weight = lm_head_config.weight + original_bias = lm_head_config.bias + original_wsf = lm_head_config.weights_scaling_factor + + lm_head_config.weight = torch.nn.functional.pad( + original_weight, + (0, 0, 0, pad_vocab_size - original_weight.shape[0]), + mode="constant", + value=0, + ) + if original_bias is not None: + lm_head_config.bias = torch.nn.functional.pad( + original_bias, + (0, pad_vocab_size - original_bias.shape[0]), + mode="constant", + value=0, + ) + + if original_wsf is not None: + assert len(original_wsf.shape) == 2, "AWQ weight scaling factor should be 2D." + pad_weights_scaling_factor = ( + torch.ones( + (pad_vocab_size, original_wsf.shape[1]), + dtype=original_wsf.dtype, + ) + / 7.0 # int4: maxbound = 7.0 + ) + + pad_weights_scaling_factor[:vocab_size, :] = original_wsf + lm_head_config.weights_scaling_factor = pad_weights_scaling_factor + + +def update_lm_head_quantization( + config: ModelConfig, lm_head: QuantLinear, inference_tensor_parallel: int = 1 +): + """Update lm_head quantization config for TRT-LLM export.""" + if not hasattr(lm_head, "weight_quantizer") or not hasattr(lm_head, "input_quantizer"): + return + weight_quantizer = lm_head.weight_quantizer + input_quantizer = lm_head.input_quantizer + + disable_lm_head_quantization = False + + assert config.layers is not None, "Expect layers in the model config." + + # check if weight shape is not multiples of block size + awq_block_size = 0 + for quantizer in SequentialQuantizer.tensor_quantizer_iterator(weight_quantizer): + if hasattr(weight_quantizer, "block_sizes") and weight_quantizer.block_sizes is not None: + awq_block_size = weight_quantizer.block_sizes[-1] + weight_quantizer = quantizer + break + if ( + awq_block_size > 0 + and weight_quantizer.is_enabled + and lm_head.weight.shape[0] % (awq_block_size * inference_tensor_parallel) != 0 + ): + disable_lm_head_quantization = True + + # disable quantizer + if disable_lm_head_quantization: + if hasattr(input_quantizer, "_pre_quant_scale"): + disable_pre_quant_scale_and_resmooth(lm_head, delete_pre_quant_scale=True) + + for quantizer in SequentialQuantizer.tensor_quantizer_iterator(lm_head.weight_quantizer): + quantizer.disable() + + input_quantizer.disable() + print("Disable lm_head quantization for TRT-LLM export due to deployment limitations.") + + else: + warn( + "Enable lm_head quantization. lm_head quantization may lead to additional accuracy loss." + ) + + +def check_weight_shape_valid(config, inference_tensor_parallel=1, training_tensor_parallel=1): + """Check if weight shape are valid with inference TP. + + This function is recurisve. + """ + + def _check_merged_weight(merged_k): + assert ( + merged_k % inference_tensor_parallel == 0 + ), f"Weights cannot be split into {inference_tensor_parallel} ranks." + + def _check_merged_weight_scaling_factor(merged_k, awq_block_size): + if awq_block_size > 0 and (merged_k // inference_tensor_parallel) % awq_block_size != 0: + raise NotImplementedError( + "Weight shape is not divisible for block size for block quantization." + ) + + def _check_merged_channel_is_valid(merged_k, awq_block_size): + _check_merged_weight(merged_k=merged_k) + _check_merged_weight_scaling_factor(merged_k=merged_k, awq_block_size=awq_block_size) + + if isinstance(config, LinearConfig): + # check weight shape + if config.linear_type == LINEAR_COLUMN: + _, k = config.weight.shape + merged_k = k * training_tensor_parallel + _check_merged_channel_is_valid(merged_k, config.awq_block_size) + elif config.linear_type == LINEAR_ROW: + k, m = config.weight.shape + merged_k = k * training_tensor_parallel + merged_m = m * training_tensor_parallel + # For int4_awq, weight scaling factors will be split as (k, (merged_m // TP) // block_size) + _check_merged_weight(merged_k=merged_k) + _check_merged_weight_scaling_factor(merged_m, config.awq_block_size) + + return + + if isinstance(config, ExpertConfig): + _, _, k = config.fc.weight.shape + merged_k = k * training_tensor_parallel + _check_merged_channel_is_valid(merged_k, config.fc.awq_block_size) + return + + if is_dataclass(config): + for field in fields(config): + check_weight_shape_valid( + getattr(config, field.name), + inference_tensor_parallel, + training_tensor_parallel, + ) + elif isinstance(config, list): + for config_i in config: + check_weight_shape_valid(config_i, inference_tensor_parallel, training_tensor_parallel) + + +def postprocess_tensors( + weights: dict[str, torch.tensor], + dtype: torch.dtype, + force_cpu: bool = True, + force_contiguous: bool = True, + force_non_view: bool = True, +): + """Make all tensors in the weights to be with the target dtype, on CPU, contiguous and own the memory.""" + for k, v in weights.items(): + if force_cpu: + v = v.cpu() + if ( + any( + k.endswith(suffix) + for suffix in ["weight", "bias", "prequant_scaling_factor", "recurrent_param"] + ) + and "router" not in k + and v.dtype + in [ + torch.float16, + torch.bfloat16, + torch.float, + ] + ): + v = v.type(dtype) + if force_contiguous: + v = v.contiguous() + if force_non_view and v._is_view(): + v = v.clone() + weights[k] = v diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py new file mode 100644 index 0000000000..dba508fc58 --- /dev/null +++ b/modelopt/torch/export/quant_utils.py @@ -0,0 +1,687 @@ +# 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. + +"""Utils for quantization including scaling factors adjustments.""" + +import warnings +from typing import Any, Optional, Union +from warnings import warn + +import torch +import torch.nn as nn + +from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection +from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.quantization.utils import is_quantized_linear + +from ..quantization.nn import SequentialQuantizer, TensorQuantizer +from .model_config import ( + KV_CACHE_FP8, + KV_CACHE_INT8, + QUANTIZATION_FP8, + QUANTIZATION_INT4_AWQ, + QUANTIZATION_INT8_SQ, + QUANTIZATION_NONE, + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + QUANTIZATION_W4A8_AWQ, +) + + +def get_scaling_factor_from_weight(weight, group_size) -> torch.tensor: + """Calculate the weight scaling factor for a given group size.""" + [n, k] = weight.shape + + if group_size != 0: + # int4_awq + if k % group_size != 0: + raise NotImplementedError( + "Weight shape is not divisible for block size for block quantization." + ) + weight = weight.reshape(n, k // group_size, group_size) + maxbound = 7.0 + else: + # int8_sq + maxbound = 127.0 + amax = weight.abs().max(dim=-1)[0].float() + + weights_scaling_factor = amax / maxbound + + # Let's filter the zeros in the scaling factor if the weights are zero + # to avoid the divided-by-zero error.. + weights_scaling_factor[weights_scaling_factor == 0] = 1.0 + + return weights_scaling_factor + + +def resmooth_and_get_scale( + merged_weights: torch.Tensor, + pre_quant_scales: list[torch.Tensor], + ranks: int, + group_size: int, + new_pre_quant_scale: torch.Tensor = None, + quantization: Optional[str] = QUANTIZATION_NONE, +): + """Resmooths weights from a single or multiple ranks and get scaling factors and amax. + + Args: + merged_weights: Merged weights from ranks. + pre_quant_scales: List of pre-quantization scales for each rank. + ranks: Number of ranks. + group_size: Group size of the quantization block. + new_pre_quant_scale (optional): If not provided, weights will be resmoothed using + the average of pre_quant_scales. + + Returns: + weights: Resmoothed weights. + weight_scaling_factors: Resmoothed scaling factors. + avg_pre_quant_scale: Calculated average of the quantization scale. + """ + if new_pre_quant_scale is None: + new_pre_quant_scale = torch.stack(pre_quant_scales).mean(dim=0) + + assert ( + len(pre_quant_scales) > 0 and new_pre_quant_scale.numel() == merged_weights.shape[1] + ), "Shape of pre_quant_scales and weights do not match." + weights = torch.chunk(merged_weights, ranks, dim=0) + + scales = [] + new_weights = [] + for i, p_scaling_factor in enumerate(pre_quant_scales): + # De smooth & Re smooth + weight = ( + weights[i] + * p_scaling_factor.type(weights[i].dtype) + / new_pre_quant_scale.type(weights[i].dtype) + ) + new_weights.append(weight) + # If NVFP4_AWQ then we view the scales as uint8 to allow for cat later + if quantization == QUANTIZATION_NVFP4_AWQ: + scale, _ = NVFP4QTensor.get_weights_scaling_factor(weight, group_size).view(torch.uint8) + else: + scale = get_scaling_factor_from_weight(weight, group_size) + scales.append(scale) + + resmoothed_scales = torch.cat(scales, dim=0) + + return ( + torch.cat(new_weights, dim=0), + resmoothed_scales.view(torch.float8_e4m3fn) + if quantization == QUANTIZATION_NVFP4_AWQ + else resmoothed_scales, # if NVFP4_AWQ we view the scales back as float8_e4m3fn after cat + new_pre_quant_scale, + ) + + +def adjust_attn_amax_values(module): + """Adjusts the amax values for the attention layers.""" + projection_prefixes = ["q", "k", "v"] + max_amax = float("-inf") + proj_layers = [] + + # Find all projection layers whose names contain 'q', 'k', or 'v' + for name, sub_module in module.named_children(): + for prefix in projection_prefixes: + if ( + prefix in name + and hasattr(sub_module, "weight_quantizer") + and hasattr(sub_module.weight_quantizer, "amax") + ): + proj_layers.append(sub_module) + max_amax = max(max_amax, sub_module.weight_quantizer.amax.item()) + + if not proj_layers: + raise ValueError( + "No projection layers with the specified prefixes ('q', 'k', 'v') have amax attributes" + ) + + assert max_amax > 0, "max_amax must be positive." + + # Set all amax values to the maximum found + for proj_layer in proj_layers: + proj_layer.weight_quantizer.amax.fill_(max_amax) + + +def get_scaling_factor(quantizer: TensorQuantizer) -> torch.Tensor: + """Returns scaling factor from the quantizer as torch.Tensor.""" + if not quantizer.is_enabled: + return None + + amax = quantizer.export_amax() + if amax is None: + return None + + # tensorrt_llm uses float as the scaling_factors. + scaling_factor = amax.float() / quantizer.maxbound + + assert torch.all(scaling_factor > 0), f"scaling factor {scaling_factor} not positive." + + return scaling_factor + + +def get_activation_scaling_factor(module: nn.Module) -> torch.Tensor: + """Returns the activation scaling factor.""" + # If NVFP4, return activation scaling factor from NVFP4QTensor + if get_quantization_format(module) in [ + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + ] and hasattr(module, "input_quantizer"): + return NVFP4QTensor.get_activation_scaling_factor(module.input_quantizer) + return ( + get_scaling_factor(module.input_quantizer) if hasattr(module, "input_quantizer") else None + ) + + +def get_weight_scaling_factor(module: nn.Module) -> torch.Tensor: + """Returns the weight scaling factor.""" + # module.weight_quantizer could be a TensorQuantizer (for algorithms except W4A8) or + # a SequentialQuantizer (for W4A8). In the latter case, we need to get the scaling factor from the + # first quantizer of the SequentialQuantizer instance. + if hasattr(module, "weight_quantizer") and isinstance( + module.weight_quantizer, SequentialQuantizer + ): + return get_scaling_factor(module.weight_quantizer[0]) + + # If NVFP4, we need to return quantized per_block scaling factors + if get_quantization_format(module) in [ + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + ] and hasattr(module, "weight_quantizer"): + return NVFP4QTensor.get_weights_scaling_factor( + module.weight, + module.weight_quantizer.block_sizes[-1], + NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(module.weight_quantizer), + )[0] + + return ( + get_scaling_factor(module.weight_quantizer) if hasattr(module, "weight_quantizer") else None + ) + + +def get_weight_scaling_factor_2(module: nn.Module) -> torch.Tensor: + """Returns the secondary weight scaling factor.""" + if get_quantization_format(module) in [ + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + ] and hasattr(module, "weight_quantizer"): + return NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(module.weight_quantizer) + if ( + not hasattr(module, "weight_quantizer") + or not isinstance(module.weight_quantizer, SequentialQuantizer) + or not module.weight_quantizer[-1].is_enabled + ): + return None + assert ( + len(module.weight_quantizer) == 2 + ), "modelopt only supports 2 sequential quantization layers for now" + return get_scaling_factor(module.weight_quantizer[-1]) + + +def get_prequant_scaling_factor(module: nn.Module) -> torch.Tensor: + """Returns the prequant scaling factor.""" + prequant_scaling_factor = ( + module.input_quantizer._pre_quant_scale.squeeze() + if hasattr(module, "input_quantizer") + and hasattr(module.input_quantizer, "_pre_quant_scale") + else None + ) + + if prequant_scaling_factor is not None: + assert torch.all( + prequant_scaling_factor > 0 + ), f"prequant scaling factor {prequant_scaling_factor} not positive." + return prequant_scaling_factor + + +def get_kv_cache_scaling_factor(qkv_modules: list[nn.Module]) -> torch.Tensor: + """Returns the kv_cache scaling factor if output quantizer is set. Else returns None by default.""" + scaling_factors = [ + get_scaling_factor(module.output_quantizer) + for module in qkv_modules + if hasattr(module, "output_quantizer") + ] + + scaling_factors = [ + scaling_factor for scaling_factor in scaling_factors if scaling_factor is not None + ] + + if not scaling_factors: + return None + + scaling_factor = torch.stack(scaling_factors).max(dim=0).values + + # For FP8, we recommend default kv cache scaling factor to be 1. + if get_kv_cache_dtype(qkv_modules) == KV_CACHE_FP8: + if scaling_factor.item() > 0.5: + warn( + f"!!!!\nWarning: Large KV activations detected: {scaling_factor.item()}, " + "Quantized KV cache may lead to higher accuracy drop.\n!!!!" + ) + scaling_factor = torch.max( + scaling_factor, + torch.tensor([1.0], dtype=torch.float, device=scaling_factor.device), + ) + return scaling_factor + + +def get_kv_cache_dtype(modules: Union[list[nn.Module], nn.Module]) -> Optional[str]: + """Returns the kv_cache dtype. + + If num_bits of output_quantizer is (4, 3) then returns FP8; if it is 8, returns int8, + otherwise returns None. + + Args: + modules (Union[list[nn.Module], nn.Module]): The module or list of modules to inspect. + + Returns: + str: The kv_cache dtype. + """ + num_bits_list = [] + + if isinstance(modules, nn.Module): + modules = [modules] + + for module in modules: + if hasattr(module, "output_quantizer") and module.output_quantizer.is_enabled: + num_bits_list.append(module.output_quantizer.num_bits) + + if (4, 3) in num_bits_list: + return KV_CACHE_FP8 + elif 8 in num_bits_list: + return KV_CACHE_INT8 + else: + return QUANTIZATION_NONE + + +def get_weight_block_size(module: nn.Module) -> int: + """Returns the weight block size.""" + if not hasattr(module, "weight_quantizer"): + return 0 + + weight_quantizer = module.weight_quantizer + + if isinstance(weight_quantizer, SequentialQuantizer): + weight_quantizer = weight_quantizer[0] + + if not weight_quantizer.is_enabled: + return 0 + + block_sizes = weight_quantizer.block_sizes + + if block_sizes: + return block_sizes[-1] + return 0 + + +def get_quantization_format(module) -> Optional[str]: + """Gets the quantization string. + + Gets the quantization string by iterating through the module and its children. + The first non-None quantization string is returned. + """ + + def _is_enabled(quantizer): + if isinstance(quantizer, SequentialQuantizer): + return any([_is_enabled(q) for q in quantizer]) + return quantizer.is_enabled + + def _get_quantization_from_linear_layer(layer): + if not hasattr(layer, "weight_quantizer") or not _is_enabled(layer.weight_quantizer): + return QUANTIZATION_NONE + w_quantizer = layer.weight_quantizer + if isinstance(w_quantizer, SequentialQuantizer): + assert ( + len(w_quantizer) == 2 + and w_quantizer[0].num_bits == 4 + and w_quantizer[1].num_bits == (4, 3) + ), "Unsupported quantizer" + assert ( + w_quantizer[0].block_sizes + and len(w_quantizer[0].block_sizes) > 0 + and w_quantizer[0].block_sizes[-1] > 0 + ), "Invalid block_sizes" + return QUANTIZATION_W4A8_AWQ + if w_quantizer.num_bits == 4: + assert ( + len(w_quantizer.block_sizes) > 0 and w_quantizer.block_sizes[-1] > 0 + ), "Invalid block_sizes" + return QUANTIZATION_INT4_AWQ + elif w_quantizer.num_bits == 8: + return QUANTIZATION_INT8_SQ + elif w_quantizer.num_bits == (4, 3): + return QUANTIZATION_FP8 + elif ( + w_quantizer.num_bits == (2, 1) + and hasattr(layer, "input_quantizer") + and hasattr(layer.input_quantizer, "_pre_quant_scale") + ): + return QUANTIZATION_NVFP4_AWQ + elif w_quantizer.num_bits == (2, 1): + return QUANTIZATION_NVFP4 + else: + raise NotImplementedError( + f"Unsupported quantizer with num_bits: {w_quantizer.num_bits}" + ) + + if is_quantized_linear(module): + return _get_quantization_from_linear_layer(module) + + for _, layer in module.named_children(): + if is_quantized_linear(layer): + quantization = _get_quantization_from_linear_layer(layer) + else: + quantization = get_quantization_format(layer) + + # Try to see if other layers has quantization + if quantization != QUANTIZATION_NONE: + return quantization + + return QUANTIZATION_NONE + + +def process_layer_quant_config(layer_config_dict): + """Processes per layer quantization information for TRTLLM export to quant_cfg.json.""" + per_layer_config: dict[str, Any] = { + "quant_algo": None, + "kv_cache_quant_algo": None, + "quantized_layers": {}, + } + layer_config: dict[str, Any] = {} + # Set of quantization formats used. + quantization_formats = set() + + for k, v in layer_config_dict.items(): + if "awq_block_size" in k: + continue + + # Get layer name for constructing quantized_layers dictionary under per_layer_config + prefix = ".".join(k.rsplit(".", 1)[:-1]) + quantization_formats.add(v) + if v == "fp8": + layer_config = {"quant_algo": "FP8"} + elif v == "int4_awq": + layer_config = { + "quant_algo": "W4A16_AWQ", + "group_size": layer_config_dict[prefix + ".awq_block_size"], + "has_zero_point": False, + "pre_quant_scale": True, + } + elif v == "w4a8_awq": + layer_config = { + "quant_algo": "W4A8_AWQ", + "group_size": layer_config_dict[prefix + ".awq_block_size"], + "has_zero_point": False, + "pre_quant_scale": True, + } + elif v == "int8_sq": + layer_config = {"quant_algo": "W8A8_SQ_PER_CHANNEL"} + elif v == "nvfp4": + layer_config = { + "quant_algo": "NVFP4", + "group_size": layer_config_dict[prefix + ".awq_block_size"], + } + else: + layer_config = {"quant_algo": v} + + per_layer_config["quantized_layers"].update({prefix: layer_config}) + + # If we have more than one quantization format, infer MIXED_PRECISION + quantization_formats.discard(None) + + if len(quantization_formats) > 1: + per_layer_config["quant_algo"] = "MIXED_PRECISION" + else: + # We return empty dictionary if we do not have more than one quantization format as + # per layer quantization information is redundant + per_layer_config = {} + + return per_layer_config + + +def to_quantized_weight( + weight: torch.Tensor, + weights_scaling_factor: torch.Tensor, + quantization: str, + weights_scaling_factor2: Optional[torch.Tensor] = None, + block_size: Optional[int] = None, +): + """Converts the weight to the quantized (packed) format.""" + if weights_scaling_factor is not None: + weights_scaling_factor = weights_scaling_factor.to(weight.device) + + if weights_scaling_factor2 is not None: + weights_scaling_factor2 = weights_scaling_factor2.to(weight.device) + + if quantization == QUANTIZATION_FP8: + if weight.dim() == 3: + # for MOE stacked weights + return (weight / weights_scaling_factor.unsqueeze(-1)).to(torch.float8_e4m3fn) + return (weight / weights_scaling_factor).to(torch.float8_e4m3fn) + + if quantization == QUANTIZATION_INT8_SQ: + return (weight / weights_scaling_factor[:, None]).round().clamp(-128, 127).to(torch.int8) + + if quantization in [QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ]: + out_dim = weight.shape[-2] + assert ( + out_dim % 2 == 0 + ), f"Cannot pack weight. Out dimension {out_dim} is not an even number." + in_dim = weight.shape[-1] + block_size = weight.shape[-1] // weights_scaling_factor.shape[-1] + int8_tensor = ( + (weight / weights_scaling_factor[..., :, torch.arange(in_dim) // block_size]) + .round() + .clamp(-8, 7) + .to(torch.int8) + ) + + if int8_tensor.dim() == 3: + # Case of MoE, where weights are stacked + transpose = int8_tensor.permute(0, 2, 1) # (experts, in_dim, out_dim) + int8_tensor = transpose.reshape( + -1, + in_dim, + out_dim // 2, + 2, + ) + int4x2_tensor = (int8_tensor[..., 0] & 0x0F) | (int8_tensor[..., 1] << 4) + # The shape of the returned weight is (experts, out_dim // 2, in_dim) + return int4x2_tensor.permute(0, 2, 1).contiguous() + + int8_tensor = int8_tensor.T.reshape(in_dim, out_dim // 2, 2) # (in_dim, out_dim) + int4x2_tensor = (int8_tensor[..., 0] & 0x0F) | (int8_tensor[..., 1] << 4) + # The shape of the returned weight is (out_dim // 2, in_dim) + return int4x2_tensor.T.contiguous() + + if quantization in [QUANTIZATION_NVFP4, QUANTIZATION_NVFP4_AWQ]: + assert block_size is not None, "Block size not passed. Unable to quantize to NVFP4 format." + assert ( + weights_scaling_factor2 is not None + ), "Weights scaling factor 2 not passed. Unable to quantize to NVFP4 format" + # If MoE reshape weights_scaling_factor2 to enable quantize operations + return NVFP4QTensor.quantize( + weight, + block_size, + weights_scaling_factor, + weights_scaling_factor2.view(-1, 1, 1) + if weights_scaling_factor2.dim() != 0 + else weights_scaling_factor2, + )[0]._quantized_data + + raise NotImplementedError(f"quantization format {quantization} not supported") + + +def from_quantized_weight( + weight: torch.Tensor, + weights_scaling_factor: torch.Tensor, + quantization: str, + torch_dtype, +): + """Converts the quantized weight to the target torch_dtype format.""" + if weight.element_size() >= 2 or weights_scaling_factor is None or not quantization: + # No need to unquantize the weight. + return weight.to(torch_dtype) + + if quantization == QUANTIZATION_FP8: + # safe tensors does not support fp8 yet. So we pack the tensors as int8 + return weight.view(torch.float8_e4m3fn).to(torch_dtype) * weights_scaling_factor.to( + torch_dtype + ) + + if quantization == QUANTIZATION_INT8_SQ: + return weight.to(torch_dtype) * weights_scaling_factor[:, None].to(torch_dtype) + + raise NotImplementedError(f"quantization format {quantization} not supported") + + +def postprocess_state_dict(state_dict: dict, maxbound: float, quantization: Optional[str]) -> dict: + """Filters out keys related to weight quantizers and updates KV cache related keys. + + Args: + state_dict (dict): The full model state_dict. + maxbound (float): The maximum bound value for the output quantizer. + quantization (Optional[str]): The KV cache quantization format. + + Returns: + dict: Filtered state_dict without unnecessary keys like '_amax' and non KV cache output quantizers. + """ + replacements = { + "k_proj.output_quantizer._amax": "k_proj.k_scale", + "v_proj.output_quantizer._amax": "v_proj.v_scale", + "input_quantizer._pre_quant_scale": "pre_quant_scale", + } + + post_state_dict = {} + + for key, value in state_dict.items(): + # Skip keys not related to quantizers + if "output_quantizer" not in key and "_amax" not in key: + post_state_dict[key] = value + continue + + # Apply replacements if the key matches any suffix in the replacements dict + for old_suffix, new_suffix in replacements.items(): + if key.endswith(old_suffix): + assert quantization == KV_CACHE_FP8, "Invalid quantization format." + assert maxbound > 0, "Maxbound must be greater than zero." + + prefix = key[: -len(old_suffix)] + if "_amax" in key: + value = value / maxbound + + # Warn if scale exceeds threshold + if value.item() > 0.5: + warnings.warn( + "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop. " + "Setting KV cache scaling factor to at least 1." + ) + + # Ensure scale is at least 1 + value.clamp_(min=1.0) + + post_state_dict[prefix + new_suffix] = value + break + + # Squeeze tensors with a leading dimension of 1 + for key, value in post_state_dict.items(): + if isinstance(value, torch.Tensor) and value.dim() == 3 and value.shape[0] == 1: + post_state_dict[key] = value.squeeze(0) + + return post_state_dict + + +def all_items_same(item_list): + """Checks if all elements in the provided list are the same.""" + return all(x == item_list[0] for x in item_list) + + +def _fuse_prequant_layernorm( + layernorm_module: torch.nn.Module, + modules: list[torch.Tensor], +): + """Scales layernorm weights with avg_pre_quant_scale of the modules list and sets pre_quant_scales to be deleted.""" + layernorm_module.weight = torch.nn.Parameter( + layernorm_module.weight * getattr(modules[0].input_quantizer, "_pre_quant_scale") + ) + # Pre_quant_scales of modules must not be exported, since they have been fused with layernorm + for module in modules: + delattr(module.input_quantizer, "_pre_quant_scale") + + +def preprocess_linear_fusion(modules: list[torch.nn.Module], resmooth_only=False): + """Preprocess the quantized linears that we plan to fuse. + + Use resmooth_only for MOE experts as each individual expert is not fused. + """ + quantization_format_list = [get_quantization_format(module) for module in modules] + assert all_items_same(quantization_format_list), "Modules have different quantization formats" + + # Activation + if hasattr(modules[0], "input_quantizer"): + # Resmooth + if modules[0].input_quantizer.pre_quant_scale is not None: + avg_prequant_scale = torch.mean( + torch.stack([module.input_quantizer.pre_quant_scale for module in modules]), + dim=0, + ) + + for module in modules: + if not torch.equal(module.input_quantizer.pre_quant_scale, avg_prequant_scale): + module.weight = nn.Parameter( + module.weight + * module.input_quantizer.pre_quant_scale.type(module.weight.dtype) + / avg_prequant_scale.type(module.weight.dtype) + ) + module.input_quantizer.pre_quant_scale = avg_prequant_scale + + # Redo weights collection + module.weight_quantizer.reset_amax() + enable_stats_collection(module.weight_quantizer) + module.weight_quantizer(module.weight) + finish_stats_collection(module.weight_quantizer) + + if resmooth_only: + return + + if modules[0].input_quantizer.is_enabled: + assert ( + modules[0].input_quantizer.amax.numel() == 1 + ), "Only support scalar input quant amax" + + input_amax = torch.max(torch.stack([module.input_quantizer.amax for module in modules])) + for module in modules: + module.input_quantizer.amax = input_amax + + # Weight + if hasattr(modules[0], "weight_quantizer"): + is_seq_quant = isinstance(modules[0].weight_quantizer, SequentialQuantizer) + + if is_seq_quant: + if modules[0].weight_quantizer[-1].is_enabled: + assert len(modules[0].weight_quantizer) == 2 + weight_amax = torch.max( + torch.stack([module.weight_quantizer[-1].amax for module in modules]) + ) + for module in modules: + module.weight_quantizer[-1].amax = weight_amax + + elif ( + modules[0].weight_quantizer.is_enabled and modules[0].weight_quantizer.amax.numel() == 1 + ): + weight_amax = torch.max( + torch.stack([module.weight_quantizer.amax for module in modules]) + ) + for module in modules: + module.weight_quantizer.amax = weight_amax diff --git a/modelopt/torch/export/tensorrt_llm_type.py b/modelopt/torch/export/tensorrt_llm_type.py new file mode 100644 index 0000000000..948138de70 --- /dev/null +++ b/modelopt/torch/export/tensorrt_llm_type.py @@ -0,0 +1,42 @@ +# 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. + +"""Code from TRT-LLM that export optimized models to the TensorRT-LLM checkpoint.""" + +from enum import IntEnum + + +# These clasess are directly copied from TRT-LLM to relex TRT-LLM dependency for checkpoint export +class LayerNormType(IntEnum): + """LayerNormType from tensorrt_llm.functional.""" + + LayerNorm = 0 + RmsNorm = 1 + GroupNorm = 2 + + +class LayerNormPositionType(IntEnum): + """LayerNormPositionType from tensorrt_llm.functional.""" + + pre_layernorm = 0 + post_layernorm = 1 + + +class MLPType(IntEnum): + """MLPType from tensorrt_llm.functional.""" + + MLP = 0 + GatedMLP = 1 + FusedGatedMLP = 2 diff --git a/modelopt/torch/export/tensorrt_llm_utils.py b/modelopt/torch/export/tensorrt_llm_utils.py new file mode 100755 index 0000000000..4241ce195d --- /dev/null +++ b/modelopt/torch/export/tensorrt_llm_utils.py @@ -0,0 +1,566 @@ +# 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. + +"""Utils for TensorRT-LLM checkpoint export. + +Some of the logics in this file are empirical and needs constant update if exceptions occur. +""" + +import logging +from pathlib import Path +from typing import TYPE_CHECKING, Any, Generator, Iterable, Optional +from warnings import warn + +from .tensorrt_llm_type import LayerNormPositionType, LayerNormType, MLPType + +if TYPE_CHECKING: + from transformers import T5Config + +from modelopt import __version__ + +from .model_config import ( + LAYERNORM_DEFAULT, + LAYERNORM_RMS, + QUANTIZATION_NONE, + DecoderLayerConfig, + MLPConfig, + ModelConfig, + MOEConfig, +) + +logger = logging.getLogger(__name__) + +# For NEMO and ENC/DEC models where the TensorRT-LLM model architecture and HF config are not aligned. +MODEL_NAME_TO_HF_ARCH_MAP = { + "llama": "LlamaForCausalLM", + "gemma": "GemmaForCausalLM", + "gpt": "GPTForCausalLM", + "t5_encoder": "EncoderModel", + "t5_decoder": "DecoderModel", + "mllama": "MLLaMAModel", +} + + +def is_tensorrt_llm_0_8_or_9(): + """Returns true if tensorrt_llm version is 0.8 or 0.9.""" + try: + import tensorrt_llm + + return tensorrt_llm.__version__.startswith(("0.8", "0.9")) + except Exception: + return False + + +def _find_layernorm_type(model_config: ModelConfig): + if model_config.ln_f: + return model_config.ln_f.layernorm_type + for layer in model_config.layers: + if layer.input_layernorm: + return layer.input_layernorm.layernorm_type + if layer.post_layernorm: + return layer.post_layernorm.layernorm_type + return LAYERNORM_DEFAULT + + +def _prefix_wildcard_summarize_exclude_modules(unquantized_layers, quantized_layers): + """Generate a summarization of the quantization layer configs using prefix wildcards. + + Prefix wildcards means we only consider wildcards that is a prefix with a star in the end. + We do not consider other wildcards such as: a*b. + """ + + def all_matching_prefix_wildcards(name): + # include all possible prefix wildcards, and the exact name itself + wildcards = set([name]) + for i in range(len(name) + 1): + wildcards.add(name[:i] + "*") + return wildcards + + def next_formated_matching_prefix_wildcards(name: str) -> Generator[list[str], None, None]: + """Enumerate formated prefix wildcards. A result may be a combination of prefix wildcards. + + Formated here means we only consider wildcards at dot split. We need two patterns. + + 1. a single wildcard: module_name* + 2. a set of 2 wildcards: {module_name, module_name.*}. We need this pattern set because + module_name* may match other modules with module_name as a prefix. + """ + for i in range(len(name)): + if name[i] == ".": + yield [name[:i] + "*"] + yield [name[:i], name[:i] + ".*"] + # in the end, itself only is a wildcard + yield [name] + + # any of the wildcard in this set cannot be present in the result + negtive_wild_candidates = set() + for layer in quantized_layers: + negtive = all_matching_prefix_wildcards(layer) + negtive_wild_candidates.update(negtive) + logger.debug( + f"Quantized layer {layer}, prefix wildcards {negtive} identified as negative wildcards" + ) + + res_summary = set() + for layer in unquantized_layers: + candidate_wildcards = [] + for wildcards in next_formated_matching_prefix_wildcards(layer): + if any([wildcard in negtive_wild_candidates for wildcard in wildcards]): + # need a more specific wildcard + logger.debug( + f"Unquantized layer {layer}, prefix wildcards {wildcards} invalidated by negative wildcards" + ) + continue + if all([wildcard in res_summary for wildcard in wildcards]): + # we get covered already, do not need to move forward, and clear candidate + logger.debug( + f"Unquantized layer {layer}, prefix wildcards {wildcards} already covered" + ) + candidate_wildcards = [] + break + # find one, now terminate the search + candidate_wildcards = wildcards + logger.debug( + f"Unquantized layer {layer}, prefix wildcards {wildcards} identified as a new match" + ) + break + res_summary.update(candidate_wildcards) + return res_summary + + +def _detect_exclude_modules(weight_keys: Iterable[str]) -> list[str]: + quantized_layers = set() + unquantized_layers = set() + + for key in weight_keys: + suffix = key.split(".")[-1] + + # Filter kv_cache_scaling_factor. + if "kv_cache_scaling_factor" in suffix: + continue + + if "_scaling_factor" in suffix: + quantized_layers.add(key.rsplit(".", 1)[0]) + else: + unquantized_layers.add(key.rsplit(".", 1)[0]) + + logger.debug(f"pre-unquantized_layers {unquantized_layers}") + logger.debug(f"pre-quantized_layers {quantized_layers}") + + unquantized_layers = unquantized_layers - quantized_layers + + logger.debug(f"unquantized_layers {unquantized_layers}") + logger.debug(f"quantized_layers {quantized_layers}") + + res_with_wildcards = _prefix_wildcard_summarize_exclude_modules( + unquantized_layers, quantized_layers + ) + return list(res_with_wildcards) + + +def _get_block_size(model_config: ModelConfig): + """Return the first block size that is not zero if any.""" + for layer in model_config.layers: + module_list = [layer.attention.qkv] if layer.attention else [] + if layer.mlp and isinstance(layer.mlp, MLPConfig): + module_list.extend( + [ + layer.mlp.fc, + layer.mlp.proj, + layer.mlp.gate, + ] + ) + if layer.mlp and isinstance(layer.mlp, MOEConfig): + module_list.extend([layer.mlp.experts.fc, layer.mlp.experts.proj]) + for m in module_list: + if m is not None and m.awq_block_size != 0: + return m.awq_block_size + return 0 + + +def _get_quant_config( + model_config: ModelConfig, + weight_keys: Iterable[str], + kv_cache_dtype: Optional[str] = None, + pp_size: int = 1, +) -> dict[str, Any]: + quantization: dict[str, Any] = {"quant_algo": None, "kv_cache_quant_algo": None} + + if model_config.quantization == "fp8": + quantization.update({"quant_algo": "FP8"}) + elif model_config.quantization == "int4_awq": + quantization.update( + { + "quant_algo": "W4A16_AWQ", + "group_size": _get_block_size(model_config), + "has_zero_point": False, + "pre_quant_scale": True, + } + ) + elif model_config.quantization == "w4a8_awq": + quantization.update( + { + "quant_algo": "W4A8_AWQ", + "group_size": _get_block_size(model_config), + "has_zero_point": False, + "pre_quant_scale": True, + } + ) + elif model_config.quantization == "int8_sq": + quantization.update( + { + "quant_algo": "W8A8_SQ_PER_CHANNEL", + } + ) + elif model_config.quantization == "nvfp4": + quantization.update( + { + "quant_algo": "NVFP4", + "group_size": _get_block_size(model_config), + } + ) + elif model_config.quantization == "nvfp4_awq": + quantization.update( + { + "quant_algo": "NVFP4_AWQ", + "group_size": _get_block_size(model_config), + "has_zero_point": False, + "pre_quant_scale": True, + } + ) + elif model_config.quantization == QUANTIZATION_NONE: + quantization.update( + { + "quant_algo": None, + } + ) + else: + quantization.update( + { + "quant_algo": model_config.quantization, + } + ) + + # Deprecate exclude modules for per layer export + if model_config.quantization != QUANTIZATION_NONE: + exclude_modules = _detect_exclude_modules(weight_keys) + # In TRT LLM, the embedding table is shared for the following models, so lm_head quantization format + # won't be automatically detected in the excluded_modules. We need to manually add it to the exclusions. + share_embedding_table = ( + True if (model_config.lm_head is None and pp_size == 1) else False, + ) + if share_embedding_table: + exclude_modules.append("lm_head") + quantization["exclude_modules"] = exclude_modules + + if kv_cache_dtype is not None: + quantization.update( + { + "kv_cache_quant_algo": kv_cache_dtype, + } + ) + + return quantization + + +def convert_to_tensorrt_llm_config( + model_config: ModelConfig, + weight_keys: Iterable[str], + hf_config=None, +): + """Convert to TensorRT-LLM checkpoint config. + + Args: + model_config: The model_config to convert. + weight_keys: The iterable of string of weights exported to the tensorrt_llm checkpoint. + hf_config: The huggingface model config. + If provided, we try to use the TensorRT-LLM's export method if available. + """ + tp_size = model_config.tensor_parallel + pp_size = model_config.pipeline_parallel + + first_attention_config = None + first_attention_decoder_config = None + for decoder_layer in model_config.layers: + first_attention_config = ( + decoder_layer.attention or decoder_layer.self_attention or decoder_layer.cross_attention + ) + + if first_attention_config is not None: + first_attention_decoder_config = decoder_layer + break + + assert ( + first_attention_config is not None and first_attention_decoder_config is not None + ), "Model must have at least one attention block" + + mapping = { + "world_size": tp_size * pp_size, + "tp_size": tp_size, + "pp_size": pp_size, + } + quantization = _get_quant_config( + model_config, weight_keys, first_attention_config.kv_cache_dtype, pp_size + ) + + decoder_type = model_config.layers[0].decoder_type + + config_architecture = model_config.architecture + if not config_architecture: + config_architecture = MODEL_NAME_TO_HF_ARCH_MAP[decoder_type] + # For T5 model + if decoder_type in ["t5"]: + # For encoder + if model_config.enc_dec == "enc": + config_architecture = MODEL_NAME_TO_HF_ARCH_MAP["t5_encoder"] + # For decoder + else: + config_architecture = MODEL_NAME_TO_HF_ARCH_MAP["t5_decoder"] + + config = { + "producer": { + "name": "modelopt", + "version": __version__, + }, + "architecture": config_architecture, + "dtype": model_config.dtype, + "logits_dtype": "float16" if model_config.dtype == "bfloat16" else model_config.dtype, + "num_hidden_layers": len(model_config.layers) * pp_size, + "num_attention_heads": model_config.num_attention_heads, + "num_key_value_heads": model_config.num_kv_heads, + "hidden_size": model_config.hidden_size, + "norm_epsilon": ( + first_attention_decoder_config.mlp_layernorm.eps + if decoder_type in ["t5"] + else first_attention_decoder_config.input_layernorm.eps + ), + "vocab_size": model_config.vocab_size, + "max_position_embeddings": model_config.max_position_embeddings, + "hidden_act": model_config.hidden_act, + "use_parallel_embedding": True, + "embedding_sharding_dim": 0, + "head_size": first_attention_decoder_config.attention_head_size, + "intermediate_size": first_attention_decoder_config.ffn_hidden_size_local * tp_size, + "position_embedding_type": ( + "alibi" if first_attention_decoder_config.use_alibi else "rope_gpt_neox" + ), + "share_embedding_table": True if (model_config.lm_head is None and pp_size == 1) else False, + "residual_mlp": first_attention_decoder_config.residual_mlp is not None, + # Model Optimizer customized fields + "bias": first_attention_config.dense.bias is not None, + "rotary_pct": first_attention_decoder_config.rotary_pct, + "rank": model_config.rank, + "decoder": first_attention_decoder_config.decoder_type, + "rmsnorm": _find_layernorm_type(model_config) == LAYERNORM_RMS, + "lm_head_bias": model_config.lm_head is not None and model_config.lm_head.bias is not None, + } + + if hf_config is not None: + try: + from tensorrt_llm.models import MODEL_MAP + from tensorrt_llm.models.modeling_utils import Mapping, QuantConfig + + model_cls = MODEL_MAP[model_config.architecture] + config_cls = getattr(model_cls, "config_class", None) + + if config_cls is not None: + tensorrt_llm_config = config_cls.from_hugging_face( + hf_config_or_dir=hf_config, + dtype=model_config.dtype, + mapping=Mapping.from_dict(mapping), + quant_config=QuantConfig.from_dict(quantization), + ).to_dict() + + config.update(tensorrt_llm_config) + # Modelopt uses parallel_embedding by default. + config.update( + { + "use_parallel_embedding": True, + "embedding_sharding_dim": 0, + "logits_dtype": "float16" + if model_config.dtype == "bfloat16" + else model_config.dtype, + } + ) + + return config + + except Exception as e: + warn( + "Cannot export tensorrt_llm checkpoint config due to the following error. " + f"Trying the manual approach: \n{e}" + ) + + if first_attention_decoder_config.rotary_base: + config["rotary_base"] = first_attention_decoder_config.rotary_base + + if first_attention_decoder_config.rope_scaling: + config["rotary_scaling"] = first_attention_decoder_config.rope_scaling + + config["mapping"] = mapping + config["quantization"] = quantization + + layernorm_type_map = {i.name: i.value for i in LayerNormType} + layernorm_position_map = {i.name: i.value for i in LayerNormPositionType} + + if decoder_type == "mpt": + config.update( + { + "clip_qkv": first_attention_config.clip_qkv, + "alibi_bias_max": model_config.layers[0].alibi_bias_max, + } + ) + elif decoder_type == "recurrentgemma": + config["conv_kernel"] = 4 + config["state_size"] = 1 + config["state_dtype"] = "float32" + config["rnn_hidden_size"] = model_config.layers[0].rnn_hidden_size + config["rnn_conv_dim_size"] = model_config.layers[0].rnn_hidden_size + config["logits_soft_cap"] = model_config.layers[0].logits_soft_cap + config["emb_scale_by_sqrt_dim"] = model_config.layers[0].emb_scale_by_sqrt_dim + config["layer_types"] = model_config.layers[0].layer_types + elif decoder_type == "t5": + config["position_embedding_type"] = "relative" + config["share_embedding_table"] = getattr(model_config, "share_embedding_table") + config["has_position_embedding"] = ( + False if not getattr(model_config, "position_embedding") else True + ) + # fallback to RmsNorm if not specified as HF might change class naming for layernorm + layernorm_type = model_config.layers[0].mlp_layernorm.layernorm_type + if not layernorm_type: + layernorm_type = "RmsNorm" + config["layernorm_type"] = layernorm_type_map[layernorm_type] + config["has_attention_qkvo_bias"] = ( + False + if not ( + model_config.layers[0].attention.qkv.bias + if model_config.enc_dec == "enc" + else model_config.layers[0].self_attention.qkv.bias + ) + else True + ) + config["has_mlp_bias"] = False if not model_config.layers[0].mlp.fc.bias else True + config["has_model_final_layernorm"] = True if model_config.ln_f else False + + mlp_type_map = {i.name: i.value for i in MLPType} + + config["mlp_type"] = mlp_type_map[ + ( + "GatedMLP" + if isinstance(model_config.layers[0].mlp, MLPConfig) + and model_config.layers[0].mlp.gate + else "MLP" + ) + ] + config["use_prompt_tuning"] = False + config["has_position_embedding"] = False if not model_config.position_embedding else True + config["has_embedding_layernorm"] = False if not model_config.ln_embed else True + config["has_embedding_scale"] = False + config["ffn_hidden_size"] = model_config.layers[0].mlp.fc.weight.shape[0] + config["q_scaling"] = 1 / config["head_size"] ** 0.5 + config["layernorm_position"] = layernorm_position_map["pre_layernorm"] + config["relative_attention"] = config["position_embedding_type"] == "relative" + config["max_distance"] = model_config.layers[0].rel_attn_max_distance + config["num_buckets"] = model_config.layers[0].rel_attn_num_buckets + config["model_type"] = "t5" + config["use_parallel_embedding"] = True + config["use_implicit_relative_attention"] = False + if model_config.enc_dec == "dec": + config["rescale_before_lm_head"] = False + config["encoder_hidden_size"] = model_config.encoder_hidden_size + config["encoder_num_heads"] = model_config.encoder_num_heads + config["encoder_head_size"] = model_config.encoder_head_size + config["skip_cross_kv"] = False + + elif decoder_type == "dbrx": + config["clip_qkv"] = first_attention_decoder_config.clip_qkv + + elif decoder_type == "mllama": + num_layers = config["num_hidden_layers"] + num_kv_heads = model_config.num_kv_heads + cross_attention_layers = set(model_config.layers[0].cross_attention_layers) + + config["num_kv_heads_per_layer"] = [ + 0 if i in cross_attention_layers else num_kv_heads for i in range(num_layers) + ] + config["num_kv_heads_per_cross_attn_layer"] = [ + num_kv_heads if i in cross_attention_layers else 0 for i in range(num_layers) + ] + + config["cross_attention"] = True + config["cross_attention_layers"] = model_config.layers[0].cross_attention_layers + config["embed_vocab_size"] = model_config.vocab_size + 8 + vision_output_dim = model_config.layers[0].vision_output_dim + config["vision_output_dim"] = vision_output_dim if vision_output_dim != 0 else 7680 + else: + raise NotImplementedError( + f"Cannot export tensorrt_llm checkpoint for model {config_architecture}. " + "It's not supported by TensorRT-LLM yet." + ) + + # For Mixtral and Arctic + if first_attention_decoder_config.moe_num_experts: + config["moe"] = { + "num_experts": first_attention_decoder_config.moe_num_experts, + "top_k": first_attention_decoder_config.moe_top_k, + "normalization_mode": 1, # ExpertScaleNormalizationMode.RENORMALIZE + } + + if decoder_type == "phi3": + config["moe"]["normalization_mode"] = 2 # ExpertScaleNormalizationMode.SPARSE_MIXER, + config["moe"]["sparse_mixer_epsilon"] = ( + first_attention_decoder_config.sparse_mixer_epsilon + ) + + config["mapping"]["moe_tp_size"] = config["mapping"]["tp_size"] + config["mapping"]["moe_ep_size"] = 1 + + # Handle Medusa decoding + # TODO (chenhany): when inference pp > 1; only last pp has medusa heads + if model_config.medusa_heads is not None: + config["base_architecture"] = config["architecture"] + config["architecture"] = "MedusaForCausalLM" + # NOTE: max_draft_len is related to the medusa tree len. Currently it is hardcoded to 63. + config["max_draft_len"] = 63 + config["num_medusa_heads"] = len(model_config.medusa_heads) + config["num_medusa_layers"] = len(model_config.medusa_heads[0].medusa_layers) + + return config + + +def prepare_enc_dec_export_dir(tensorrt_llm_config: dict[str, Any], export_root: Path): + """Prepare the export directory for encoder-decoder model.""" + # For encoder + if tensorrt_llm_config["architecture"] == "EncoderModel": + export_dir = export_root.joinpath("encoder") + # For decoder + else: + export_dir = export_root.joinpath("decoder") + return export_dir + + +def prepare_enc_dec_decoder_layer( + layer_config: DecoderLayerConfig, + model_config: "T5Config", + enc_dec: str, + layers: list[DecoderLayerConfig], +): + """Prepare the config for each decoder layer of encoder-decoder model.""" + layer_config.rel_attn_max_distance = model_config.relative_attention_max_distance + layer_config.rel_attn_num_buckets = model_config.relative_attention_num_buckets + if enc_dec == "enc" and layer_config.attention.rel_attn_table is None: + layer_config.attention.rel_attn_table = layers[0].attention.rel_attn_table + elif enc_dec == "dec" and layer_config.self_attention.rel_attn_table is None: + layer_config.self_attention.rel_attn_table = layers[0].self_attention.rel_attn_table diff --git a/modelopt/torch/export/transformer_engine.py b/modelopt/torch/export/transformer_engine.py new file mode 100644 index 0000000000..26ea66f236 --- /dev/null +++ b/modelopt/torch/export/transformer_engine.py @@ -0,0 +1,94 @@ +# Adapted from https://github.com/huggingface/accelerate/blob/c2f193a25c5449c0e34408bbbffa61ce420b00f6/ +# src/accelerate/utils/transformer_engine.py + +# Copyright 2021 The HuggingFace Team. All rights reserved. +# +# 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. + +# 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. + +"""Convert the Model Optimizer quantized model to the transformer_engine.""" + +import torch +from torch import nn + +__all__ = ["convert_to_transformer_engine"] + + +def _convert_model(model): + """Recursively converts the linear layers to their `transformers_engine` counterpart.""" + import transformer_engine.pytorch as te + + with te.fp8_model_init(enabled=True): + for name, module in model.named_children(): + if type(module).__name__ == "QuantLinear": + # Return early if the linear layer weights are not multiples of 16 + if any(p % 16 != 0 for p in module.weight.shape): + return + has_bias = module.bias is not None + te_module = te.Linear( + module.in_features, + module.out_features, + bias=has_bias, + params_dtype=module.weight.dtype, + ) + te_module.weight.copy_(module.weight) + if has_bias: + te_module.bias.copy_(module.bias) + + input_quantizer = module.input_quantizer + quant_max = input_quantizer.maxbound + input_scale_inv = input_quantizer.export_amax().float()[0] / quant_max + input_scale = 1 / input_scale_inv + + weight_quantizer = module.weight_quantizer + quant_max = weight_quantizer.maxbound + weight_scale_inv = weight_quantizer.export_amax().float()[0] / quant_max + weight_scale = 1 / weight_scale_inv + + # Update the TE layers with the Model Optimizer scaling factors + te_module.fp8_meta["scaling_fwd"].scale_inv[0].copy_(input_scale_inv) + te_module.fp8_meta["scaling_fwd"].scale_inv[1].copy_(weight_scale_inv) + te_module.fp8_meta["scaling_fwd"].scale[0].copy_(input_scale) + te_module.fp8_meta["scaling_fwd"].scale[1].copy_(weight_scale) + + setattr(model, name, te_module) + else: + _convert_model(module) + + +def convert_to_transformer_engine(model: nn.Module): + """Converts the `Model Optimizer` quantized model to the `transformers_engine`.""" + import transformer_engine.common.recipe as te_recipe + from transformer_engine.pytorch import fp8_autocast + + with torch.no_grad(): + _convert_model(model) + + fp8_recipe = te_recipe.DelayedScaling() + model.forward = fp8_autocast(enabled=True, fp8_recipe=fp8_recipe)(model.forward) + + return model diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py new file mode 100644 index 0000000000..98d9969946 --- /dev/null +++ b/modelopt/torch/export/unified_export_hf.py @@ -0,0 +1,363 @@ +# 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. + +"""Code that export quantized Hugging Face models for deployment.""" + +import json +import tempfile +import warnings +from collections import defaultdict +from pathlib import Path +from typing import Any, Optional, Union + +import torch +import torch.nn as nn + +from modelopt import __version__ + +from .layer_utils import get_experts_list, is_layernorm, is_moe, is_quantlinear +from .model_config import ( + QUANTIZATION_FP8, + QUANTIZATION_NONE, + QUANTIZATION_NVFP4, + QUANTIZATION_NVFP4_AWQ, + QUANTIZATION_W4A8_AWQ, +) +from .quant_utils import ( + _fuse_prequant_layernorm, + get_activation_scaling_factor, + get_kv_cache_dtype, + get_quantization_format, + get_weight_block_size, + get_weight_scaling_factor, + get_weight_scaling_factor_2, + postprocess_state_dict, + preprocess_linear_fusion, + process_layer_quant_config, + to_quantized_weight, +) + +__all__ = ["export_hf_checkpoint"] + +SPECULATIVE_DECODING_MODULE_NAMES = ["medusa_heads", "eagle_module", "drafter"] + + +def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): + """Group modules that take the same input and register shared parameters in module.""" + # TODO: Handle DBRX MoE + input_to_linear = defaultdict(list) + output_to_layernorm = defaultdict(None) + quantization_format = get_quantization_format(model) + + def _input_hook(module, input, output): + """Update dictionary with list of all modules that share the same input.""" + # TODO: Handle DBRX MoE case + input_to_linear[input[0]].append(module) + + def _output_hook(module, input, output): + """Update dictionary with mapping of layernorms and their outputs.""" + output_to_layernorm[output] = module + + handles = [] + model_type = type(model).__name__.lower() + + for name, module in model.named_modules(): + # For MoE models update pre_quant_scale to average pre_quant_scale amongst experts + if is_moe(module) and ("awq" in quantization_format): + # update_experts_avg_prequant_scale(module) + grouped_experts = get_experts_list(module, model_type) + for modules in grouped_experts: + preprocess_linear_fusion(modules, resmooth_only=True) + + # Attach hook to layernorm modules that need to be fused + if is_layernorm(module): + module.name = name + handle = module.register_forward_hook(_output_hook) + handles.append(handle) + + elif "QuantLinear" in type(module).__name__: + module.name = name + handle = module.register_forward_hook(_input_hook) + handles.append(handle) + + with torch.no_grad(): + fake_input = torch.ones([1, 1], dtype=torch.long).to(model.device) + # Run forward pass so that all modules sharing the same input are collected using forward hook. + model(fake_input) + for handle in handles: + handle.remove() + + for tensor, modules in input_to_linear.items(): + quantization_format = get_quantization_format(modules[0]) + if len(modules) > 1 and quantization_format not in [QUANTIZATION_FP8, QUANTIZATION_NONE]: + # Fuse modules that have the same input + preprocess_linear_fusion(modules) + + # Fuse layernorms + if tensor in output_to_layernorm.keys() and "awq" in get_quantization_format(modules[0]): + # Pre quant scale of modules is already updated to avg_pre_quant_scale + _fuse_prequant_layernorm(output_to_layernorm[tensor], modules) + + +def _export_hf_checkpoint( + model: nn.Module, dtype: Optional[torch.dtype] = None +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + """Exports the torch model to the packed checkpoint with original HF naming. + + The packed checkpoint will be consumed by the TensorRT-LLM unified converter. + + Args: + model: the torch model. + dtype: the weights data type to export the unquantized layers or the default model data type if None. + + Returns: + post_state_dict: Dict containing quantized weights + quant_config: config information to export hf_quant_cfg.json + per_layer_quantization: Dict containing layerwise quantization information to export quant_cfg.json + in mixed_precision case. + """ + if dtype is None: + dtype = model.config.torch_dtype + else: + warnings.warn( + f"Model's original dtype ({model.config.torch_dtype}) differs from target dtype " + f"({dtype}), which may lead to numerical errors." + ) + + # Base model layers + layer_pool = { + f"model.layers.{name}": sub_module + for name, sub_module in model.model.layers.named_modules() + } + # NOTE: Speculative decoding models have extra modules that may be quantized + # Need to add these modules to the layer_pool + for key in SPECULATIVE_DECODING_MODULE_NAMES: + if hasattr(model, key): + for name, sub_module in getattr(model, key).named_modules(): + layer_pool.update({f"{key}.{name}": sub_module}) + + # Layer config dict holds quantization format of each layer. + # It also holds awq_block_size information for applicable layers. + layer_config_dict = {} + + # Resmooth and requantize fused layers + # TODO: Handle mixed precision + requantize_resmooth_fused_llm_layers(model) + + kv_cache_max_bound = 0 + kv_cache_format = QUANTIZATION_NONE + + for name, sub_module in layer_pool.items(): + if is_quantlinear(sub_module): + quantization_format = get_quantization_format(sub_module) + block_size = get_weight_block_size(sub_module) + + # Construct per layer config dictionary + layer_config_dict.update({name + ".quantization": quantization_format}) + layer_config_dict.update({name + ".awq_block_size": block_size}) + + if quantization_format == QUANTIZATION_FP8: + # Convert amax to float32 + sub_module.weight_quantizer._amax = sub_module.weight_quantizer._amax.to( + torch.float32 + ) + + if sub_module.weight_quantizer._amax.dim() == 1: + weight_scaling_factor = torch.tensor( + sub_module.weight_quantizer.amax.item() + / sub_module.weight_quantizer.maxbound + ) + else: + # Per-channel amax + weight_scaling_factor = torch.tensor( + sub_module.weight_quantizer.amax / sub_module.weight_quantizer.maxbound + ) + + sub_module.register_buffer( + "weight_scale", + weight_scaling_factor, + ) + + if hasattr(sub_module.input_quantizer, "_amax"): + sub_module.input_quantizer._amax = sub_module.input_quantizer._amax.to( + torch.float32 + ) + + sub_module.register_buffer( + "input_scale", + get_activation_scaling_factor(sub_module).squeeze(), + ) + + if hasattr(sub_module.output_quantizer, "_amax"): + sub_module.output_quantizer._amax = sub_module.output_quantizer._amax.to( + torch.float32 + ) + + if quantization_format in [ + QUANTIZATION_NVFP4_AWQ, + QUANTIZATION_NVFP4, + QUANTIZATION_W4A8_AWQ, + ]: + # Register weight_scale_2 + sub_module.register_buffer( + "weight_scale_2", + get_weight_scaling_factor_2(sub_module).reshape(1), + ) + + if quantization_format not in [QUANTIZATION_FP8, QUANTIZATION_NONE]: + # Register weight_scale and input_scale + sub_module.register_buffer("weight_scale", get_weight_scaling_factor(sub_module)) + sub_module.register_buffer("input_scale", get_activation_scaling_factor(sub_module)) + + # Check if quantization format is None, to support auto_quant + if quantization_format != QUANTIZATION_NONE: + quantized_weight = to_quantized_weight( + sub_module.weight.to(dtype), + sub_module.weight_scale, + quantization_format, # type:ignore [arg-type] + sub_module.weight_scale_2 if hasattr(sub_module, "weight_scale_2") else None, + block_size, + ) + sub_module.weight = nn.Parameter(quantized_weight, requires_grad=False) + + # Find kv cache quant format + if kv_cache_format == QUANTIZATION_NONE: + kv_cache_format = get_kv_cache_dtype(sub_module) + else: + assert kv_cache_format == get_kv_cache_dtype( + sub_module + ), "Do not support mixed precision kv cache quantization" + if kv_cache_format != QUANTIZATION_NONE: + kv_cache_max_bound = sub_module.output_quantizer.maxbound + + quantized_state_dict = model.state_dict() + + # Process per layer quantization config dict + per_layer_quantization = process_layer_quant_config(layer_config_dict) + + quantized_state_dict = postprocess_state_dict( + quantized_state_dict, kv_cache_max_bound, kv_cache_format + ) + # Create the quantization config + # TODO: add support for customized mixed precision config + quant_config: dict[str, Any] = { + "producer": { + "name": "modelopt", + "version": __version__, + }, + "quantization": {"quant_algo": None, "kv_cache_quant_algo": None}, + } + + if quantization_format == "fp8": + quant_config["quantization"].update({"quant_algo": "FP8", "exclude_modules": ["lm_head"]}) + # TODO: add info about per-channel weight scale and dynamic per token input scale + elif quantization_format == "int4_awq": + quant_config["quantization"].update( + { + "quant_algo": "W4A16_AWQ", + "group_size": block_size, + "has_zero_point": False, + "pre_quant_scale": True, + "exclude_modules": ["lm_head"], + } + ) + elif quantization_format == "nvfp4": + quant_config["quantization"].update( + { + "quant_algo": "NVFP4", + "group_size": block_size, + "exclude_modules": ["lm_head"], + } + ) + elif quantization_format == "nvfp4_awq": + quant_config["quantization"].update( + { + "quant_algo": "NVFP4_AWQ", + "group_size": block_size, + "has_zero_point": False, + "pre_quant_scale": True, + "exclude_modules": ["lm_head"], + } + ) + else: + quant_config["quantization"].update( + { + "quant_algo": ( + quantization_format if quantization_format != QUANTIZATION_NONE else None + ), + } + ) + + if kv_cache_format is not None: + quant_config["quantization"].update( + { + "kv_cache_quant_algo": kv_cache_format, + } + ) + + return quantized_state_dict, quant_config, per_layer_quantization + + +def export_hf_checkpoint( + model: nn.Module, + dtype: Optional[torch.dtype] = None, + export_dir: Union[Path, str] = tempfile.gettempdir(), +): + """Exports the torch model to unified checkpoint and saves to export_dir. + + Args: + model: the torch model. + dtype: the weights data type to export the unquantized layers or the default model data type if None. + export_dir: the target export path. + """ + export_dir = Path(export_dir) + export_dir.mkdir(parents=True, exist_ok=True) + try: + post_state_dict, hf_quant_config, per_layer_quantization = _export_hf_checkpoint( + model, dtype + ) + + # If auto_quant is used, save per layer quantization information in quant_cfg.json + if per_layer_quantization: + # Update auto quant related information in quantization.json + per_layer_quantization["kv_cache_quant_algo"] = hf_quant_config["quantization"][ + "kv_cache_quant_algo" + ] + + # Update auto quant related informaion in hf_quant_config.json + # We remove group_size, has_zero_point and pre_quant_scale information from config.json + hf_quant_config["quantization"] = { + k: per_layer_quantization[k] for k in ("quant_algo", "kv_cache_quant_algo") + } + + with open(f"{export_dir}/quant_cfg.json", "w") as file: + json.dump(per_layer_quantization, file, indent=4) + + # Save config + with open(f"{export_dir}/hf_quant_config.json", "w") as file: + json.dump(hf_quant_config, file, indent=4) + + # Save model + model.save_pretrained(export_dir, state_dict=post_state_dict) + + except Exception as e: + fallback_model_path = f"{export_dir}/modelopt_model.pth" + torch.save(model.state_dict(), fallback_model_path) + warnings.warn( + "Cannot export model to the model_config. The modelopt-optimized model state_dict" + f" (including the quantization factors) is saved to {fallback_model_path} using" + " torch.save for further inspection." + ) + raise e diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py new file mode 100644 index 0000000000..66d28564f2 --- /dev/null +++ b/modelopt/torch/export/unified_export_megatron.py @@ -0,0 +1,692 @@ +# 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. + +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + + +"""Code that export quantized Megatron Core models for deployment.""" + +import json +import os +import tempfile +from collections import OrderedDict +from pathlib import Path +from typing import Any, Dict, Optional, Tuple, Union + +import torch +from huggingface_hub import snapshot_download +from safetensors.torch import safe_open, save_file + +from modelopt import __version__ +from modelopt.torch.utils import import_plugin + +from .mcore_hf_export_map import all_mcore_hf_export_mapping, all_mcore_hf_import_mapping +from .model_config import QUANTIZATION_FP8 +from .quant_utils import ( + get_activation_scaling_factor, + get_kv_cache_scaling_factor, + get_quantization_format, + get_weight_block_size, + get_weight_scaling_factor, + get_weight_scaling_factor_2, + to_quantized_weight, +) + +with import_plugin("transformers", verbose=False): + import transformers + +has_mcore = False +with import_plugin("megatron"): + from megatron.core.models.gpt import GPTModel + from megatron.core.parallel_state import ( + get_pipeline_model_parallel_rank, + get_pipeline_model_parallel_world_size, + ) + + has_mcore = True + +__all__ = ["export_mcore_gpt_to_hf", "import_mcore_gpt_from_hf"] + + +def get_quantized_state( + module: torch.nn.Module, + dtype: torch.dtype = torch.float16, +) -> Tuple[Dict[str, torch.Tensor], str, int]: + """Return a state_dict, quantization format, and block_size of the module. + + Args: + module: The target module to perform real quantization. + dtype: The default data type. + + Returns: + Tuple: state_dict, quantization format, and block_size of the module. + """ + name_to_value = {} + + weight = module.weight.to(dtype).cpu() + name_to_value["weight"] = weight + + if hasattr(module, "bias") and module.bias is not None: + name_to_value["bias"] = module.bias.to(dtype).cpu() + + qformat: str = get_quantization_format(module) + block_size = get_weight_block_size(module) + + # Getting the weight scales + weight_scale = get_weight_scaling_factor(module) + weight_scale_2 = get_weight_scaling_factor_2(module) + if weight_scale is not None: + name_to_value["weight_scale"] = weight_scale + + if weight_scale_2 is not None: + name_to_value["weight_scale_2"] = weight_scale_2 + + # Getting the input scale + input_scale = get_activation_scaling_factor(module) + if input_scale is not None: + name_to_value["input_scale"] = input_scale + # TODO (chenhany): support AWQ with pre_quant_scale + if hasattr(module.input_quantizer, "_pre_quant_scale"): + raise ValueError("Detect pre_quant_scale! SmoothQuant/AWQ are not yet supported!") + + if hasattr(module, "output_quantizer"): + output_scale = get_kv_cache_scaling_factor([module]) + if output_scale is not None: + name_to_value["output_scale"] = output_scale + + return name_to_value, qformat, block_size + + +class GPTModelExporter: + """Megatron Core GPTModel Exporter. + + The Exporter is created by `export_mcore_gpt_to_hf` to host attributes + and methods that export a quantized Megatron Core GPTModel to the Hugging + Face unified checkpoint. + + Args: + model: The Megatron Core GPTModel instance. + pretrained_model_name_or_path: Can be either: the *model id* of a + pretrained model hosted inside a model repo on huggingface.co; or + a *directory* containing model weights saved using + [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + dtype: The weights data type to export the unquantized layers. + """ + + def __init__( + self, + model: torch.nn.Module, + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, + dtype=torch.bfloat16, + ): + """Create a GPTModel exporter instance.""" + if not isinstance(model, GPTModel): + raise ValueError("Input to GPTModelExport must be a megatron.core.models.GPTModel!") + + if pretrained_model_name_or_path is None: + self.arch = "GPTModel" + else: + self._hf_config = transformers.AutoConfig.from_pretrained(pretrained_model_name_or_path) + # Update hf_config + self._hf_config.num_hidden_layers = model.config.num_layers + self._hf_config.hidden_size = model.config.hidden_size + self._hf_config.head_dim = model.config.kv_channels + self._hf_config.num_attention_heads = model.config.num_attention_heads + self._hf_config.num_key_value_heads = model.config.num_query_groups + self._hf_config.intermediate_size = model.config.ffn_hidden_size + self.arch = self._hf_config.architectures[0] + + self.all_rules = self._populate_rule_book() + self.rules = self.all_rules[self.arch] + self.model = model + self.dtype = dtype + self._state_dict = OrderedDict() + self._hf_pretrained_model_name = pretrained_model_name_or_path + + def save_pretrained(self, save_directory: Union[str, os.PathLike]): + """Save a unified checkpoint which can be deploied by vLLM and TensorRT-LLM. + + Args: + save_directory: Directory to which to save. Will be created if it doesn't exist. + """ + pp_rank = get_pipeline_model_parallel_rank() + pp_size = get_pipeline_model_parallel_world_size() + + meta_filename = "model-{:05d}-of-{:05d}.json".format(pp_rank + 1, pp_size) + ckpt_filename = "model-{:05d}-of-{:05d}.safetensors".format( + pp_rank + 1, + pp_size, + ) + + # Main export process + state_dict = self.state_dict + quantization = None + kv_cache_quantization = None + total_size = 0 + weight_map = {} + + for key, val in state_dict.items(): + total_size += val.numel() * val.element_size() + weight_map[key] = ckpt_filename + if "weight_scale" in key: + quantization = "FP8" + if "k_scale" in key or "output_scale" in key: + kv_cache_quantization = "FP8" + + if torch.distributed.get_rank() == 0: + if self._hf_pretrained_model_name: + self._hf_config.save_pretrained(save_directory) + try: + generation_config = transformers.GenerationConfig.from_pretrained( + self._hf_pretrained_model_name + ) + generation_config.save_pretrained(save_directory) + except OSError: + pass + try: + tokenizer = transformers.AutoTokenizer.from_pretrained( + self._hf_pretrained_model_name + ) + tokenizer.save_pretrained(save_directory) + except OSError: + pass + else: + os.makedirs(save_directory, exist_ok=True) + + # Barrier to ensure the export_dir has been created. + torch.distributed.barrier() + + # Write multi-file metadata + with open(save_directory + "/" + meta_filename, "w") as f: + json.dump( + {"metadata": {"total_size": total_size}, "weight_map": weight_map}, f, indent=4 + ) + + save_file(self.state_dict, save_directory + "/" + ckpt_filename, metadata={"format": "pt"}) + + # Barrier to ensure all ranks have written the metadata + torch.distributed.barrier() + + if torch.distributed.get_rank() == 0: + hf_quant_config = { + "producer": { + "name": "modelopt", + "version": __version__, + }, + "quantization": { + "quant_algo": quantization, + "kv_cache_quant_algo": kv_cache_quantization, + }, + } + with open(save_directory + "/hf_quant_config.json", "w") as f: + json.dump(hf_quant_config, f, indent=4) + safetensor_index = { + "metadata": {"total_size": 0}, + "weight_map": {}, + } + for pp_rank in range(get_pipeline_model_parallel_world_size()): + meta_filename = "model-{:05d}-of-{:05d}.json".format( + pp_rank + 1, + get_pipeline_model_parallel_world_size(), + ) + with open(save_directory + "/" + meta_filename) as f: + shard = json.load(f) + safetensor_index["metadata"]["total_size"] += shard["metadata"]["total_size"] + safetensor_index["weight_map"].update(shard["weight_map"]) + + with open(save_directory + "/model.safetensors.index.json", "w") as f: + json.dump(safetensor_index, f, indent=4) + + @property + def state_dict(self): + """Return the real quantized state_dict.""" + if len(self._state_dict) == 0: + self._get_state_dict() + return self._state_dict + + def _populate_rule_book(self): + all_rules = {} + + def _custom_mapping_to_lambda(mapping): + method_map = { + "name_remapping": self._name_remapping, + "qkv_slicing": self._qkv_slicing, + "gated_mlp_slicing": self._gated_mlp_slicing, + } + func = method_map[mapping.func_name] + prefix = mapping.target_name_or_prefix + func_kwargs = mapping.func_kwargs + return lambda m, *args: func(m, prefix.format(*args), **func_kwargs) + + for arch, mappings in all_mcore_hf_export_mapping.items(): + all_rules[arch] = {k: _custom_mapping_to_lambda(v) for (k, v) in mappings.items()} + + return all_rules + + def _get_weight_scales(self, quantized_state: Dict[str, Any], qformat: str): + weight_scale = quantized_state.get("weight_scale", None) + weight_scale_2 = quantized_state.get("weight_scale_2", None) + + if weight_scale is not None: + weight_scale = weight_scale.clone().detach() + if qformat == QUANTIZATION_FP8 and weight_scale.numel() == 1: + weight_scale = weight_scale.squeeze() + if weight_scale_2 is not None: + weight_scale_2 = weight_scale_2.clone().detach() + + return weight_scale, weight_scale_2 + + def _name_remapping(self, module, prefix, skip_output_scale=True): + name_to_value, qformat, block_size = get_quantized_state(module, self.dtype) + + weight = name_to_value.pop("weight") + weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) + + if weight_scale is None: + self._state_dict[prefix + "weight"] = weight + else: + self._state_dict[prefix + "weight"] = to_quantized_weight( + weight, + weight_scale, + qformat, + weight_scale_2, + block_size, + ) + + for key, val in name_to_value.items(): + if "output_scale" == key and skip_output_scale: + continue + else: + self._state_dict[prefix + key] = val + + def _gated_mlp_slicing( + self, module, prefix, gate_proj_name="gate_proj", up_proj_name="up_proj" + ): + name_to_value, qformat, block_size = get_quantized_state(module, self.dtype) + + weight = name_to_value.pop("weight") + weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) + + gate_proj_prefix = prefix + gate_proj_name + "." + up_proj_prefix = prefix + up_proj_name + "." + + ffn_hidden_size = module.config.ffn_hidden_size + gate_proj_weight = weight[:ffn_hidden_size, :] + up_proj_weight = weight[ffn_hidden_size:, :] + + if weight_scale is None: + self._state_dict[gate_proj_prefix + "weight"] = gate_proj_weight + self._state_dict[up_proj_prefix + "weight"] = up_proj_weight + else: + if len(weight_scale.shape) == 0: + gate_proj_weight_scale = weight_scale.detach().clone() + up_proj_weight_scale = weight_scale.detach().clone() + else: + gate_proj_weight_scale = weight_scale[:ffn_hidden_size] + up_proj_weight_scale = weight_scale[ffn_hidden_size:] + self._state_dict[gate_proj_prefix + "weight"] = to_quantized_weight( + gate_proj_weight, + gate_proj_weight_scale, + qformat, + weight_scale_2, + block_size, + ) + self._state_dict[up_proj_prefix + "weight"] = to_quantized_weight( + up_proj_weight, + up_proj_weight_scale, + qformat, + weight_scale_2, + block_size, + ) + self._state_dict[gate_proj_prefix + "weight_scale"] = gate_proj_weight_scale + self._state_dict[up_proj_prefix + "weight_scale"] = up_proj_weight_scale + + # weight and weight_scale have been pop out. + for key, val in name_to_value.items(): + gate_proj_key = gate_proj_prefix + key + up_proj_key = up_proj_prefix + key + if "output_scale" == key: + continue + else: + self._state_dict[gate_proj_key] = val.detach().clone() + self._state_dict[up_proj_key] = val.detach().clone() + + def _qkv_slicing( + self, + module, + prefix, + q_proj_name="q_proj", + k_proj_name="k_proj", + v_proj_name="v_proj", + k_scale_name="k_scale", + v_scale_name="v_scale", + ): + name_to_value, qformat, block_size = get_quantized_state(module, self.dtype) + + q_proj_prefix = prefix + q_proj_name + "." + k_proj_prefix = prefix + k_proj_name + "." + v_proj_prefix = prefix + v_proj_name + "." + + config = module.config + hidden_size = config.hidden_size + num_query_groups = config.num_query_groups + head_num = config.num_attention_heads + head_size = hidden_size // head_num + heads_per_group = head_num // num_query_groups + qkv_total_dim = head_num + 2 * num_query_groups + + weight = name_to_value.pop("weight").reshape([qkv_total_dim, head_size, hidden_size]) + weight_scale, weight_scale_2 = self._get_weight_scales(name_to_value, qformat) + + q_slice = torch.cat( + [ + torch.arange((heads_per_group + 2) * i, (heads_per_group + 2) * i + heads_per_group) + for i in range(num_query_groups) + ] + ) + k_slice = torch.arange(heads_per_group, qkv_total_dim, (heads_per_group + 2)) + v_slice = torch.arange(heads_per_group + 1, qkv_total_dim, (heads_per_group + 2)) + ## Example of slices + ## 7b: num_query_groups = head_num = 32, + ## q_slice = [0, 3, 6, 9 , ... 90, 93] + ## k_slice = [1, 4, 7, 10, ... 91, 94] + ## v_slice = [2, 5, 8, 11, ... 92, 95] + ## 70b (with GQA): num_query_groups = 8, head_num = 64 + ## q_slice = [0, 1, .. 6, 7, 10, 11, .. 16, 17, 20, 21, .. 67, 70, ... 76, 77] + ## k_slice = [8, 18, 28, ... 68, 78] + ## v_slice = [9, 19, 29, ... 69, 79] + slices = [q_slice, k_slice, v_slice] + prefixes = [q_proj_prefix, k_proj_prefix, v_proj_prefix] + + proj_weights = [weight[s].reshape(-1, hidden_size) for s in slices] + proj_keys = [p + "weight" for p in prefixes] + + if weight_scale is None: + for key, weight in zip(proj_keys, proj_weights): + self._state_dict[key] = weight + else: + if len(weight_scale.shape) > 0: + # AWQ per-block or per-channel scaling + proj_weight_scales = [weight_scale[s] for s in slices] + else: + # per-tensor scaling + proj_weight_scales = [weight_scale.detach().clone()] * 3 + + for weight, scale, key in zip(proj_weights, proj_weight_scales, proj_keys): + quantized_weight = to_quantized_weight( + weight, + scale, + qformat, + weight_scale_2, + block_size, + ) + self._state_dict[key] = quantized_weight + self._state_dict[key + "_scale"] = scale + + # weight and weight_scale have been pop out. + for key, val in name_to_value.items(): + q_proj_key = q_proj_prefix + key + k_proj_key = k_proj_prefix + key + v_proj_key = v_proj_prefix + key + if "output_scale" == key: + self._state_dict[prefix + k_scale_name] = val.detach().clone() + self._state_dict[prefix + v_scale_name] = val.detach().clone() + else: + self._state_dict[q_proj_key] = val.detach().clone() + self._state_dict[k_proj_key] = val.detach().clone() + self._state_dict[v_proj_key] = val.detach().clone() + + def _get_state_dict(self): + model = self.model + + # Embedding + if hasattr(model, "embedding"): + self.rules["word_embeddings"](model.embedding.word_embeddings) + + # Final layernorm + if hasattr(model.decoder, "final_layernorm") and model.decoder.final_layernorm: + self.rules["final_layernorm"](model.decoder.final_layernorm) + + # Output layer + if hasattr(model, "output_layer"): + self.rules["output_layer"](model.output_layer) + + # Decoder layers + for layer in model.decoder.layers: + layer_id = layer.layer_number - 1 + self.rules["input_layernorm"](layer.input_layernorm, layer_id) + self.rules["linear_qkv"](layer.self_attention.linear_qkv, layer_id) + self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id) + self.rules["pre_mlp_layernorm"](layer.pre_mlp_layernorm, layer_id) + self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id) + self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id) + + if hasattr(model, "medusa_heads"): + for head_id, head in enumerate(model.medusa_heads): + self.rules["medusa_heads.lm_head"](head.lm_head, head_id) + for layer_id, layer in enumerate(head.medusa_layers): + self.rules["medusa_heads.medusa_layers.linear"](layer.linear, head_id, layer_id) + + if hasattr(model, "eagle_module"): + self.rules["eagle_module.fc"](model.eagle_module.fc) + for layer in model.eagle_module.decoder.layers: + layer_id = layer.layer_number - 1 + self.rules["input_layernorm"](layer.input_layernorm, layer_id) + self.rules["linear_qkv"](layer.self_attention.linear_qkv, layer_id) + self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id) + self.rules["pre_mlp_layernorm"](layer.pre_mlp_layernorm, layer_id) + self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id) + self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id) + + +def export_mcore_gpt_to_hf( + model: torch.nn.Module, + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]] = None, + dtype: torch.dtype = torch.float16, + export_dir: Union[Path, str] = tempfile.gettempdir(), +): + """Export Megatron Core GPTModel to unified checkpoint and save to export_dir. + + Args: + model: The Megatron Core GPTModel instance. + pretrained_model_name_or_path: Can be either: the *model id* of a + pretrained model hosted inside a model repo on huggingface.co; or + a *directory* containing model weights saved using + [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + dtype: The weights data type to export the unquantized layers. + export_dir: The target export path. + """ + exporter = GPTModelExporter(model, pretrained_model_name_or_path, dtype) + exporter.save_pretrained(export_dir) + + +class GPTModelImporter: + """Megatron Core GPTModel HuggingFace Importer. + + The Importer is created by `import_mcore_gpt_from_hf` to host attributes + and methods that import a Megatron Core GPTModel from a supported Hugging + Face model. + + Args: + model: The Megatron Core GPTModel instance. + pretrained_model_name_or_path: Can be either: the *model id* of a + pretrained model hosted inside a model repo on huggingface.co; or + a *directory* containing model weights saved using + [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + dtype: The weights data type to export the unquantized layers. + """ + + def __init__( + self, + model: torch.nn.Module, + pretrained_model_name_or_path: str, + dtype=torch.bfloat16, + ): + """Create a GPTModel importer instance.""" + self._hf_config = transformers.AutoConfig.from_pretrained(pretrained_model_name_or_path) + pretrained_model_path = Path(pretrained_model_name_or_path) + if not pretrained_model_path.is_dir(): + pretrained_model_path = tempfile.gettempdir() + "/" + pretrained_model_name_or_path + snapshot_download( + repo_id=pretrained_model_name_or_path, + local_dir=pretrained_model_path, + ) + self.arch = self._hf_config.architectures[0] + self.all_rules = self._populate_rule_book() + self.rules = self.all_rules[self.arch] + self.model = model + self.pretrained_model_path = pretrained_model_path + self.dtype = dtype + + def _populate_rule_book(self): + """The rule book maps each state_dict key to a Callable.""" + all_rules = {} + + def _custom_mapping_to_lambda(mapping): + method_map = { + "name_remapping": self._name_remapping, + "qkv_merging": self._qkv_merging, + "gated_mlp_merging": self._gated_mlp_merging, + } + func = method_map[mapping.func_name] + prefix = mapping.target_name_or_prefix + func_kwargs = mapping.func_kwargs + return lambda m, *args: func(m, prefix.format(*args), **func_kwargs) + + for arch, mappings in all_mcore_hf_import_mapping.items(): + all_rules[arch] = {k: _custom_mapping_to_lambda(v) for (k, v) in mappings.items()} + + return all_rules + + def _get_safetensor(self, key): + """Get a safetensor from the sharded checkpoint.""" + safetensors_file = Path(self.pretrained_model_path + "/model.safetensors") + safetensors_index_file = Path(self.pretrained_model_path + "/model.safetensors.index.json") + if safetensors_file.is_file(): + pass + elif safetensors_index_file.is_file(): + with open(safetensors_index_file) as f: + safetensors_index = json.load(f) + safetensors_file = Path( + self.pretrained_model_path + "/" + safetensors_index["weight_map"][key] + ) + else: + raise ValueError("Only safetensors (single of multi- files) are supported.") + + with safe_open(safetensors_file, framework="pt") as f: + tensor = f.get_tensor(key) + return tensor + + def _name_remapping(self, module, prefix): + state_dict = {} + for key, val in module.state_dict().items(): + if "extra_state" in key: + state_dict[key] = val + else: + tensor = self._get_safetensor(prefix + key) + state_dict[key] = tensor.to(dtype=self.dtype).to(device=val.device) + module.load_state_dict(state_dict) + + def _gated_mlp_merging( + self, module, prefix, gate_proj_name="gate_proj", up_proj_name="up_proj" + ): + gate_proj = self._get_safetensor(prefix + gate_proj_name + ".weight") + up_proj = self._get_safetensor(prefix + up_proj_name + ".weight") + tensor = torch.cat((gate_proj, up_proj), dim=0) + state_dict = {"weight": tensor.to(dtype=self.dtype).to(device=module.weight.device)} + module.load_state_dict(state_dict) + + def _qkv_merging( + self, + module, + prefix, + q_proj_name="q_proj", + k_proj_name="k_proj", + v_proj_name="v_proj", + ): + config = module.config + hidden_size = config.hidden_size + num_query_groups = config.num_query_groups + head_num = config.num_attention_heads + head_size = hidden_size // head_num + heads_per_group = head_num // num_query_groups + qkv_total_dim = head_num + 2 * num_query_groups + q_slice = torch.cat( + [ + torch.arange((heads_per_group + 2) * i, (heads_per_group + 2) * i + heads_per_group) + for i in range(num_query_groups) + ] + ) + k_slice = torch.arange(heads_per_group, qkv_total_dim, (heads_per_group + 2)) + v_slice = torch.arange(heads_per_group + 1, qkv_total_dim, (heads_per_group + 2)) + + q_proj = self._get_safetensor(prefix + q_proj_name + ".weight").reshape( + -1, head_size, hidden_size + ) + k_proj = self._get_safetensor(prefix + k_proj_name + ".weight").reshape( + -1, head_size, hidden_size + ) + v_proj = self._get_safetensor(prefix + v_proj_name + ".weight").reshape( + -1, head_size, hidden_size + ) + tensor = module.weight.detach().clone().reshape([qkv_total_dim, head_size, hidden_size]) + + tensor[q_slice] = q_proj.to(dtype=tensor.dtype).to(device=tensor.device) + tensor[k_slice] = k_proj.to(dtype=tensor.dtype).to(device=tensor.device) + tensor[v_slice] = v_proj.to(dtype=tensor.dtype).to(device=tensor.device) + + module.load_state_dict({"weight": tensor.reshape(-1, hidden_size)}) + + def _import_state_dict(self): + model = self.model + + # Embedding + if hasattr(model, "embedding"): + self.rules["word_embeddings"](model.embedding.word_embeddings) + + # Final layernorm + if hasattr(model.decoder, "final_layernorm") and model.decoder.final_layernorm: + self.rules["final_layernorm"](model.decoder.final_layernorm) + + # Output layer + if hasattr(model, "output_layer"): + self.rules["output_layer"](model.output_layer) + + # Decoder layers + for layer in model.decoder.layers: + layer_id = layer.layer_number - 1 + self.rules["input_layernorm"](layer.input_layernorm, layer_id) + self.rules["linear_qkv"](layer.self_attention.linear_qkv, layer_id) + self.rules["linear_proj"](layer.self_attention.linear_proj, layer_id) + self.rules["pre_mlp_layernorm"](layer.pre_mlp_layernorm, layer_id) + self.rules["linear_fc1"](layer.mlp.linear_fc1, layer_id) + self.rules["linear_fc2"](layer.mlp.linear_fc2, layer_id) + + +def import_mcore_gpt_from_hf( + model: torch.nn.Module, + pretrained_model_path: str, + dtype: torch.dtype = torch.float16, +): + """Import GPTModel state_dict from supported HuggingFace pretrained model path. + + Args: + model: The Megatron Core GPTModel instance. + pretrained_model_path: A path to a *directory* containing model weights saved using + [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + dtype: The weights data type to import. + """ + importer = GPTModelImporter(model, pretrained_model_path, dtype) + importer._import_state_dict() diff --git a/modelopt/torch/nas/__init__.py b/modelopt/torch/nas/__init__.py new file mode 100644 index 0000000000..47838d1835 --- /dev/null +++ b/modelopt/torch/nas/__init__.py @@ -0,0 +1,27 @@ +# 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. + +"""Main module for NAS/Pruning-based model design and optimization.""" + +# isort: off +# Import dynamic modules first to avoid KeyError: 'nn.Conv1d is not registered for a dynamic module!' +from . import modules + +# isort: on + +from . import config, hparams, mode, plugins, utils +from .algorithms import * +from .conversion import * +from .utils import * diff --git a/modelopt/torch/nas/_algorithms.py b/modelopt/torch/nas/_algorithms.py new file mode 100644 index 0000000000..d1590c65c3 --- /dev/null +++ b/modelopt/torch/nas/_algorithms.py @@ -0,0 +1,430 @@ +# 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. + +"""High-level search and model design algorithms to help you optimize your model.""" + +import copy +from typing import Any, Callable, Optional, Union + +import numpy as np +from numpy.typing import NDArray +from scipy.interpolate import LinearNDInterpolator, interp1d +from scipy.spatial import QhullError +from torch import nn + +from modelopt.torch.opt.searcher import ConstraintsDict, Deployment, LimitsTuple +from modelopt.torch.opt.utils import is_configurable +from modelopt.torch.utils import distributed as dist +from modelopt.torch.utils import param_num_from_forward, random, unwrap_model + +from .search_space import SampleFunc +from .utils import ( + _SearchSpaceUnwrapped, + get_subnet_config, + inference_flops, + sample, + sample_and_reset, +) + +ConstraintsRes = dict[str, float] +ConstraintEvalFunc = Callable[[Optional[ConstraintsRes]], float] + + +class ConstraintInterpolator: + """A piece-wise, linear, d-dim interpolator to interpolate some model property like latency. + + The interpolator is based on scipy.interpolate.LinearNDInterpolator, which is based on an + d-dimensional linear interpolation from unstructured d-dimensional data points. The interpolator + uses a Delaunay triangulation as intermediate data representation. + + Note that for a successful triangulation/interpolation, at least d+1 linearly independent points + are needed that form a convex hull around the desired query point. + + Every time a new data point is queried, the interpolator will be dynamically re-generated as + needed. + """ + + def __init__( + self, + model: nn.Module, + points_funcs: dict[str, ConstraintEvalFunc], + value_func: ConstraintEvalFunc, + ): + """Initialize interpolator with points_func and value_func to interpolate.""" + assert len(points_funcs) > 0, "Need at least one point function to interpolate from!" + self.model = model + self._points_funcs = points_funcs + self._value_func = value_func + self.points = np.zeros((0, len(points_funcs))) + self._vals = np.zeros(0) + self._interpolator = None + + # In collect mode, data points are added to the grid even if the interpolator could be used + # in order to improve the interpolator's quality. + self.collect_mode = False + + def _add_point(self, point: Union[list, NDArray]) -> float: + """Add point to interpolation table and return value.""" + # compute new value + val = self._value_func(self.model) + + # add point/va. to our database of points + self.points = np.vstack((self.points, np.asarray(point)[None])) + self._vals = np.hstack((self._vals, val)) + + # dynamically re-generate the interpolator + try: + if len(self._points_funcs) > 1: + self._interpolator = LinearNDInterpolator(self.points, self._vals) + elif len(self._vals) > 1: # can only use interp1d with 2+ points collected + self._interpolator = interp1d(self.points[:, 0], self._vals, bounds_error=False) + else: + self._interpolator = None # not enough points collected yet + except QhullError: + self._interpolator = None # not enough linearly independent points collected yet + + # return new point + return val + + def __call__(self, pre_computed: Optional[dict[str, float]] = None) -> float: + """Interpolate based on the provided model query or compute a new, exact data point.""" + # compute new point but only query as necessary + pre_computed = pre_computed or {} + point = [ + pre_computed[key] if key in pre_computed else func(self.model) + for key, func in self._points_funcs.items() + ] + + # check if we can use the interpolator or need/want to query the value directly + if self._interpolator is None or self.collect_mode: + val = self._add_point(point) + else: + val = self._interpolator(point) + # check if we landed outside the convex hull (val == NaN) and if so add it as new point + if np.isnan(val): + val = self._add_point(point) + + # return value + return np.asarray(val).item() + + +class ConstraintsFunc: + """A Functor class to check if sub-net satisfied all provided constraints. + + We intentionally expose some attributes like `limits` s.t. we can modify it manually. + """ + + _sample_points_dict: dict[tuple[str, ...], dict[str, SampleFunc]] = { + ("flops",): {"min": min, "centroid": random.centroid, "max": max}, + ("flops", "flops_min_depth"): { + "min": min, + "min_regular_centroid_depth": {"*": min, "*.depth": random.centroid}, + "min_regular_max_depth": {"*": min, "*.depth": max}, + "centroid_regular_min_depth": {"*": random.centroid, "*.depth": min}, + "centroid": random.centroid, + "centroid_regular_max_depth": {"*": random.centroid, "*.depth": max}, + "max_regular_min_depth": {"*": max, "*.depth": min}, + "max_regular_centroid_depth": {"*": max, "*.depth": random.centroid}, + "max": max, + }, + } + + def __init__( + self, + model: nn.Module, + constraints: ConstraintsDict, + dummy_input: Union[Any, tuple], + deployment: Optional[Deployment] = None, + fast_eval: bool = True, + dp_group=None, + ): + self.model = model + self.dummy_input = dummy_input + self.deployment = deployment + self._fast_eval = fast_eval + + # Getting data parallel group for + self.dp_group = dp_group + + # initialize latency interpolator + keys_for_interpolation = ("flops",) + if ConstraintsFunc.is_configurable(self.model, "depth"): + keys_for_interpolation += ("flops_min_depth",) + self._latency_interpolator = ConstraintInterpolator( + self.model, + points_funcs={k: self.constraint_eval_funcs[k] for k in keys_for_interpolation}, + value_func=self._get_true_latency, + ) + # set fast/regular mode for latency interpolator + self._latency_interpolator.collect_mode = not self.fast_eval + + # set limit at the end with setter to use sanity checks on constraints + self._limits = {} + self.limits = constraints + + @staticmethod + def is_configurable(model: nn.Module, hparam_suffix: Optional[str] = None): + return hparam_suffix is None or any( + k.rpartition(".")[-1] == hparam_suffix for k in get_subnet_config(model, True).keys() + ) + + def get_sample_points(self) -> dict[str, SampleFunc]: + """Get the sample points required for latency interpolator.""" + return self._sample_points_dict[tuple(self._latency_interpolator._points_funcs)] + + @property + def fast_eval(self) -> bool: + """Return boolean indicator whether we use faster interpolation for latency.""" + return self._fast_eval + + @fast_eval.setter + def fast_eval(self, value: bool): + """Set bool whether we should use faster interpolation for latency.""" + self._fast_eval = value + self._latency_interpolator.collect_mode = not value + + @property + def limits(self) -> dict[str, LimitsTuple]: + """Get limits for all constraints.""" + return self._limits + + @limits.setter + def limits(self, constraints: ConstraintsDict) -> None: + # check that no unexpected keys are found in provided constraints dict. + if not constraints.keys() <= self.constraint_eval_funcs.keys(): + raise KeyError( + f"Found unexpected keys:\n{constraints.keys() - self.constraint_eval_funcs.keys()}", + f"\nCurrently supported constraint keys:\n{self.constraint_eval_funcs.keys()}", + ) + + # check that deployment is provided when latency is used as a constraint. + if "latency" in constraints and self.deployment is None: + raise ValueError( + "Deployment config must be provided when latency is used as a constraint." + ) + + # making sure we add both "flops" and "params" by default as fake constraints + # since it's cheap and can provide useful insights during constraint evaluation. + # Moreover, we need it for latency interpolation anyway. + constraints_complete: ConstraintsDict = { + k: None for k in ("flops", "params") if k in self.constraint_eval_funcs + } + constraints_complete.update(**constraints) + + if "latency" in constraints: + fake_constraints = {k: None for k in self._latency_interpolator._points_funcs.keys()} + constraints_complete.update(fake_constraints) + + self._limits = self.get_expanded_limits(constraints_complete) + + @property + def effective_limits(self) -> dict[str, LimitsTuple]: + return {k: v for k, v in self.limits.items() if v != self.trivial_limit} + + @property + def trivial_limit(self) -> LimitsTuple: + return (0, float("inf")) + + def get_expanded_limits(self, constraints: ConstraintsDict) -> dict[str, LimitsTuple]: + """Extract and standardize limits from constraint dictionary.""" + # standardize limits to LimitsTuple: (lower_bound, upper_bound) + expanded_limits = {} + for k, lim in copy.deepcopy(constraints).items(): + if lim is None: + expanded_limits[k] = self.trivial_limit + continue + + if isinstance(lim, str): + assert lim.endswith("%") + fraction_value = float(lim[:-1]) / 100 + + # Use _SearchSpaceUnwrapped to avoid reseting the calibration statistics + ss = _SearchSpaceUnwrapped(self.model) + config = ss.config() + ss.sample(random.original) + orig_value = self.constraint_eval_funcs[k]() + ss.select(config) + + lim = orig_value * fraction_value + + expanded_limits[k] = (0, lim) # lower bound is always 0 by default + + return expanded_limits + + def set_rel_lower_bounds(self, rel_lower: float = 0.85) -> None: + """Set relative lower bounds for all constraints.""" + for k, (lb, ub) in self._limits.items(): + assert lb == 0, "Lower bound has already been set!" + if self._limits[k] != self.trivial_limit: + self._limits[k] = (rel_lower * ub, ub) + + def _get_params(self, _: Optional[ConstraintsRes] = None) -> float: + """Get number of model parameters from forward pass.""" + return param_num_from_forward(self.model, args=self.dummy_input, unit=1.0) + + def _get_flops(self, _: Optional[ConstraintsRes] = None) -> float: + """Get inference FLOPs.""" + return float(inference_flops(self.model, dummy_input=self.dummy_input, unit=1.0)) + + def _get_flops_min_depth(self, _: Optional[ConstraintsRes] = None) -> float: + """Get inference FLOPs with depth set to minimum.""" + with sample_and_reset(self.model, {"*.depth": min}): + flops = float(inference_flops(self.model, dummy_input=self.dummy_input, unit=1.0)) + return flops + + def _get_true_latency(self, _: Optional[ConstraintsRes] = None) -> float: + """Get true inference latency.""" + from modelopt.torch._deploy import get_latency + + assert self.deployment is not None, "Deployment config must be provided for latency!" + return get_latency(self.model, self.dummy_input, self.deployment) + + def _get_latency(self, precomputed: Optional[ConstraintsRes] = None) -> float: + """Get inference latency from interpolator.""" + return self._latency_interpolator(precomputed) + + @property + def constraint_eval_funcs(self) -> dict[str, ConstraintEvalFunc]: + """Get constraint eval fns.""" + return { + "params": self._get_params, + "flops": self._get_flops, + "flops_min_depth": self._get_flops_min_depth, + "latency": self._get_latency, + } + + def __call__(self) -> tuple[bool, ConstraintsRes]: + return self.evaluate_constraints() + + def evaluate_constraints(self) -> tuple[bool, ConstraintsRes]: + """Check if sub-net satisfied all provided constraints. + + Returns: A tuple (is_satisfied, vals) where + is_satisfied is an indicator (bool) whether all constraints are satisfied, + vals is a dictionary containing all evaluated constraints up to 1st unsatisfied constraint. + + The sub-net is profiled after exporting it in eval mode without BN. + """ + + def check_constraint(val: float, limits: LimitsTuple) -> bool: + """Check if val falls within limits. `None` indicates open-ended intervals.""" + return val >= limits[0] and val <= limits[1] + + is_satisfied = True + vals: ConstraintsRes = {} + + if dist.is_master(group=self.dp_group): + # construct sub-net without DDP and BN for more accurate profiling + model = unwrap_model(self.model) + + # put model into eval mode + is_training = model.training + model.eval() + + # check constraint (with early stopping if not satisfied) + # follow the order of the constraint_eval_funcs + for criterion, score_func in self.constraint_eval_funcs.items(): + if criterion not in self.limits: + continue + val = score_func(vals) + vals[criterion] = val + is_satisfied &= check_constraint(val, self.limits[criterion]) + + # put back into original mode + model.train(is_training) + + is_satisfied, vals = dist.broadcast([is_satisfied, vals]) + return is_satisfied, vals + + +class StatisticsEstimator: + """A class to randomly sample subnets & estimate statistics for the provided constraints. + + The class can: + Estimate median and standard deviation of the constraints + Plot the distribution of constraints + + Args: + constraints_func: An object of class ``ConstraintsFunc`` used to define and evaluate the + constraints. + num_of_samples: An integer that specifies the number of random samples of model + used to estimate the constraint statsitics. + """ + + def __init__( + self, + constraints_func: ConstraintsFunc, + num_of_samples: int = 100, + ) -> None: + self.constraints_func = constraints_func + self.num_of_samples = num_of_samples + self._sample_dict = {} + self.reset() + + def reset(self): + """Resets statistics estimator history""" + self._sample_dict = {k: [] for k in self.constraints_func.limits.keys()} + + @property + def num_of_collected_samples(self): + return len(next(iter(self._sample_dict.values()))) + + def _sample_random_subnets(self, model) -> dict: + with random._deterministic_seed(): + max_num_samples = self.num_of_samples if is_configurable(model) else 1 + for _ in range(self.num_of_collected_samples, max_num_samples): + sample(model) + info = self.constraints_func()[1] + for constraint in self.constraints_func.limits.keys(): + self._sample_dict[constraint].append(info[constraint]) + + return self._sample_dict + + def median(self, model: nn.Module) -> dict[str, float]: + """Estimates median constraint values for the model.""" + median_dict = {} + sample_dict = self._sample_random_subnets(model) + for constraint, constraint_samples in sample_dict.items(): + median_dict[constraint] = float(np.median(constraint_samples)) + return median_dict + + def std(self, model: nn.Module) -> dict[str, float]: + """Estimates std constraint values for the model.""" + std_dict = {} + sample_dict = self._sample_random_subnets(model) + for constraint, constraint_samples in sample_dict.items(): + std_dict[constraint] = float(np.std(constraint_samples)) + return std_dict + + def median_subnet(self, model: nn.Module) -> dict: + """Estimates a median/representative subnet for the model. + + Returns model nas.config dict + """ + raise NotImplementedError + + +# Mapping module type to a specific constraint function. Plugins may change this mapping. +MODULE_TYPE_TO_CONSTRAINTS_FUNC = {} + + +def get_constraints_func(model, *args, **kwargs) -> ConstraintsFunc: + """Getting the standard constraints function or a special one from the plugins.""" + # Initialize special constraints functor if MODULE_TYPE_TO_CONSTRAINTS_FUNC is not empty. + for module_type, func in MODULE_TYPE_TO_CONSTRAINTS_FUNC.items(): + if isinstance(model, module_type): + return func(model, *args, **kwargs) + # Initialize the standard constraints functor + return ConstraintsFunc(model, *args, **kwargs) diff --git a/modelopt/torch/nas/_patch.py b/modelopt/torch/nas/_patch.py new file mode 100644 index 0000000000..1b3c94c1e5 --- /dev/null +++ b/modelopt/torch/nas/_patch.py @@ -0,0 +1,262 @@ +# 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 types +from abc import ABC, abstractmethod +from contextlib import contextmanager +from typing import Any, Callable, Iterator, Optional + +import torch.nn as nn + +from modelopt.torch.opt.searcher import ForwardLoop + +from .utils import is_modelopt_patches_enabled + +PatchData = dict[str, Any] +PatchMethods = dict[str, Callable] + + +class PatchManager(ABC): + """A standard interface to handle the monkey patching of a model. + + The class provides two main interfaces, ``patch`` and ``unpatch``, which can be used to properly + overwrite methods and add new attributes to the model. They will all be removed upon calling + ``unpatch()``. + """ + + _patch_data_key = "_modelopt_patch_data" + _patch_cls_key = "_modelopt_patch_cls" + + def __init__(self, model: nn.Module) -> None: + self._model = model + + # find submodule with patch (if it exists) + self._modelopt_name, self._modelopt_module = self._get_named_patched_module(model) + + @classmethod + def _get_named_patched_module( + cls, model: nn.Module + ) -> tuple[Optional[str], Optional[nn.Module]]: + """Return the name of the patched module and the patched module itself.""" + for name, module in model.named_modules(): + if hasattr(module, cls._patch_data_key): + return name, module + return None, None + + @classmethod + def is_patched(cls, model: nn.Module) -> bool: + """Return whether the model is patched.""" + return cls._get_named_patched_module(model)[1] is not None + + @abstractmethod + def _get_default_patch_data(self) -> PatchData: + """Return the default patch data of the model.""" + + @classmethod + def get_manager(cls, model: nn.Module) -> "PatchManager": + """Return the patch manager of the model using the "correct" class.""" + _, modelopt_module = cls._get_named_patched_module(model) + assert modelopt_module is not None, "Model is not patched!" + return getattr(modelopt_module, cls._patch_cls_key)(model) + + @property + def patch_data(self) -> PatchData: + """Return the patch data of the model.""" + assert self._modelopt_module is not None, "Model is not patched!" + return self.patch_data_or_empty + + @property + def patch_data_or_empty(self) -> PatchData: + """Return the patch data of the model or an empty dictionary.""" + if self._modelopt_module is None: + return {} + return getattr(self._modelopt_module, self._patch_data_key) + + def patch(self) -> None: + """Patch model in-place to be compatible with subsequent Model Optimizer tasks.""" + # sanity check that there is no patch data in another submodule + assert ( + self._modelopt_name is None + ), f"Model Optimizer patch detected in {self._modelopt_name}." + + # initialize patch data + setattr(self._model, self._patch_data_key, self._get_default_patch_data()) + setattr(self._model, self._patch_cls_key, type(self)) + + # monkey-patch methods + self._set_patched_methods() + + # correctly set _modelopt_name, _modelopt_module + self._modelopt_name, self._modelopt_module = "", self._model + + def unpatch(self) -> None: + """Remove and delete patching from the model. + + For example: + + .. code-block:: python + + PatchManager(model).unpatch() + model.forward(x) # no patched (auto-) operations will be executed anymore. + """ + # sanity check that there is patch data + assert ( + self._modelopt_name is not None + ), f"No Model Optimizer patch detected in {self._model}." + + # remove patch-related artifacts + delattr(self._modelopt_module, self._patch_data_key) + delattr(self._modelopt_module, self._patch_cls_key) + self._unset_patched_methods() + self._modelopt_name, self._modelopt_module = None, None + + def reset_before_sample(self) -> None: + """Call reset hook before sample-related operations (sample & select).""" + self._hook_pre_sample() + + def call_post_eval(self, forward_loop: Optional[ForwardLoop] = None) -> None: + """Call post-eval hook explicitly. + + Args: + forward_loop: A ``Callable`` that takes a model as input and runs a forward loop on it. + """ + self._hook_post_eval(forward_loop) + + def _hook_pre_sample(self) -> None: + """Optional hook to be called before sample-related operations (sample & select).""" + pass + + def _hook_post_eval(self, forward_loop: Optional[ForwardLoop] = None) -> None: + """Optional hook that is called after eval() (or train(False)) to calibrate the model.""" + pass + + def _hook_pre_forward(self, *args, **kwargs) -> None: + """Optional hook that is called after the original forward function is called.""" + pass + + def _hook_post__replicate_for_data_parallel(self) -> None: + """Optional hook to be called after _replicate_for_data_parallel.""" + self._set_patched_methods() + + @staticmethod + def hooked_forward(mod: nn.Module, *args, **kwargs): + """The forward method with hooks.""" + if is_modelopt_patches_enabled(): + PatchManager.get_manager(mod)._hook_pre_forward(*args, **kwargs) + return getattr(mod, "_modelopt_unhooked_forward")( + mod, *args, **kwargs + ) # call original forward + + @staticmethod + def hooked_train(mod: nn.Module, mode: bool = True) -> nn.Module: + """Sets the model into train or eval mode according to flag.""" + ret = getattr(mod, "_modelopt_unhooked_train")(mod, mode) # call original train + + # call model calibration directly here if we go into eval mode + # Note that train might have already been called from prep_for_eval in which case we + # don't want to call model calibration again!!! + with _modelopt_eval_recursion_guard(mod) as guard: + if not guard and not mode and is_modelopt_patches_enabled(): + PatchManager.get_manager(mod)._hook_post_eval() + + return ret + + @staticmethod + def hooked__replicate_for_data_parallel(mod: nn.Module) -> nn.Module: + mod = getattr(mod, "_modelopt_unhooked__replicate_for_data_parallel")(mod) + if is_modelopt_patches_enabled(): + PatchManager.get_manager(mod)._hook_post__replicate_for_data_parallel() + + def _get_methods_for_patching(self) -> PatchMethods: + """Returns a lookup for methods to be patched together with the getters for the patch.""" + return { + "forward": self.hooked_forward, + "train": self.hooked_train, + "_replicate_for_data_parallel": self.hooked__replicate_for_data_parallel, + } + + def _set_patched_methods(self): + for name, func_hooked in self._get_methods_for_patching().items(): + func_unhooked = getattr(self._model, name).__func__ + if func_unhooked is func_hooked: + continue + setattr(self._model, f"_modelopt_unhooked_{name}", func_unhooked) + setattr(self._model, name, types.MethodType(func_hooked, self._model)) + + def _unset_patched_methods(self): + for name, func_hooked in self._get_methods_for_patching().items(): + func_current = getattr(self._model, name).__func__ + if func_current is not func_hooked: + continue + func_unhooked = getattr(self._model, f"_modelopt_unhooked_{name}") + delattr(self._model, f"_modelopt_unhooked_{name}") + if func_unhooked is getattr(type(self._model), name): + delattr(self._model, name) + else: + setattr(self._model, name, types.MethodType(func_unhooked, self._model)) + + +@contextmanager +def _modelopt_eval_recursion_guard(model: nn.Module) -> Iterator[bool]: + """Context manager to guard against infinite recursion when calling eval/prep_for_eval. + + Args: + model: The model to guard. + + For example: + + .. code-block:: python + + with _modelopt_eval_recursion_guard(model): + model.eval() + """ + guard_key = "_modelopt_eval_recursion_guard" + has_key = hasattr(model, guard_key) + original_guard: bool = getattr(model, guard_key, False) + setattr(model, guard_key, True) + try: + yield original_guard + finally: + # depending on whether key already existed we either delete it or reset it to original val + if has_key: + setattr(model, guard_key, original_guard) + else: + delattr(model, guard_key) + + +def prep_for_eval(model: nn.Module, forward_loop: Optional[ForwardLoop] = None) -> nn.Module: + """Calibrate model for evaluation and enable eval(). + + Args: + model: A nn.Module that might be dynamic. + forward_loop: A ``Callable`` that takes a model as input and runs a pre-defined forward + loop on it using real data. + + .. note:: + + This function should only be explicitly called once after conversion when the model is + immediately used for evaluation without prior training! + """ + # regular eval mode with recursion guard + # this is to ensure that model calibration won't be called again from within eval() + with _modelopt_eval_recursion_guard(model): + ret = model.eval() + + # call post-eval hook explicitly + if PatchManager.is_patched(model): + PatchManager.get_manager(model).call_post_eval(forward_loop) + + # keep return value consistent with native eval() + return ret diff --git a/modelopt/torch/nas/algorithms.py b/modelopt/torch/nas/algorithms.py new file mode 100644 index 0000000000..da4984d2ab --- /dev/null +++ b/modelopt/torch/nas/algorithms.py @@ -0,0 +1,452 @@ +# 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. + +"""High-level search and model design algorithms to help you optimize your model.""" + +import copy +from collections import OrderedDict +from typing import Any, Optional, Union + +import torch +from rich.console import Console +from rich.table import Table +from torch import nn + +from modelopt.torch.opt.conversion import get_mode +from modelopt.torch.opt.searcher import ConstraintsDict, Deployment, SearchConfig +from modelopt.torch.opt.utils import is_configurable, search_space_size +from modelopt.torch.utils import clear_cuda_cache, num2hrb + +from ._algorithms import ConstraintsFunc, ConstraintsRes, StatisticsEstimator, get_constraints_func +from .conversion import export +from .search_space import SampleFunc +from .utils import ( + get_subnet_config, + no_modelopt_patches, + print_search_space_summary, + sample_and_reset, +) + +__all__ = ["profile", "search"] + + +def search( + model: nn.Module, + constraints: ConstraintsDict, + dummy_input: Union[Any, tuple], + config: Optional[SearchConfig] = None, +) -> tuple[nn.Module, dict[str, Any]]: + """Search a given prunable model for the best sub-net and return the search model. + + The best sub-net maximizes the score given by ``score_func`` while satisfying the + ``constraints``. + + Args: + model: The converted model to be searched. + constraints: The dictionary from constraint name to upper bound the searched model has to satisfy. + Currently, we support ``flops`` and ``params`` as constraints. + The constraints dictionary generally takes the following form: + + .. code-block:: python + + constraints = {"params": 5.0e6, "flops": 4.5e8} + + We recommend to simply provide the most relevant constraint, e.g., flops: + + .. code-block:: python + + constraints = {"flops": 4.5e8} + + You can also provide a percentage value instead of absolute value, e.g., + + .. code-block:: python + + # search for a model with <= 60% of the original model flops + constraints = {"flops": "60%"} + + dummy_input: Arguments of ``model.forward()``. This is used for exporting and calculating + inference-based metrics, such as latency/FLOPs. The format of ``dummy_inputs`` follows + the convention of the ``args`` argument in + `torch.onnx.export `_. + Specifically, ``dummy_input`` can be: + + #. a single argument (``type(dummy_input) != tuple``) corresponding to + + .. code-block:: python + + model.forward(dummy_input) + + #. a tuple of arguments corresponding to + + .. code-block:: python + + model.forward(*dummy_input) + + #. a tuple of arguments such that ``type(dummy_input[-1]) == dict`` corresponding to + + .. code-block:: python + + model.forward(*dummy_input[:-1], **dummy_input[-1]) + + .. warning:: + + In this case the model's ``forward()`` method **cannot** contain keyword-only + arguments (e.g. ``forward(..., *, kw_only_args)``) or variable keyword arguments + (e.g. ``forward(..., **kwargs)``) since these cannot be sorted into positional + arguments. + + .. note:: + + In order to pass a dict as last non-keyword argument, you need to use a tuple as + ``dummy_input`` and add an *empty* dict as the last element, e.g., + + .. code-block:: python + + dummy_input = (x, {"y": y, "z": z}, {}) + + The empty dict at the end will then be interpreted as the keyword args. + + See `torch.onnx.export `_ + for more info. + + Note that if you provide a ``{arg_name}`` with batch size ``b``, the results will be + computed based on batch size ``b``. + config: Additional optional arguments to configure the search. Currently, we support: + + * ``checkpoint``: Path to save/restore checkpoint with dictionary containing intermediate + search state. If provided, the intermediate search state will be automatically + restored before search (if exists) and stored/saved during search. + * ``verbose``: Whether to print detailed search space profiling and search stats during search. + * ``forward_loop``: A ``Callable`` that takes a model as input and runs a forward loop + on it. It is recommended to choose the data loader used inside the forward loop + carefully to reduce the runtime. Cannot be provided at the same time as + ``data_loader`` and ``collect_func``. + * ``data_loader``: An iterator yielding batches of data for calibrating the + normalization layers in the model or compute gradient scores. It is recommended to use + the same data loader as for training but with significantly fewer iterations. Cannot + be provided at the same time as ``forward_loop``. + * ``collect_func``: A ``Callable`` that takes a batch of data from the data loader as + input and returns the input to ``model.forward()`` as described in + :meth:`run_forward_loop `. Cannot + be provided at the same time as ``forward_loop``. + * ``max_iter_data_loader``: Maximum number of iterations to run the data loader. + * ``score_func``: A callable taking the model as input and returning a single accuracy/score + metric (float). This metric will be maximized during search. + + .. note:: + + The ``score_func`` is required for ``autonas`` and ``fastnas`` modes. It will be + evaluated on models in eval mode (``model.eval()``). + * ``loss_func``: A ``Callable`` which takes the model output (i.e output of ``model.forward()``) + and the batch of data as its inputs and returns a scalar loss. + This is a required argument if the model is converted via ``gradnas`` mode. + + It should be possible to run a backward pass on the loss value returned by this method. + + ``collect_func`` will be used to gather the inputs to ``model.forward()`` + from a batch of data yielded by``data_loader``. + + ``loss_func`` should support the following usage: + + .. code-block:: python + + for i, batch in enumerate(data_loader): + if i >= max_iter_data_loader: + break + + # Assuming collect_func returns a tuple of arguments + output = model(*collect_func(batch)) + + loss = loss_func(output, batch) + loss.backward() + + .. note:: + + Additional configuration options may be added by individual algorithms. Please + refer to the documentation of the individual algorithms for more information. + + Returns: A tuple (subnet, state_dict) where + subnet is the searched subnet (nn.Module), which can be used for subsequent tasks like + fine-tuning, state_dict contains the history and detailed stats of the search procedure. + + .. note:: + + The given model is modified (exported) in-place to match the best subnet found by the + search algorithm. The returned subnet is thus a reference to the same model instance as the + input model. + """ + # retrieve searcher class + mode = get_mode(model) + assert mode is not None, "Model is not converted to a search space. Please convert first." + searcher_cls = mode.search_algorithm + + # run search, select best, and get config, metric, constraint of search result + searcher = searcher_cls() + + searcher.search(model, constraints, dummy_input, config) + + # export model in-place + model = export(model) + + clear_cuda_cache() + + # return the optimized & exported model and searcher's state_dict (detailed stats) + return model, searcher.state_dict() + + +@torch.no_grad() +@no_modelopt_patches() +def profile( + model: nn.Module, + dummy_input: Optional[Union[Any, tuple]] = None, + constraints: Optional[Union[ConstraintsDict, ConstraintsFunc]] = None, + deployment: Optional[Deployment] = None, + strict: bool = False, + verbose: bool = True, + use_centroid: bool = False, +) -> tuple[bool, dict[str, dict]]: + """Profile statistics of the search space of a converted model or a regular model. + + Args: + model: The model to be profiled. Can be converted or not. + dummy_input: Arguments of ``model.forward()``. This is used for exporting and calculating + inference-based metrics, such as latency/FLOPs. The format of ``dummy_inputs`` follows + the convention of the ``args`` argument in + `torch.onnx.export `_. + Specifically, ``dummy_input`` can be: + + #. a single argument (``type(dummy_input) != tuple``) corresponding to + + .. code-block:: python + + model.forward(dummy_input) + + #. a tuple of arguments corresponding to + + .. code-block:: python + + model.forward(*dummy_input) + + #. a tuple of arguments such that ``type(dummy_input[-1]) == dict`` corresponding to + + .. code-block:: python + + model.forward(*dummy_input[:-1], **dummy_input[-1]) + + .. warning:: + + In this case the model's ``forward()`` method **cannot** contain keyword-only + arguments (e.g. ``forward(..., *, kw_only_args)``) or variable keyword arguments + (e.g. ``forward(..., **kwargs)``) since these cannot be sorted into positional + arguments. + + .. note:: + + In order to pass a dict as last non-keyword argument, you need to use a tuple as + ``dummy_input`` and add an *empty* dict as the last element, e.g., + + .. code-block:: python + + dummy_input = (x, {"y": y, "z": z}, {}) + + The empty dict at the end will then be interpreted as the keyword args. + + See `torch.onnx.export `_ + for more info. + + Note that if you provide a ``{arg_name}`` with batch size ``b``, the results will be + computed based on batch size ``b``. + constraints: The dictionary from constraint name to upper bound the searched model has to satisfy. + Currently, we support ``flops`` and ``params`` as constraints. + The constraints dictionary generally takes the following form: + + .. code-block:: python + + constraints = {"params": 5.0e6, "flops": 4.5e8} + + .. note:: + + We recommend to simply provide the most relevant constraint, e.g., flops: + + .. code-block:: python + + constraints = {"flops": 4.5e8} + + Note that you can also provide a percentage value instead of absolute value, e.g., + + .. code-block:: python + + # search for a model with <= 60% of the original model flops + constraints = {"flops": "60%"} + + strict: Raise an error if constraints are not satisfiable. + verbose: Print detailed profiling results. + use_centroid: By default, profile reports median of the evaluation results from + randomly sampled subnets (instead of the evaluation result from deterministic + centroid subnet). Set use_centroid to True to use the deterministic centroid for + profiling. + + Returns: A tuple (is_all_sat, stats) where + is_all_sat is a bool indicating whether all constraints can be satisfied. + stats is a dictionary containing statistics for the search space if the model is converted, + e.g., the FLOPs and params for the min, centroid, max subnets and their max/min ratios, + size of the search space, number of configurable hparams. + """ + # pre-fetch some low-level properties properties of the search space + is_configurable_ = is_configurable(model) + size_of_search_space = search_space_size(model) + num_configurable = len(get_subnet_config(model, configurable=True)) + + # custom print function that tracks print statements + detailed_msg = "" + + def _print(msg: str, verb: bool = verbose): + nonlocal detailed_msg + if detailed_msg: + detailed_msg += "\n" + detailed_msg += msg + if verb: + print(msg) + + # construct constraints checker function from constraints dict + if isinstance(constraints, ConstraintsFunc): + assert_msg = "No deployment or dummy_input allowed for constraints as ConstraintFunc obj!" + assert deployment is None and dummy_input is None, assert_msg + assert constraints.model is model, "Model of constraints does not match the given model!" + constraints_func = constraints + else: + assert dummy_input is not None, "Please provide a dummy_input for the constraints." + constraints = copy.deepcopy(constraints) or {} + if deployment is not None and "latency" not in constraints: + constraints["latency"] = None + constraints_func = get_constraints_func(model, constraints, dummy_input, deployment, False) + + # retrieve sample points and print message + options: dict[str, SampleFunc] + if is_configurable_: + options = constraints_func.get_sample_points() + msg = f"Profiling the following subnets from the given model: {tuple(options)}." + else: + msg = "Profiling full net, since the model is not a converted model." + options = OrderedDict([("full", max)]) + _print(msg="\n" + msg + "\n" + "-" * 80) + + # extract limits + limits = constraints_func.limits + + # sample through some distinct networks and collect some stats + # note that during that process we definitely don't want fast constraint evaluation! + fast_eval = constraints_func.fast_eval + constraints_func.fast_eval = False + info: dict[str, ConstraintsRes] = {} + for subnet_key, sample_method in options.items(): + # avoid changing the active subnet before and after nas.profile + with sample_and_reset(model, sample_method): + info[subnet_key] = constraints_func()[1] + + # Use statistical median values. + # it does not have a pre-defined sampling method. + # To further speed up, we can reduce the number of samples. + if "centroid" in options and use_centroid is False: + constraints_func.fast_eval = True + statistics_estimator = StatisticsEstimator(constraints_func, num_of_samples=30) + with sample_and_reset(model): + info["centroid"] = statistics_estimator.median(model) + + constraints_func.fast_eval = fast_eval + + # for each sub net and constraint report the constraint evaluation and where it falls within the + # limits + is_all_sat = True + + stats = info.copy() + if is_configurable_: + stats["max/min ratio"] = {} + + # options that we also print + options_keys_for_print = [k for k in options if k in {"min", "centroid", "max", "full"}] + + console = Console() + with console.capture() as capture: + table_profile = Table(title="\nProfiling Results") + table_profile.add_column("Constraint", width=12) + + table_constraints = Table(title="\nConstraints Evaluation") + table_constraints.add_column("Constraint", width=12) + table_constraints.add_column("Upper Bound", width=12) + table_constraints.add_column("Satisfiable Upper Bound", width=12) + + for subnet in options_keys_for_print: + table_profile.add_column(subnet, width=12) + + if is_configurable_: + table_profile.add_column("max/min ratio") + + for constraint, (lb, ub) in limits.items(): + print_info = [f"{constraint:8}"] + is_sat_lower = False + is_sat_upper = False + # check constraint satisfaction + for subnet in options_keys_for_print: + subnet_info = info[subnet] + val = subnet_info[constraint] + is_sat_lower |= val >= lb + is_sat_upper |= val <= ub + print_info.append(f"{num2hrb(val)}") + + if is_configurable_: + # add max/min ratio + ratio_val = info["max"][constraint] / info["min"][constraint] + print_info.append(f"{ratio_val:4.2f}") + stats["max/min ratio"].update({constraint: ratio_val}) + # table_profile.add_row() + table_profile.add_row(*print_info) + + is_all_sat &= is_sat_lower and is_sat_upper + # print update + if constraint in constraints_func.effective_limits: + print_info = [f"{constraint:8}"] + print_info.extend([f"{num2hrb(ub)}", f"{str(is_sat_upper):5}"]) + table_constraints.add_row(*print_info) + + console.print(table_profile) + if len(constraints_func.effective_limits) > 0: + console.print(table_constraints) + + print_info = str(capture.get()) + _print(print_info) # type: ignore[arg-type] + + if is_configurable_: + if verbose: + print_search_space_summary(model) + # add seach space info + _print(f"Number of configurable hparams: {num_configurable}") + stats["number of configurable hparams"] = num_configurable + _print(f"Total size of the search space: {size_of_search_space:8.2e}") + stats["search space size"] = size_of_search_space + + # check if all constraints are satisfiable and raise error if not and strict + if is_all_sat: + if constraints_func.effective_limits: + _print("Note: all constraints can be satisfied within the search space!\n") + elif is_configurable_: + warn_msg = "" if verbose else f"\n{detailed_msg}\n" + warn_msg += "NOT all constraints can be satisfied within the search space, see above!" + if strict: + raise ValueError(warn_msg) + else: + _print(f"Warning: {warn_msg}\n", verb=True) + + return is_all_sat, stats diff --git a/modelopt/torch/nas/autonas.py b/modelopt/torch/nas/autonas.py new file mode 100644 index 0000000000..29d89297a4 --- /dev/null +++ b/modelopt/torch/nas/autonas.py @@ -0,0 +1,637 @@ +# 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. + +"""Entrypoints for AutoNAS mode.""" + +import copy +import hashlib +import json +import warnings +from abc import ABC, abstractmethod +from collections import defaultdict, deque +from functools import partial +from typing import Any, Optional + +import torch +import torch.nn as nn +import tqdm +from torch.nn.modules.batchnorm import _BatchNorm + +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.conversion import ApplyModeError, ModelLikeModule +from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict +from modelopt.torch.opt.searcher import BaseSearcher, ForwardLoop, SearchConfig, SearchStateDict +from modelopt.torch.opt.utils import is_configurable +from modelopt.torch.utils import ( + compare_dict, + get_model_attributes, + is_channels_last, + num2hrb, + random, + run_forward_loop, + stats, + torch_detach, + torch_to, + unwrap_model, +) + +from ._algorithms import ConstraintsFunc, get_constraints_func +from ._patch import PatchData, PatchManager, _modelopt_eval_recursion_guard, prep_for_eval +from .config import ExportConfig +from .search_space import SearchSpace, generate_search_space +from .utils import MODELOPT_BN_CALIB_ITERS, MODELOPT_QUEUE_MAXLEN, get_subnet_config, sample, select + +__all__ = [ + "AutoNASPatchManager", + "EvolveSearcher", + "IterativeSearcher", + "RandomSearcher", + "convert_autonas_searchspace", + "convert_searchspace", + "export_searchspace", + "restore_autonas_searchspace", + "restore_export", + "restore_searchspace", + "update_autonas_metadata", +] + + +class AutoNASPatchManager(PatchManager): + """A class to handle the monkey patching of the model for automode.""" + + def _get_default_patch_data(self) -> PatchData: + """Return the default patch data of the model.""" + return { + "queue": deque(maxlen=MODELOPT_QUEUE_MAXLEN), # queue for reset_bn data batches cache + "fill": False, # force push to queue if set to True + "count": 0, # internal count for number of forward passes + "train_count_since_sample": MODELOPT_BN_CALIB_ITERS, # step count since last sample + "autosampled": False, # indicator whether current sub-net was auto-sampled + } + + def _hook_post_eval(self, forward_loop: Optional[ForwardLoop] = None) -> None: + """Optional hook that is called after eval() (or train(False)) to calibrate the model.""" + # get the model and patch_data reference + model = self._model + patch_data = self.patch_data + + # automode-related attributes + fill = patch_data["fill"] + + # BN calibration related in autonas/fastnas mode + if is_configurable(model): + # indicates that model is being trained in which case we want max sub-net for eval + if patch_data["autosampled"]: + sample(model, sample_func=max) + + # check current count since last sample operation + count_since_sample = patch_data["train_count_since_sample"] + + # logic to determine whether and how to calibrate BN statistics + queue = patch_data["queue"] + if forward_loop is None and count_since_sample < MODELOPT_BN_CALIB_ITERS: + if len(queue) < queue.maxlen and self._is_modelopt_queue_needed(model): + warnings.warn( + "modelopt data queue not filled! Inference results can be inaccurate." + ) + else: + forward_loop = partial( + run_forward_loop, + data_loader=queue, + collect_func=lambda x: x, + ) + elif forward_loop is not None: + # in this case we want to fill up the queue since new data is provided. + patch_data["fill"] = True + + # call reset_bn in eval mode + if forward_loop is not None: + with _modelopt_eval_recursion_guard(model): + self._reset_bn(model, forward_loop) + patch_data["train_count_since_sample"] = MODELOPT_BN_CALIB_ITERS + + # resetting modelopt_fill + patch_data["fill"] = fill + + @staticmethod + def _is_modelopt_queue_needed(model) -> bool: + return any(isinstance(m, _BatchNorm) and m.track_running_stats for m in model.modules()) + + @classmethod + def _reset_bn(cls, model: nn.Module, forward_loop: ForwardLoop): + """Calibrate BN statistics by the provided data loader. + + Args: + model: The model to be calibrated + forward_loop: A ``Callable`` that takes a model as input and runs a pre-defined forward + loop on it using data that is suitable for BN calibration. + """ + # skip if there is no batch normalization layer in the network that + # requires tracking running stats + if not cls._is_modelopt_queue_needed(model): + return + + is_training = model.training + momentum = {} + + # modify all batch norms to track cumulative moving average + for name, m in model.named_modules(): + if isinstance(m, _BatchNorm): + m.reset_running_stats() + m.train() + momentum[name] = m.momentum + m.momentum = None + + # do some forward passes to track BN statistics + with torch.no_grad(): + forward_loop(model) + + # reset network into its original state + model.train(is_training) + for name, m in model.named_modules(): + if name in momentum: + m.momentum = momentum[name] + + def _hook_pre_sample(self) -> None: + """Optional hook to be called before sample-related operations (sample & select).""" + patch_data = self.patch_data_or_empty + patch_data["train_count_since_sample"] = 0 + patch_data["autosampled"] = False + + @property + def sample_during_training(self) -> bool: + """Indicates whether we should sample a new subnet during training.""" + return True + + def _hook_pre_forward(self, *args, **kwargs) -> None: + """Optional hook that is called before the original forward function is called.""" + patch_data = self.patch_data + mod = self._model + if not (mod.training or patch_data["fill"]): + return + + # run sample step and update metadata related to sampling + def process_data(data): + return torch_to(torch_detach(data), device="cpu", non_blocking=True) + + if mod.training or patch_data["fill"]: + # sample new model configuration in training mode and indicate it was auto-sampled + if mod.training: + if self.sample_during_training: + sample(mod) + patch_data["autosampled"] = True + patch_data["train_count_since_sample"] += 1 # update train mode count + # push detached CPU copy of data to queue (once full only update periodically) + queue = patch_data["queue"] + if len(queue) < queue.maxlen or patch_data["count"] % 100 == 0 or patch_data["fill"]: + queue.append((*process_data(args), process_data(kwargs))) + patch_data["count"] += 1 + + +class IterativeSearcher(BaseSearcher, ABC): + """Base class for iterative search algorithms.""" + + iter_num: int + num_satisfied: int + constraints_func: ConstraintsFunc + candidate: dict[str, Any] + best: SearchStateDict + samples: dict[str, Any] + history: dict[str, Any] + best_history: dict[str, Any] + + @property + def default_search_config(self) -> SearchConfig: + """Get the default config for the searcher.""" + return { + **super().default_search_config, + "num_iters": 5000, + "max_iter_data_loader": 50, # plenty to calibrate BN statistics + } + + @property + def default_state_dict(self) -> SearchStateDict: + """Return default state dict.""" + return { + "iter_num": 0, + "num_satisfied": 0, + "candidate": {}, + "best": {"metric": -float("inf"), "constraints": None}, + "samples": {}, + "history": {"metric": [], "constraints": defaultdict(list)}, + "best_history": {"iter_num": [], "candidates": []}, + } + + def sanitize_search_config(self, config: Optional[SearchConfig]) -> SearchConfig: + """Sanitize the search config dict.""" + config = super().sanitize_search_config(config) + assert config["score_func"] is not None, "Please provide `score_func`!" + return config + + def _sample(self) -> dict[str, Any]: + return {"config": sample(self.model)} + + @abstractmethod + def sample(self) -> dict[str, Any]: + """Sample and select new sub-net configuration and return configuration.""" + raise NotImplementedError + + def before_search(self) -> None: + """Ensure that the model is actually configurable and ready for eval.""" + super().before_search() + + # Initialize the constraints functor + self.constraints_func = get_constraints_func( + self.model, self.constraints, self.dummy_input, deployment=self.deployment + ) + + # Set the model to max subnet and put model into eval mode with potential calibaration + sample(self.model, max) + prep_for_eval(self.model, self.forward_loop) + + # check for configurability of the models; otherwise we don't need to run the search. + if not is_configurable(self.model): + warnings.warn( + "Provided model does not contain configurable hparams with multiple choices! " + "Running only one iteration." + ) + self.config["num_iters"] = 1 + + # do a sanity check and profile the constraints + # This way we also fill up the interpolation table for latency with min, centroid, max! + from .algorithms import profile # TODO: hack until we refactor files + + profile( + self.model, + constraints=self.constraints_func, + strict=True, + verbose=self.config["verbose"], + use_centroid=True, + ) + + def run_search(self) -> None: + """Run iterative search loop.""" + num_iters = self.config["num_iters"] + verbose = self.config["verbose"] + + # run search loop + if verbose: + pbar = tqdm.trange( + self.iter_num, + num_iters, + initial=self.iter_num, + total=num_iters, + position=0, + leave=True, + ) + + for self.iter_num in range(self.iter_num, num_iters): + self.before_step() + self.run_step() + self.after_step() + + if verbose: + info = { + "num_satisfied": self.num_satisfied, + "metric": self.candidate["metric"], + "constraints": self.candidate["constraints"], + "best_subnet_metric": self.best["metric"], + "best_subnet_constraints": self.best["constraints"], + } + + # display the full stats only once a while + if len(self.history["metric"]) == 100: + info["metric/stats"] = stats(self.history["metric"]) + info["constraints/stats"] = { + name: stats(vals) for name, vals in self.history["constraints"].items() + } + self.history["metric"].clear() + self.history["constraints"].clear() + + def _recursive_format(obj, fmt): + if isinstance(obj, float): + return num2hrb(obj) + if isinstance(obj, dict): + return {k: _recursive_format(v, fmt) for k, v in obj.items()} + return obj + + pbar.update() + info = _recursive_format(info, fmt="{:.4g}") + pbar.set_description(f'[num_satisfied] = {info["num_satisfied"]}') + + if self.early_stop(): + break + + if verbose: + pbar.close() + print(f'[best_subnet_constraints] = {info["best_subnet_constraints"]}') + + def after_search(self) -> None: + """Select best model.""" + super().after_search() + select(self.model, self.best["config"]) + # if self.has_score: + # final_score = self.eval_score(silent=False) + # print(f"Final score for the searched model: {final_score}") + + def before_step(self) -> None: + """Run before each iterative step.""" + pass + + def run_step(self) -> None: + """The main routine of each iterative step.""" + # sample and select the candidate + self.candidate = self.sample() + + # obtain the independent config and hparams + # TODO: consider whether we want to eventually bring back "active_config" from omnimizer + # to improve the efficiency of the search and avoid redundant sampling. This was used to + # avoid testing two architectures that only differed in modules that were not active, i.e., + # due to reduced depth. + independent_config = self._configurable_config(self.model) + + # serialize and hash the candidate + buffer = json.dumps({"config": independent_config}, sort_keys=True) + ckey = hashlib.sha256(buffer.encode()).hexdigest() + + if ckey not in self.samples: + # check constraints + self.candidate["is_satisfied"], self.candidate["constraints"] = self.constraints_func() + + # evaluate the metric + if self.candidate["is_satisfied"]: + self.candidate["metric"] = self.eval_score() + else: + self.candidate["metric"] = -float("inf") + + self.samples[ckey] = copy.deepcopy(self.candidate) + else: + self.candidate = copy.deepcopy(self.samples[ckey]) + + self.num_satisfied += int(self.candidate["is_satisfied"]) + is_best = self.candidate["metric"] >= self.best["metric"] or self.iter_num == 0 + self.candidate["is_best"] = is_best + + # update the stats if satisfied + if self.candidate["is_satisfied"]: + self.history["metric"].append(self.candidate["metric"]) + for name, val in self.candidate["constraints"].items(): + self.history["constraints"][name].append(val) + + # update the best if necessary + if is_best: + self.best = copy.deepcopy(self.candidate) + self.best_history["iter_num"].append(self.iter_num) + self.best_history["candidates"].append(self.best) + + # save the state dict + if self.iter_num % 100 == 0 or self.iter_num == self.config["num_iters"] - 1 or is_best: + self.save_search_checkpoint() + + def after_step(self) -> None: + """Run after each iterative step.""" + pass + + def early_stop(self) -> bool: + """Check if we should early stop the search if possible.""" + return False + + def _configurable_config(self, model) -> dict[str, Any]: + """Returns the config dict of the configurable hyperparameters of the model.""" + return get_subnet_config(model, configurable=True) + + +class RandomSearcher(IterativeSearcher): + """An iterative searcher that samples subnets randomly.""" + + def sample(self) -> dict[str, Any]: + """Random sample new subset during each steo.""" + return self._sample() + + +class EvolveSearcher(IterativeSearcher): + """An iterative searcher that uses an evolutionary algorithm to optimize the subnet config.""" + + population: list[dict[str, Any]] + candidates: list[dict[str, Any]] + + @property + def default_search_config(self) -> SearchConfig: + """Default search config contains additional algorithm parameters.""" + return { + **super().default_search_config, + "population_size": 100, + "candidate_size": 25, + "mutation_prob": 0.1, + } + + @property + def default_state_dict(self) -> SearchStateDict: + """Return default state dict.""" + return { + **super().default_state_dict, + "population": [], + "candidates": [], + } + + def _mutate(self, input: dict[str, Any]) -> dict[str, Any]: + output = self._sample() + # only considers independent hparams + output["config"] = self._configurable_config(self.model) + for var in output: # pylint: disable=C0206 + for name in output[var]: + if random.random() > self.config["mutation_prob"]: + output[var][name] = input[var][name] + # returns reduced, independent config + return output + + def _crossover(self, inputs: list[dict[str, Any]]) -> dict[str, Any]: + output = {"config": {}} + # only considers independent hparams + for name in self._configurable_config(self.model): + output["config"][name] = random.choice(inputs)["config"][name] + # returns reduced, independent config + return output + + def sample(self) -> dict[str, Any]: + """Sampling a new subnet involves random sampling, mutation, and crossover.""" + if not self.candidates: + return self._sample() + if len(self.population) < self.config["population_size"] // 2: + output = self._mutate(random.choice(self.candidates)) + else: + output = self._crossover(random.sample(self.candidates, 2)) + # returns full config + select(self.model, output["config"]) + output["config"] = get_subnet_config(self.model) + return output + + def before_search(self) -> None: + """Set the lower bound of the constraints to 0.85 * upper bound before search.""" + super().before_search() + self.constraints_func.set_rel_lower_bounds(rel_lower=0.85) + + def before_step(self) -> None: + """Update candidates and population before each iterative step.""" + if len(self.population) >= self.config["population_size"]: + all_candidates = sorted(self.population, key=lambda x: x["metric"], reverse=True) + self.candidates = all_candidates[: self.config["candidate_size"]] + self.population = [] + + def after_step(self) -> None: + """Update population after each iterative step.""" + if self.candidate["is_satisfied"]: + self.population.append(copy.deepcopy(self.candidate)) + + +def convert_searchspace( + model: nn.Module, config: ModeloptBaseConfig, patch_manager_type: type[PatchManager] +) -> ConvertReturnType: + """Convert given model into a search space.""" + # initialize the true module if necessary + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + + # check if the model is in channels-last format + # NOTE: not supported because we have not implemented channel sorting for channels-last format + if is_channels_last(model): + raise RuntimeError(f"Channels-last format in {type(model).__name__} is not supported!") + + # put together metadata before starting to modify model + metadata = {"model_attributes": get_model_attributes(model)} + + # now convert the model to a search space + search_space = generate_search_space(model, rules=config.model_dump()) + + # sanity check if any module has been converted + if not search_space.is_configurable(): + raise ApplyModeError( + "The model does not contain any configurable hyperparameters! Please check the" + " documentation for modules and config and how to get a configurable model." + ) + + # activate the max subnet that can be used for subsequent tasks like training + # (note it should be activated already but we do it here again to be sure) + sample(model, max) + + # search space requires a patch + patch_manager_type(model).patch() + + # get current config store in metadata + metadata["subnet_config"] = get_subnet_config(model) + + # return converted model as well as metadata + return model, metadata + + +def convert_autonas_searchspace(model: nn.Module, config: ModeloptBaseConfig) -> ConvertReturnType: + """Convert search space for AutoNAS mode with correct patch manager.""" + return convert_searchspace(model, config, AutoNASPatchManager) + + +def restore_autonas_searchspace( + model: nn.Module, config: ModeloptBaseConfig, metadata: MetadataDict +) -> nn.Module: + """Restore search space for AutoNAS mode with correct patch manager.""" + return restore_searchspace(model, config, metadata, AutoNASPatchManager) + + +def restore_searchspace( + model: nn.Module, + config: ModeloptBaseConfig, + metadata: MetadataDict, + patch_manager: type[PatchManager], +) -> nn.Module: + """Restore a search space from the given model.""" + # retrieve metadata model attributes + model_attributes = metadata["model_attributes"] + + # initialize the true module if necessary + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + + # set up train/eval mode accordingly + is_training = model.training + model.train(model_attributes["training"]) + + # run regular convert entrypoint + model, metadata_new = convert_searchspace(model, config, patch_manager) + + # compare new model attributes with provided model attributes + new_attributes = metadata_new["model_attributes"] + unmatched_keys = compare_dict(model_attributes, new_attributes) + if len(unmatched_keys) > 0: + error_msg = ["Following keys in your model do not match the checkpoint:"] + error_msg += [ + f"{k}: Original={model_attributes.get(k, 'N/A')}; New={new_attributes.get(k, 'N/A')}" + for k in unmatched_keys + ] + raise ApplyModeError("\n\t".join(error_msg)) + + # ensure model is back in mode it started out in + model.train(is_training) + + # select subnet config stored in metadata + select(model, metadata["subnet_config"]) + + # return converted model + return model + + +def update_autonas_metadata( + model: nn.Module, config: ModeloptBaseConfig, metadata: MetadataDict +) -> None: + """Update subnet config to current subnet config of model.""" + metadata["subnet_config"] = get_subnet_config(model) + + +def export_searchspace(model: nn.Module, config: ExportConfig) -> ConvertReturnType: + """Export a subnet configuration of the search space to a regular model.""" + # sanity check to avoid DP/DDP here in the entrypoint + model = unwrap_model(model, raise_error=True) + + # store config from model if we can find it for a future convert/restore process + subnet_config = get_subnet_config(model) + + # Check for patching and calibration + if PatchManager.is_patched(model): + manager = PatchManager.get_manager(model) + if config.calib: + manager.call_post_eval() + manager.unpatch() + + # export model in-place + model = SearchSpace(model).export() + + # construct metadata + metadata = { + "subnet_config": subnet_config, + } + + return model, metadata + + +def restore_export(model: nn.Module, config: ExportConfig, metadata: MetadataDict) -> nn.Module: + """Restore & export the subnet configuration of the search space to a regular model.""" + # select subnet config provided in metadata + select(model, metadata["subnet_config"], strict=config["strict"]) + + # run export + model, metadata_new = export_searchspace(model, config) + + # double check metadata + unmatched_keys = compare_dict(metadata, metadata_new) + if unmatched_keys: + raise ApplyModeError(f"Unmatched metadata={unmatched_keys}!") + + return model diff --git a/modelopt/torch/nas/config.py b/modelopt/torch/nas/config.py new file mode 100644 index 0000000000..b9ff884997 --- /dev/null +++ b/modelopt/torch/nas/config.py @@ -0,0 +1,94 @@ +# 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. + +"""Default configurations for NAS modes.""" + +from pydantic import create_model + +from modelopt.torch.opt.config import ( + ModeloptBaseConfig, + ModeloptField, + get_kwargs_for_create_model_with_rules, +) + +from .registry import DMRegistry + +__all__ = ["AutoNASConfig", "ExportConfig"] + + +def _get_ratio_list(): + return (0.5, 0.67, 1.0) + + +def _conv_config(): + return { + "channels_ratio": _get_ratio_list(), + "kernel_size": (), + "channel_divisor": 32, + } + + +def _norm_lin_config(): + return { + "features_ratio": _get_ratio_list(), + "feature_divisor": 32, + } + + +AutoNASConfig: type[ModeloptBaseConfig] = create_model( + "AutoNASConfig", + **get_kwargs_for_create_model_with_rules( + registry=DMRegistry, + default_rules={ + "nn.Conv1d": _conv_config(), + "nn.Conv2d": _conv_config(), + "nn.Conv3d": _conv_config(), + "nn.ConvTranspose1d": _conv_config(), + "nn.ConvTranspose2d": _conv_config(), + "nn.ConvTranspose3d": _conv_config(), + "nn.Linear": _norm_lin_config(), + "nn.BatchNorm1d": _norm_lin_config(), + "nn.BatchNorm2d": _norm_lin_config(), + "nn.BatchNorm3d": _norm_lin_config(), + "nn.SyncBatchNorm": _norm_lin_config(), + "nn.InstanceNorm1d": _norm_lin_config(), + "nn.InstanceNorm2d": _norm_lin_config(), + "nn.InstanceNorm3d": _norm_lin_config(), + "nn.LayerNorm": _norm_lin_config(), + "nn.GroupNorm": {k: v for k, v in _conv_config().items() if k != "kernel_size"}, + "nn.Sequential": {"min_depth": 0}, + }, + doc='Configuration for the ``"autonas"`` mode.', + ), +) + + +class ExportConfig(ModeloptBaseConfig): + """Configuration for the export mode. + + This mode is used to export a model after NAS search. + """ + + strict: bool = ModeloptField( + default=True, + title="Strict export", + description="Enforces that the subnet configuration must exactly match during export.", + ) + + calib: bool = ModeloptField( + default=False, + title="Calibration", + description="Whether to calibrate the subnet before exporting.", + ) diff --git a/modelopt/torch/nas/conversion.py b/modelopt/torch/nas/conversion.py new file mode 100644 index 0000000000..3f82706aa0 --- /dev/null +++ b/modelopt/torch/nas/conversion.py @@ -0,0 +1,119 @@ +# 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. + +"""Main APIs+entrypoints for model pruning.""" + +from torch import nn + +from modelopt.torch.opt.conversion import apply_mode +from modelopt.torch.opt.mode import ModeLike +from modelopt.torch.utils import ModelLike, unwrap_model + +from .mode import NASModeRegistry + +__all__ = ["export", "convert"] + + +def convert(model: ModelLike, mode: ModeLike) -> nn.Module: + """Convert a regular PyTorch model into a model that supports design space optimization. + + Args: + model: A model-like object. Can be an nn.Module, a model class type, or a tuple. + Tuple must be of the form ``(model_cls,)`` or ``(model_cls, args)`` or ``(model_cls, args, kwargs)``. + Model will be initialized as ``model_cls(*args, **kwargs)``. + mode: A (list of) string(s) or Mode(s) or a list of tuples containing the mode and its + config indicating the desired mode(s) (and configurations) for the convert + process. Modes set up the model for different algorithms for model optimization. The + following modes are available: + + * :class:`"autonas"`: The ``model`` will + be converted into a search space and set up to automatically perform operations + required for AutoNAS-based model training, evaluation, and search. The mode's config + is described in :class:`AutoNASConfig`. + + If the mode argument is specified as a dictionary, the keys should indicate the mode and + the values specify the per-mode configuration. If not provided, then default + configuration would be used. + + Returns: + A converted model with the original weights preserved that can be used for model + optimization. + + .. note:: + + Note that model wrappers (such as DataParallel/DistributedDataParallel) are `not` supported + during the convert process. Please wrap the model after the convert process. + + .. note:: + ``convert()`` relies on `monkey patching `_ to + augment the ``forward()``, ``eval()``, and ``train()`` methods of ``model`` as well as + augment individual modules to make them dynamic. This renders the conversion incompatible + with other monkey patches to those methods and modules! Note that ``convert()`` is still + fully compatible with inheritance. + + .. note:: + + #. Configs can be customized for individual layers using + `glob expressions `_ on + qualified submodule names, e.g., as shown for ``nn.Conv2d`` in the above example. + + #. Keys in the config that appear earlier in a dict have lower priority, e.g., + ``backbone.stages.1.0.spatial_conv`` will have ``out_channels_ratio`` + ``[0.334, 0.5, 0.667, 1.0]``, not ``[1.0]``. + + #. Config entries without layer qualifiers are also supported, e.g., as shown for + ``nn.Sequential`` in the above example. + + #. Mixed usage of configurations with and without layer qualifiers is supported for + *different* layers, e.g., as shown for ``nn.Conv2d`` and ``nn.Sequential`` in the above + example. For a specific layer type, only configurations with *or* without layer + qualifiers are supported. + + #. Use ``*`` as a wildcard matching any layer. + """ + # apply mode and handle model-like object with wrapper + return apply_mode(model, mode, registry=NASModeRegistry) + + +def export(model: nn.Module, strict: bool = True, calib: bool = False) -> nn.Module: + """Export a pruned subnet to a regular model. + + Args: + model: The pruned subnet to be exported. + strict: Raise an error when the config does not contain all necessary keys. + calib: Whether to calibrate the subnet to be exported. + + Returns: + The current active subnet in regular PyTorch model format. + + .. note:: + + if model is a wrapper such as DistributedDataParallel, it will be unwrapped, e.g., + ``model.module`` will be returned. + """ + # unwrap a DP/DDP model + model = unwrap_model( + model, + warn=True, + msg=( + f"Unwrapping a {type(model).__name__} model for export! Note that the export is" + " in-place and the model wrapper should be re-created after export since the wrapper" + " might not support changing parameters after initialization." + ), + ) + + # apply export mode and return model + config = {"strict": strict, "calib": calib} + return apply_mode(model, [("export", config)], registry=NASModeRegistry) diff --git a/modelopt/torch/nas/hparams/__init__.py b/modelopt/torch/nas/hparams/__init__.py new file mode 100644 index 0000000000..5bed52c2db --- /dev/null +++ b/modelopt/torch/nas/hparams/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""Module containing implementations for various types of hyperparameters.""" + +from .concat import * +from .container import * diff --git a/modelopt/torch/nas/hparams/concat.py b/modelopt/torch/nas/hparams/concat.py new file mode 100644 index 0000000000..c266f07216 --- /dev/null +++ b/modelopt/torch/nas/hparams/concat.py @@ -0,0 +1,281 @@ +# 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. + +"""Hparam that represent concat.""" + +from collections import defaultdict +from itertools import product +from math import prod +from typing import Callable, Iterator, Optional + +import numpy as np +import torch + +from modelopt.torch.trace import ConcatSymbol, Symbol + +from ..traced_hp import TracedHp, TracedHpRegistry + +__all__ = ["ConcatTracedHp"] + + +@TracedHpRegistry.register(ConcatSymbol) +class ConcatTracedHp(TracedHp): + """Concat hparam that can stitch together input hparams.""" + + _inputs: list[TracedHp] + _sum_to_combo: dict[int, tuple[int]] + _hp_start_idx: torch.LongTensor + + @staticmethod + def _choice_combos(*all_choices) -> Iterator[tuple[int, ...]]: + # compute total number of combinations + n_combos = prod(len(c_list) for c_list in all_choices) + + # we don't wanna iterate over more than that to keep it fast + n_max = 2e5 + + # if we have less than n_max combinations, we can iterate over all of them + if n_combos <= n_max: + # use itertools.product and iterate over all combinations + yield from product(*all_choices) + return + + # otherwise, we use a pruned set of combinations based on immediate vicinity of each value + + # range of choices + c_min = min(min(c_list) for c_list in all_choices) + c_max = max(max(c_list) for c_list in all_choices) + + # figure out for each list what's the closest index for each value in the choice range + idx_lookup = { + c: [min(range(len(c_list)), key=lambda i: abs(c_list[i] - c)) for c_list in all_choices] + for c in range(c_min, c_max + 1) + } + + # solve equation for r such that n_combos = (c_max - c_min) * (2r+1)^n \approx n_max + # use r = max(1, r) to avoid edge cases + r = max(int((n_max / (c_max - c_min)) ** (1 / len(all_choices)) - 1) // 2, 1) + + # yield reduced number of combinations (note that we might have duplicates; hence we filter) + # each combination is based on the product of choices centered around the value closest to + # the target value c. We use r to control the size of the neighborhood. + processed = set() + for c in range(c_min, c_max + 1): + all_choices_reduced = [] + for i, c_list in enumerate(all_choices): + idx_closest = idx_lookup[c][i] + all_choices_reduced.append(c_list[max(0, idx_closest - r) : idx_closest + r + 1]) + for combo in product(*all_choices_reduced): + if combo not in processed: + processed.add(combo) + yield combo + + def _get_sum_to_combo(self) -> dict[int, tuple[int, ...]]: + """Return dict mapping sum to individual hparam active values.""" + # If two hparams are the same, we need their values to be the same and we need to track that + hps_unique = list(set(self._inputs)) + idxs_unique = [hps_unique.index(hp) for hp in self._inputs] + + # Keep the combination with the smallest difference between the individual values + # e.g. for sum 48, prefer [24, 24] over [16, 32] or [32, 16] + sum_to_combo = {} + sum_to_cost = defaultdict(lambda: float("inf")) + for combo in self._choice_combos(*[hp.choices for hp in hps_unique]): + # get full combo and sum + combo_full = tuple(combo[idx] for idx in idxs_unique) + s = sum(combo_full) + # check cost (sum of absolute differences between average and individual values) + cost = sum(abs(c - s / len(combo_full)) for c in combo_full) + # update if better + if cost < sum_to_cost[s]: + sum_to_cost[s] = cost + sum_to_combo[s] = combo_full + return sum_to_combo + + @property # type: ignore[override] + def active(self) -> int: + """Return the sum of active values of all hparams.""" + active = 0 + for hp in self._inputs: + active += hp.active + return active + + @active.setter + def active(self, val: Optional[int]): + """Set the active value with a sanity check for choices and dynamic hparams.""" + val = self.original if val is None else val + assert val in self._choices, f"val = {val}, choices = {self.choices}" + assert val in self._sum_to_combo, f"Invalid value: {val} not in {self._sum_to_combo.keys()}" + for hp_in, v in zip(self._inputs, self._sum_to_combo[val]): + assert v in hp_in.choices + hp_in._active = v + + @property + def active_slice(self) -> TracedHp.ActiveSlice: + """Return the currently active sorted indices or slice corresponding to the active value. + + Example: + hp1(max=4, active=2), hp2(max=4, active=4) + Then the active indices in the concatenated hp are [0, 1, 4, 5, 6, 7] + """ + # get vanilla slice first (with sanity check) + as_simple = super().active_slice + assert isinstance(as_simple, slice), "ConcatTracedHp does not support its own slice order!" + + # get the list of slices as active indices in the concatenated hp + as_hps = [hp.active_slice for hp in self._inputs] + as_hps = [torch.arange(_as.stop)[_as] if isinstance(_as, slice) else _as for _as in as_hps] + + # concatenate the slices with correct offsets + active_slice = torch.cat([_as + self._hp_start_idx[i] for i, _as in enumerate(as_hps)]) + + # check if active_slice corresponds to the vanilla slice + if torch.equal(active_slice, torch.arange(as_simple.stop)[as_simple]): + return as_simple + + return active_slice + + def _get_importance(self) -> TracedHp.Importance: + # compute the regular importance (this hparam + registered dependencies) + imp_total = super()._get_importance() + + # now we add in the importance from the individual hparams contributing to the concat node + # concat inputs are not registered as dependencies, so we need to add them here + # note that we use ``hp.importance`` instead of the private API ``hp._get_importance()`` + # here. Why? We can have the situation where only one of the concat inputs is searchable. + # We need to account for that and compute the "partial" importances hence! + imps_cat = [hp.importance if hp.is_configurable else None for hp in self._inputs] + + # clean up importances that are None + imps_cat = [imp or torch.zeros(hp.max) for hp, imp in zip(self._inputs, imps_cat)] + + # Let's split imp_total + imps_split = torch.tensor_split(imp_total, self._hp_start_idx[1:-1]) + + # Added splitted imps + imps = [imp + imp_cat.to(imp) for imp, imp_cat in zip(imps_split, imps_cat)] + + # We need to aggregate between split importances when the come from the same hparam! + imps = [ + sum([imp_ for imp_, hp in zip(imps, self._inputs) if hp is self._inputs[i]]) + for i, imp in enumerate(imps) + ] + + # concatenate and return + return torch.cat(imps) + + def _split_order(self, order: torch.Tensor) -> list[torch.Tensor]: + """Split the order tensor into the individual order tensors for each input hparam.""" + return [ + ( + order[torch.logical_and(order >= start, order < start + hp.max)] - start + if hp.is_configurable and hp.is_sortable + else torch.arange(hp.max, device=order.device) + ) + for hp, start in zip(self._inputs, self._hp_start_idx[:-1]) + ] + + @torch.no_grad() + def _enforce_order(self, order: Optional[torch.Tensor] = None) -> None: + """Enforcing the order. + + Note that we only support intra-group sorting but **not** inter-group sorting. Inter-group + sorting would not be compatible with the concat operation and is thus not well-defined. + + E.g., consider a concat operation with multiple inputs from convs. Then a channel that is + currently assigned to the first group (i.e. first conv) cannot be resorted to appear in any + other group/conv. Hence, we need to ensure that we maintain inter-group consistency. + """ + # vanilla case + if order is None: + for hp in self._inputs: + hp._enforce_order(order) # private API is fine here since it is just None + return super()._enforce_order(order) + + # we have to group the order by the hparams to enforce order on each group separately with + # the correct offset substracted + orders_split = self._split_order(order) + + # ensure that orders are consistent if they come from the same hparam + for hp, order in zip(self._inputs, orders_split): + assert all( + torch.equal(order, order_) + for hp_, order_ in zip(self._inputs, orders_split) + if hp is hp_ + ) + + # enforce order now ... + # note that we use the public API here for extra sanity checks + for hp, partial_order in zip(self._inputs, orders_split): + if hp.is_configurable and hp.is_sortable: + hp.enforce_order(partial_order) + + # recombine the order into a single tensor without inter-group sorting for dependencies + order_intra_only = torch.cat( + [partial + start for partial, start in zip(orders_split, self._hp_start_idx[:-1])] + ) + + # sort dependencies now + super()._enforce_order(order_intra_only) + self._slice_order = None # we don't store slice_order in a ConcatTracedHp + + def _resolve_dependencies( + self, sym: Symbol, get_hp: Callable[[Symbol], TracedHp] + ) -> dict[Symbol, TracedHp]: + # check that we have the correct symbol + assert type(sym) is ConcatSymbol, f"Sym should be {type(self)}: {sym}." + + # construct input hparams + # NOTE: input syms are hidden and corresponding registered symbol should be last dependency! + # --> if that's not the case we cannnot retrieve the hps needed. + assert all(isinstance(s, ConcatSymbol.Input) for s in sym.input_syms) + assert all(s._dependencies[-1]._parent == s for s in sym.input_syms) + self._inputs = [get_hp(s._dependencies[-1]) for s in sym.input_syms] + + def _get_hp_with_input(s: Symbol) -> TracedHp: + assert type(sym) is ConcatSymbol + for s_, hp in zip(sym.input_syms, self._inputs): + if s is s_: + return hp + return get_hp(s) + + # resolve regular dependencies of all input hparams + mapping: dict[Symbol, TracedHp] = {} + for hp_in, sym_in in zip(self._inputs, sym.input_syms): + mapping.update(hp_in._resolve_dependencies(sym_in, _get_hp_with_input)) + + # remove input syms from mapping... + for sym_in in set(sym.input_syms): + mapping.pop(sym_in) + + # resolve regular dependencies + mapping.update(super()._resolve_dependencies(sym, get_hp)) + + # compute sum_to_combo values and make sure they are consistent with choices + sum_to_combo = {s: val for s, val in self._get_sum_to_combo().items() if s in self.choices} + self._sum_to_combo = sum_to_combo + self.choices = list(self._sum_to_combo) + + # update choices for inputs from sum_to_combo and make them non-configurable + for i, hp_in in enumerate(self._inputs): + hp_in.choices = [combo[i] for combo in sum_to_combo.values()] + hp_in._is_configurable = False + + # compute start indices for input hp + # NOTE: the last index is the length of the concatenated hp + hp_start_idx = np.concatenate(([0], np.cumsum([hp.max for hp in self._inputs]))) + self._hp_start_idx = torch.asarray(hp_start_idx, dtype=torch.long) + + return mapping diff --git a/modelopt/torch/nas/hparams/container.py b/modelopt/torch/nas/hparams/container.py new file mode 100644 index 0000000000..c3f1840662 --- /dev/null +++ b/modelopt/torch/nas/hparams/container.py @@ -0,0 +1,48 @@ +# 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. + +"""Hparam for Depth.""" + +from typing import Callable + +from modelopt.torch.trace import Symbol, SymDepth + +from ..traced_hp import TracedHp, TracedHpRegistry + +__all__ = ["DepthHparam"] + + +@TracedHpRegistry.register(SymDepth) +class DepthHparam(TracedHp): + """Hparam describing depth.""" + + def _resolve_dependencies( + self, sym: Symbol, get_hp: Callable[[Symbol], TracedHp] + ) -> dict[Symbol, TracedHp]: + assert isinstance(sym, SymDepth), f"Unexpected type {type(sym)} for {sym}!" + assert not sym._dependencies, "Depth should not have any dependencies!" + assert not sym._parent, "Depth should not have any parents!" + + # we can only skip layers when all layers after are skippable + choices = {sym.max_depth} + for i in range(sym.max_depth - 1, -1, -1): + if not sym.is_skippable(i): + break + choices.add(i) + + # record choices and skippable layers + self.choices = list(choices & set(self.choices)) + + return super()._resolve_dependencies(sym, get_hp) diff --git a/modelopt/torch/nas/mode.py b/modelopt/torch/nas/mode.py new file mode 100644 index 0000000000..0509631b42 --- /dev/null +++ b/modelopt/torch/nas/mode.py @@ -0,0 +1,135 @@ +# 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. + +"""Module implementing and describing modes that can be used during the NAS convert process. + +Check out :meth:`mtn.convert ` to learn more about modes. +""" + +from typing import Optional + +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.mode import ( + ConvertEntrypoint, + RestoreEntrypoint, + UpdateEntrypoint, + _ModeDescriptor, + _ModeRegistryCls, +) +from modelopt.torch.opt.searcher import BaseSearcher + +from .autonas import ( + EvolveSearcher, + convert_autonas_searchspace, + export_searchspace, + restore_autonas_searchspace, + restore_export, + update_autonas_metadata, +) +from .config import AutoNASConfig, ExportConfig + +__all__ = ["AutoNASModeDescriptor", "ExportModeDescriptor"] + +NASModeRegistry = _ModeRegistryCls() + + +@NASModeRegistry.register_mode +class AutoNASModeDescriptor(_ModeDescriptor): + """Class to describe the ``"autonas"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "autonas" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return AutoNASConfig + + @property + def next_modes(self) -> Optional[set[str]]: + """Modes that must immediately follow this mode.""" + return {"export", "kd_loss", "quantize", "sparse_magnitude", "sparse_gpt"} + + @property + def export_mode(self) -> Optional[str]: + """The mode that corresponds to the export mode of this mode.""" + return "export" + + @property + def search_algorithm(self) -> type[BaseSearcher]: + """Specifies the search algorithm to use for this mode (if any).""" + return EvolveSearcher + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return convert_autonas_searchspace + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_autonas_searchspace + + @property + def update_for_save(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models state before saving.""" + return update_autonas_metadata + + @property + def update_for_new_mode(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models state before new mode.""" + return update_autonas_metadata + + +@NASModeRegistry.register_mode +class ExportModeDescriptor(_ModeDescriptor): + """Class to describe the ``"export"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "export" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return ExportConfig + + @property + def is_export_mode(self) -> bool: + """Whether the mode is an export mode. + + Returns: + True if the mode is an export mode, False otherwise. Defaults to False. + """ + return True + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return export_searchspace + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_export diff --git a/modelopt/torch/nas/modules/__init__.py b/modelopt/torch/nas/modules/__init__.py new file mode 100644 index 0000000000..3830617957 --- /dev/null +++ b/modelopt/torch/nas/modules/__init__.py @@ -0,0 +1,21 @@ +# 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. + +"""The dynamic equivalents for torch.nn modules.""" + +from .container import * +from .conv import * +from .linear import * +from .norm import * diff --git a/modelopt/torch/nas/modules/container.py b/modelopt/torch/nas/modules/container.py new file mode 100644 index 0000000000..0297c74437 --- /dev/null +++ b/modelopt/torch/nas/modules/container.py @@ -0,0 +1,98 @@ +# 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. + +"""Dynamic sequential implementation (variable depth based on torch.nn.modules.container).""" + +from collections import OrderedDict +from typing import Any, Callable + +from torch import nn + +from modelopt.torch.opt.dynamic import DynamicModule + +from ..registry import DMRegistry +from ..traced_hp import TracedHp + +__all__ = ["_DynamicSequential"] + + +def _activate_depth(func: Callable) -> Callable: + """A decorator for enabling dynamic depth within methods in _DynamicSequential. + + The decorated method (and methods within the decorated method) will see any attributes with + dynamic depth enabled. + """ + + def func_with_dynamic_depth(self: "_DynamicSequential", *args, **kwargs) -> Any: + """Call func and return value wrapped with temporarily enabling dynamic depth.""" + val = self._dynamic_depth + self._dynamic_depth = True + ret = func(self, *args, **kwargs) + self._dynamic_depth = val + return ret + + return func_with_dynamic_depth + + +@DMRegistry.register({nn.Sequential: "nn.Sequential"}) +class _DynamicSequential(DynamicModule): + """An ``nn.Sequential`` layer with dynamic hyperparams and variable ``depth``.""" + + _dynamic_depth: bool + + @_activate_depth + def forward(self, input): + """Forward with activated variable depth.""" + return super().forward(input) + + @_activate_depth + def export(self) -> nn.Module: + """Export with activated variable depth.""" + return super().export() + + @_activate_depth + def __repr__(self): + """__repr__ with dynamic depth enabled -> str(self) will only show active subnet.""" + return super().__repr__() + + @_activate_depth + def extra_repr(self): + """extra_repr with dynamic depth enabled -> str(self) will only show active subnet.""" + return super().extra_repr() + + @staticmethod + def _get_modules(mod: "_DynamicSequential", modules: OrderedDict) -> OrderedDict: + if mod.depth < len(modules) and mod._dynamic_depth: + return OrderedDict(list(modules.items())[: mod.depth]) + return modules + + def _setup(self): + # register temp attribute to keep track of whether dynamic depth is on + self._register_temp_attribute("_dynamic_depth", False) + + # register hyperparameters + self._register_hparam("depth", TracedHp(list(range(0, len(self) + 1)))) + + # register _modules as a dynamic attribute + self._register_dynamic_attribute("_modules", self._get_modules) + + def modify(self, *, min_depth: int = 0): + """Modify the dynamic choices of the module according to provided keyword arguments. + + Args: + min_depth: The minimum depth of the module. + """ + hp = self.get_hparam("depth") + hp.choices = [d for d in hp.choices if d >= min_depth] diff --git a/modelopt/torch/nas/modules/conv.py b/modelopt/torch/nas/modules/conv.py new file mode 100644 index 0000000000..8aff7490bf --- /dev/null +++ b/modelopt/torch/nas/modules/conv.py @@ -0,0 +1,253 @@ +# 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. + +"""Dynamic Conv implementations based on torch.nn.modules.conv.""" + +import itertools +from typing import Optional, Union + +import torch +from torch import nn + +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.utils import make_divisible, val2tuple + +from ..registry import DMRegistry +from ..traced_hp import TracedHp +from .utils import get_sliced_tensor, get_sliced_tensor_by_slices + +__all__ = ["_DynamicConvNd", "_DynamicConvTransposeNd"] + + +@DMRegistry.register({nn.Conv1d: "nn.Conv1d", nn.Conv2d: "nn.Conv2d", nn.Conv3d: "nn.Conv3d"}) +class _DynamicConvNd(DynamicModule): + @property + def ndim(self): + return len(self.kernel_size) + + @staticmethod + def _assert_input_format( + mod: "_DynamicConvNd", input: Union[tuple[torch.Tensor], torch.Tensor] + ) -> None: + if isinstance(input, tuple): + assert len(input) == 1, f"Expected single input, but got {input}." + input = input[0] + if input.dim() != mod.ndim + 2: + raise ValueError(f"Expected batched input (e.g. NCHW vs CHW), but got {input.shape}.") + + @staticmethod + def _get_padding( + mod: "_DynamicConvNd", padding: Union[str, tuple[int, ...]] + ) -> Union[str, tuple[int, ...]]: + """New padding such that output size stays the same as with max kernel size and padding.""" + if isinstance(padding, str): + return padding + + active_padding = [] + max_kernel_size = mod.get_hparam("kernel_size").max + for dim in range(mod.ndim): + max_ks = max_kernel_size[dim] # type: ignore[index] + active_ks = mod.kernel_size[dim] + assert max_ks % 2 == active_ks % 2, (max_ks, active_ks) + + max_p = padding[dim] + active_p = max_p - (max_ks - active_ks) * mod.dilation[dim] // 2 + assert active_p >= 0, (max_p, active_p) + + active_padding.append(active_p) + return tuple(active_padding) + + @staticmethod + def _get_bias(mod: "_DynamicConvNd", bias: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + return get_sliced_tensor(mod, bias, "out_channels") + + @staticmethod + def _get_channel_order() -> list[str]: + return ["out_channels", "in_channels"] + + @staticmethod + def _get_weight(mod: "_DynamicConvNd", weight: torch.Tensor) -> torch.Tensor: + slices = [] + + # get channel hparams in correct order + hps_c = [mod.get_hparam(c_name) for c_name in mod._get_channel_order()] + c_active = [hp.active for hp in hps_c] + + # add slice for 1st channel dimension + slices.append(hps_c[0].active_slice) + + # slice for 2nd channel dimension depends on groups + slice_1 = hps_c[1].active_slice + groups = mod.groups + if groups == 1: + slices.append(slice_1) + elif all(c % groups == 0 for c in c_active): + assert slice_1 == slice(c_active[1]) or c_active[0] == c_active[1] == groups + slices.append(slice(c_active[1] // groups)) + else: + raise NotImplementedError(f"Group size {groups} not supported!") + + # kernel size slices + max_kernel_size = mod.get_hparam("kernel_size").max + active_kernel_size = mod.kernel_size + assert isinstance(max_kernel_size, tuple) + for max_ks, active_ks in zip(max_kernel_size, active_kernel_size): + a = max_ks // 2 - active_ks // 2 + b = max_ks // 2 + active_ks // 2 + max_ks % 2 + slices.append(slice(a, b)) + + return get_sliced_tensor_by_slices(weight, slices) + + @staticmethod + def _get_groups(mod: "_DynamicConvNd", groups: int) -> int: + # retrieve channel hparams + hp_oc = mod.get_hparam("out_channels") + hp_ic = mod.get_hparam("in_channels") + + # vanilla cases + if groups == 1 or (hp_oc.active == hp_oc.original and hp_ic.active == hp_ic.original): + return groups + + # only support out_channels == in_channels if grouped + assert hp_oc.active == hp_ic.active, (hp_oc.active, hp_ic.active) + + # compute current number of groups + assert hp_oc.original % groups == 0, (hp_oc.original, groups) + active_groups = (hp_oc.active * groups) // hp_oc.original + + # sanity checks + assert isinstance(active_groups, int) + assert hp_oc.active % active_groups == 0, (hp_oc.active, active_groups) + assert hp_ic.active % active_groups == 0, (hp_ic.active, active_groups) + + return active_groups + + def _estimate_importance(self) -> TracedHp.Importance: + # for group > 1, we do not know how to handle it yet + if self.groups > 1: + return None + weight = self._parameters["weight"] # retrieve full weight tensor + c_in = weight.shape[1] + return torch.norm(torch.reshape(weight.detach().transpose(0, 1), (c_in, -1)), dim=1) + + def _setup(self): + # only support ungrouped conv or grouped conv with in_channels == out_channels + if self.groups == 1: + oc_step = ic_step = 1 + elif self.out_channels == self.in_channels: + oc_step = ic_step = self.out_channels // self.groups + else: + oc_step = self.out_channels + ic_step = self.in_channels + oc_choices = [c for c in range(1, self.out_channels + 1) if c % oc_step == 0] + ic_choices = [c for c in range(1, self.in_channels + 1) if c % ic_step == 0] + + # construct choices for kernel size hyperparameter + ks_set = {self.kernel_size} + # TODO: padding mode other than "zeros" is not supported since in other padding mode, conv + # will use "_reversed_padding_repeated_twice" instead of "padding" in the forward function. + if self.padding_mode == "zeros": + for ks in itertools.product(*[range(1, k + 1) for k in self.kernel_size]): + for d in range(self.ndim): + # check kernel size: smaller than original, same parity + if ks[d] > self.kernel_size[d] or ks[d] % 2 != self.kernel_size[d] % 2: + break + + # check padding: non-negative + if ( + isinstance(self.padding, tuple) + and self.padding[d] < (self.kernel_size[d] - ks[d]) * self.dilation[d] // 2 + ): + break + else: + ks_set.add(ks) + ks_choices = list(ks_set) + + # register hyperparameters + self._register_hparam("out_channels", TracedHp(oc_choices, self.out_channels)) + self._register_hparam("in_channels", TracedHp(ic_choices, self.in_channels)) + self._register_hparam("kernel_size", TracedHp(ks_choices, self.kernel_size)) + + # We restrict the input to be batched (e.g. NCHW) format so searchable_tensor_dims can be 1. + # Ideally searchable_tensor_dims for Conv2d would be -3, but its more common to use + # torch.cat([...], dim=1), and enforcing batched input would allow us to support these. + hook_handle = self.register_forward_pre_hook(self._assert_input_format) + self._register_temp_attribute( + "_hook_handle", hook_handle, del_hook=lambda m, n: m._hook_handle.remove() + ) + + # register dynamic attributes + self._register_dynamic_attribute("padding", self._get_padding) + self._register_dynamic_attribute("weight", self._get_weight) + self._register_dynamic_attribute("bias", self._get_bias) + self._register_dynamic_attribute("groups", self._get_groups) + + # register importance for in_channels + self.get_hparam("in_channels").register_importance(self._estimate_importance) + + def modify( + self, + *, + channels_ratio: Optional[tuple[float, ...]] = None, + channel_divisor: int = 1, + kernel_size: tuple[Union[int, tuple[int, ...]], ...] = (), + ): + """Modify the dynamic choices of the module according to provided keyword arguments. + + Args: + channels_ratio: The ratios of the desired number of out/in channels over original + number of out/in channels. + channel_divisor: The divisor of the out/in channels. + kernel_sizes: The desired kernel sizes. + """ + # modify both in_channels and out_channels + channels = ["in_channels", "out_channels"] + choices: set[float] + for channel in channels: + hp = self.get_hparam(channel) + if channels_ratio is not None: + assert isinstance(hp.original, int) + choices = {r * hp.original for r in channels_ratio} + else: + choices = set(hp.choices) # type: ignore[arg-type] + choices_rounded: set[int] = {int(make_divisible(c, channel_divisor)) for c in choices} + hp.choices = list(set(hp.choices) & choices_rounded | {hp.original}) + + # modify kernel size choices + hp_ks = self.get_hparam("kernel_size") + kernel_size = hp_ks.choices if kernel_size is None else kernel_size + ks_choices = {val2tuple(ks, self.ndim) for ks in kernel_size} + hp_ks.choices = list(set(hp_ks.choices) & ks_choices | {hp_ks.original}) + + +@DMRegistry.register( + { + nn.ConvTranspose1d: "nn.ConvTranspose1d", + nn.ConvTranspose2d: "nn.ConvTranspose2d", + nn.ConvTranspose3d: "nn.ConvTranspose3d", + } +) +class _DynamicConvTransposeNd(_DynamicConvNd): + @staticmethod + def _get_channel_order() -> list[str]: + return ["in_channels", "out_channels"] + + def _estimate_importance(self) -> TracedHp.Importance: + # for group > 1, we do not know how to handle it yet + if self.groups > 1: + return None + weight = self._parameters["weight"] # retrieve full weight tensor + c_in = weight.shape[0] + return torch.norm(torch.reshape(weight.detach(), (c_in, -1)), dim=1) diff --git a/modelopt/torch/nas/modules/linear.py b/modelopt/torch/nas/modules/linear.py new file mode 100644 index 0000000000..29aaa7c4a9 --- /dev/null +++ b/modelopt/torch/nas/modules/linear.py @@ -0,0 +1,79 @@ +# 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. + +"""Dynamic linear implementations based on torch.nn.modules.linear.""" + +from typing import Optional + +import torch +from torch import nn + +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.utils import make_divisible + +from ..registry import DMRegistry +from ..traced_hp import TracedHp +from .utils import get_sliced_tensor + +__all__ = ["_DynamicLinear"] + + +@DMRegistry.register({nn.Linear: "nn.Linear"}) +class _DynamicLinear(DynamicModule): + """An ``nn.Linear`` layer with dynamic hyperparams.""" + + @staticmethod + def _get_weight(mod: "_DynamicLinear", weight: torch.Tensor) -> torch.Tensor: + return get_sliced_tensor(mod, weight, "out_features", "in_features") + + @staticmethod + def _get_bias(mod: "_DynamicLinear", bias: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + return get_sliced_tensor(mod, bias, "out_features") + + def _estimate_importance(self) -> TracedHp.Importance: + return self._parameters["weight"].detach().norm(dim=0) + + def _setup(self): + # register hyperparameters + self._register_hparam("in_features", TracedHp(list(range(1, self.in_features + 1)))) + self._register_hparam("out_features", TracedHp(list(range(1, self.out_features + 1)))) + + # register dynamic attributes of the class + self._register_dynamic_attribute("weight", self._get_weight) + self._register_dynamic_attribute("bias", self._get_bias) + + # register importance for in_features + self.get_hparam("in_features").register_importance(self._estimate_importance) + + def modify( + self, *, features_ratio: Optional[tuple[float, ...]] = None, feature_divisor: int = 1 + ): + """Modify the dynamic choices of the module according to provided keyword arguments. + + Args: + features_ratio: The ratios of the desired number of output/input features over original + number of output/input features. + feature_divisor: The divisor of the output/input features. + """ + # modify both in_features and out_features + features = ["in_features", "out_features"] + for feature in features: + hp = self.get_hparam(feature) + if features_ratio is not None: + choices = {r * hp.original for r in features_ratio} + else: + choices = set(hp.choices) + choices = {int(make_divisible(c, feature_divisor)) for c in choices} + hp.choices = list(set(hp.choices) & choices | {hp.original}) diff --git a/modelopt/torch/nas/modules/norm.py b/modelopt/torch/nas/modules/norm.py new file mode 100644 index 0000000000..60b6611a25 --- /dev/null +++ b/modelopt/torch/nas/modules/norm.py @@ -0,0 +1,198 @@ +# 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. + +"""Dynamic norm implementations based on norm modules in torch.nn.modules.""" + +from typing import Optional, Sequence, Union + +import torch +from torch import nn + +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.utils import make_divisible + +from ..registry import DMRegistry +from ..traced_hp import TracedHp +from .utils import get_sliced_tensor + +__all__ = ["_DynamicBatchNorm", "_DynamicInstanceNorm", "_DynamicLayerNorm", "_DynamicGroupNorm"] + + +class _DynamicBatchInstance(DynamicModule): + """Dynamic base class for batch norm and instance norm layers. + + NOTE: Don't use this class for instance checks. Use _DynamicBatchNorm or _DynamicInstanceNorm + instead! + """ + + @staticmethod + def _cut_to_active_features(mod: "_DynamicBatchInstance", value: Optional[torch.Tensor]): + return get_sliced_tensor(mod, value, "num_features") + + def _setup(self): + # register hyperparameters + self._register_hparam("num_features", TracedHp(list(range(1, self.num_features + 1)))) + + # register dynamic attributes + dyn_attrs = ["running_mean", "running_var", "weight", "bias"] + for attr in dyn_attrs: + self._register_dynamic_attribute(attr, self._cut_to_active_features) + + def modify( + self, *, features_ratio: Optional[tuple[float, ...]] = None, feature_divisor: int = 1 + ): + """Modify the dynamic choices of the module according to provided keyword arguments. + + Args: + features_ratio: The ratios of the desired number of features over original number of + features. + feature_divisor: The divisor of the number of features. + """ + hp = self.get_hparam("num_features") + if features_ratio is not None: + choices = {r * hp.original for r in features_ratio} + else: + choices = set(hp.choices) + choices = {int(make_divisible(c, feature_divisor)) for c in choices} + hp.choices = list(set(hp.choices) & choices | {hp.original}) + + +@DMRegistry.register( + { + nn.BatchNorm1d: "nn.BatchNorm1d", + nn.BatchNorm2d: "nn.BatchNorm2d", + nn.BatchNorm3d: "nn.BatchNorm3d", + nn.SyncBatchNorm: "nn.SyncBatchNorm", + } +) +class _DynamicBatchNorm(_DynamicBatchInstance): + """Just syntactic sugar so we have a common base class for batch norm only.""" + + +@DMRegistry.register( + { + nn.InstanceNorm1d: "nn.InstanceNorm1d", + nn.InstanceNorm2d: "nn.InstanceNorm2d", + nn.InstanceNorm3d: "nn.InstanceNorm3d", + } +) +class _DynamicInstanceNorm(_DynamicBatchInstance): + """Just syntactic sugar so we have a common base class for instance norm only.""" + + +@DMRegistry.register({nn.LayerNorm: "nn.LayerNorm"}) +class _DynamicLayerNorm(DynamicModule): + """An ``nn.LayerNorm`` layer with dynamic hyperparams.""" + + @staticmethod + def _get_normalized_shape( + mod: "_DynamicLayerNorm", value: Sequence[Union[int, TracedHp]] + ) -> tuple: + return tuple(value[:-1]) + (mod.num_features,) + + @staticmethod + def _cut_to_active_features( + mod: "_DynamicLayerNorm", value: Optional[torch.Tensor] + ) -> Optional[torch.Tensor]: + if value is None: + return value + nf_slice = mod.get_hparam("num_features").active_slice + return value[..., nf_slice].contiguous() + + def _setup(self): + # construct normalized shape with Hparam as last dimension + normalized_shape = list(self.normalized_shape) + normalized_shape[-1] = TracedHp(list(range(1, normalized_shape[-1] + 1))) + + # register the hyperparameter with a new name + self._register_hparam("num_features", normalized_shape[-1]) + + # register dynamic attributes + dyn_attrs = ["weight", "bias"] + for attr in dyn_attrs: + self._register_dynamic_attribute(attr, self._cut_to_active_features) + + self._register_dynamic_attribute("normalized_shape", self._get_normalized_shape) + + def modify( + self, *, features_ratio: Optional[tuple[float, ...]] = None, feature_divisor: int = 1 + ): + """Modify the dynamic choices of the module according to provided keyword arguments. + + Args: + features_ratio: The ratios of the desired number of features over original number of + features. + feature_divisor: The divisor of the number of features. + """ + hp = self.get_hparam("num_features") + if features_ratio is not None: + choices = {r * hp.original for r in features_ratio} + else: + choices = set(hp.choices) + choices = {int(make_divisible(c, feature_divisor)) for c in choices} + hp.choices = list(set(hp.choices) & choices | {hp.original}) + + +@DMRegistry.register({nn.GroupNorm: "nn.GroupNorm"}) +class _DynamicGroupNorm(DynamicModule): + """An ``nn.GroupNorm`` layer with dynamic hyperparams.""" + + _group_size: int + + @staticmethod + def _get_num_groups(mod: "_DynamicGroupNorm", value: int) -> int: + return mod.num_channels // mod._group_size + + @staticmethod + def _cut_to_active_channels(mod: "_DynamicGroupNorm", value: Optional[torch.Tensor]): + return get_sliced_tensor(mod, value, "num_channels") + + def _setup(self): + # register num_channels as hyperparameter + group_size = self.num_channels // self.num_groups + choices = [ + c + for c in range(self.num_groups, self.num_channels + 1) + if c % self.num_groups == 0 and c % group_size == 0 + ] + self._register_hparam("num_channels", TracedHp(choices, original=self.num_channels)) + + # register num_groups as a dynamic attribute so group size is same + self._register_temp_attribute("_group_size", group_size) + self._register_dynamic_attribute("num_groups", self._get_num_groups) + + # register dynamic attributes + dyn_attrs = ["weight", "bias"] + for attr in dyn_attrs: + self._register_dynamic_attribute(attr, self._cut_to_active_channels) + + def modify( + self, *, channels_ratio: Optional[tuple[float, ...]] = None, channel_divisor: int = 1 + ): + """Modify the dynamic choices of the module according to provided keyword arguments. + + Args: + channels_ratio: The ratios of the desired number of out/in channels over original + number of out/in channels. + channel_divisor: The divisor of the out/in channels. + """ + hp = self.get_hparam("num_channels") + choices: list[int] + if channels_ratio is not None: + choices = {r * hp.original for r in channels_ratio} + else: + choices = set(hp.choices) + choices = {int(make_divisible(c, channel_divisor)) for c in choices} + hp.choices = list(set(hp.choices) & choices | {hp.original}) diff --git a/modelopt/torch/nas/modules/utils.py b/modelopt/torch/nas/modules/utils.py new file mode 100644 index 0000000000..f3c08ea2a5 --- /dev/null +++ b/modelopt/torch/nas/modules/utils.py @@ -0,0 +1,63 @@ +# 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. + +"""Internal module for utility functions.""" + +from typing import Optional + +import torch + +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.opt.hparam import Hparam + + +def get_sliced_tensor_by_slices( + tensor: Optional[torch.Tensor], + slices: list[Hparam.ActiveSlice], +) -> Optional[torch.Tensor]: + """Get the tensor based on the active slice.""" + if tensor is None: + return tensor + + # check if we can return the original tensor + if all( + isinstance(s, slice) and (s.stop is None or s.indices(s.stop) == (0, tensor.shape[i], 1)) + for i, s in enumerate(slices) + ): + return tensor + + # slice tensor with minimal number of index operations + tensor_sliced = tensor + for i, _ in enumerate(slices): + if sum(not isinstance(s, slice) for s in slices) < 2: + tensor_sliced = tensor_sliced[slices] + break + tensor_sliced = tensor_sliced[slices[: i + 1]] + slices[i] = slice(None) # replace with a vanilla slice ("[:]") for next slicing iteration + + # return sliced, contiguous tensor + return tensor_sliced.contiguous() + + +def get_sliced_tensor( + mod: DynamicModule, + tensor: Optional[torch.Tensor], + *hp_names: Optional[str], +) -> Optional[torch.Tensor]: + """Get the tensor based on the slices.""" + slices = [ + mod.get_hparam(hp_name).active_slice if hp_name else slice(None) for hp_name in hp_names + ] + return get_sliced_tensor_by_slices(tensor, slices) diff --git a/modelopt/torch/nas/plugins/__init__.py b/modelopt/torch/nas/plugins/__init__.py new file mode 100644 index 0000000000..2d28ba7f9f --- /dev/null +++ b/modelopt/torch/nas/plugins/__init__.py @@ -0,0 +1,30 @@ +# 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. + +"""Handles nas plugins for third-party modules.""" + +from modelopt.torch.utils import import_plugin + +from .torch import * + +with import_plugin("megatron"): + from modelopt.core.torch.nas.plugins.megatron import * + + +with import_plugin("transformer engine"): + from .transformer_engine import * + +with import_plugin("transformers"): + from .transformers import * diff --git a/modelopt/torch/nas/plugins/torch.py b/modelopt/torch/nas/plugins/torch.py new file mode 100644 index 0000000000..a79fba540b --- /dev/null +++ b/modelopt/torch/nas/plugins/torch.py @@ -0,0 +1,71 @@ +# 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. + +"""Pytorch plugins needed for NAS and dynamic modules.""" + +import torch.nn as nn + +from ..modules import _DynamicBatchNorm +from ..registry import DMRegistry + +__all__ = [] + + +def wrapped_convert_sync_batchnorm(cls: type[nn.Module], /, module: nn.Module, *args, **kwargs): + """Extend the original convert_sync_batchnorm to handle dynamic modules. + + This method ensures that _DynamicBatchNorm instances are correctly + converted into DynamicSyncBatchNorm instances. + + .. note:: + + We explicitly use ``del`` here following the original implementation! + """ + # handle vanilla case + if not isinstance(module, _DynamicBatchNorm): + return cls.old_convert_sync_batchnorm(module, *args, **kwargs) + + # maintain hparam objects + hparams = dict(module.named_hparams()) + + # set all hparams to max value and retain old active values + val_actives = {name: hp.active for name, hp in hparams.items()} + for hp in hparams.values(): + hp.active = hp.max + + # export current module + module = module.export() + + # convert exported DynamicBN to SyncBN using the original method + module = cls.old_convert_sync_batchnorm(module, *args, **kwargs) + + # ensure that we indeed got a SyncBN + assert isinstance(module, nn.SyncBatchNorm), f"Expected SyncBatchNorm, got {type(module)}!" + + # convert SyncBN to DynamicSyncBN + module = DMRegistry.convert(module) + + # re-use hparams from original dynamic BN to maintain consistency and re-assign active val + for hp_name, hp in hparams.items(): + setattr(module, hp_name, hp) + hp.active = val_actives[hp_name] + + # return the new module + return module + + +# hook into convert_sync_batchnorm +nn.SyncBatchNorm.old_convert_sync_batchnorm = nn.SyncBatchNorm.convert_sync_batchnorm +nn.SyncBatchNorm.convert_sync_batchnorm = classmethod(wrapped_convert_sync_batchnorm) diff --git a/modelopt/torch/nas/plugins/transformer_engine.py b/modelopt/torch/nas/plugins/transformer_engine.py new file mode 100644 index 0000000000..0f710097b0 --- /dev/null +++ b/modelopt/torch/nas/plugins/transformer_engine.py @@ -0,0 +1,42 @@ +# 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. + +"""Plugin to add NAS support for Transformer Engine modules.""" + +import transformer_engine as te + +from ..modules import _DynamicLayerNorm +from ..registry import DMRegistry +from ..traced_hp import TracedHp + +__all__ = ["_DynamicTENorm"] + + +@DMRegistry.register( + {te.pytorch.LayerNorm: "te.pytorch.LayerNorm", te.pytorch.RMSNorm: "te.pytorch.RMSNorm"} +) +class _DynamicTENorm(_DynamicLayerNorm): + """A ``te.pytorch.{Layer/RMS}Norm`` layer with dynamic hyperparams.""" + + def _setup(self): + hidden_size = self.weight.shape[-1] + + # register the hyperparameter with a new name + self._register_hparam("num_features", TracedHp(list(range(1, hidden_size + 1)))) + + # register dynamic attributes + self._register_dynamic_attribute("weight", self._cut_to_active_features) + if hasattr(self, "bias"): # Bias is not present in RMSNorm + self._register_dynamic_attribute("bias", self._cut_to_active_features) diff --git a/modelopt/torch/nas/plugins/transformers.py b/modelopt/torch/nas/plugins/transformers.py new file mode 100644 index 0000000000..2cb249e739 --- /dev/null +++ b/modelopt/torch/nas/plugins/transformers.py @@ -0,0 +1,207 @@ +# 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. + +"""Plugin to add NAS support for HuggingFace Transformers models.""" + +from typing import Callable, Optional, Sequence + +import torch +from torch import nn +from transformers.models.bert.modeling_bert import BertAttention, BertSelfAttention +from transformers.models.gptj.modeling_gptj import GPTJAttention + +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.opt.hparam import HPType +from modelopt.torch.trace import Symbol +from modelopt.torch.trace.plugins.transformers import SymAttentionHead +from modelopt.torch.utils import make_divisible + +from ..config import AutoNASConfig +from ..registry import DMRegistry +from ..traced_hp import TracedHp, TracedHpRegistry + +__all__ = ["_DynamicAttention", "_DynamicBertAttention", "_DynamicGPTJAttention"] + + +@TracedHpRegistry.register(SymAttentionHead) +class AttentionHeadTracedHp(TracedHp): + _hp_hidden_dim: TracedHp # QKV out_features + _hidden_dim_per_head: int + + @TracedHp.active.setter # type: ignore[attr-defined] + def active(self, val: Optional[HPType]): + super(AttentionHeadTracedHp, type(self)).active.fset(self, val) # type: ignore[attr-defined] + with self._hp_hidden_dim._force_configurable(): + self._hp_hidden_dim.active = self.active * self._hidden_dim_per_head + + @TracedHp.choices.setter # type: ignore[attr-defined] + def choices(self, val: Sequence[HPType]): + super(AttentionHeadTracedHp, type(self)).choices.fset(self, val) # type: ignore[attr-defined] + with self._hp_hidden_dim._force_configurable(): + self._hp_hidden_dim.choices = [c * self._hidden_dim_per_head for c in self.choices] + + def _get_importance(self) -> TracedHp.Importance: + imp_hidden_dim = self._hp_hidden_dim._get_importance() + assert imp_hidden_dim is not None + imp_heads = imp_hidden_dim.reshape(self.max, -1).sum(dim=1) + return imp_heads + + def _enforce_order(self, order: Optional[torch.Tensor] = None) -> None: + """Convert order of heads to order of hidden_dim.""" + if order is None: + order_hidden_dim = None + else: + h_to_d_slice = torch.arange(self._hp_hidden_dim.max, device=order.device).reshape( + self.max, -1 + ) + order_hidden_dim = h_to_d_slice[order].flatten() + + super()._enforce_order(order) + self._hp_hidden_dim._enforce_order(order_hidden_dim) + + def _resolve_dependencies( + self, sym: Symbol, get_hp: Callable[[Symbol], TracedHp] + ) -> dict[Symbol, TracedHp]: + # get first dependency which is the hidden dim hparam and make it not configurable + sym_hidden_dim = sym._dependencies.pop(0) + self._hp_hidden_dim = get_hp(sym_hidden_dim) + self._hp_hidden_dim._is_configurable = False + + # set up constrained choices hidden dim hp + self._hidden_dim_per_head = self._hp_hidden_dim.original // self.original + with self._hp_hidden_dim._force_configurable(): + self._hp_hidden_dim.choices = [c * self._hidden_dim_per_head for c in self.choices] + + # handle remaining dependencies to regular way + mapping = super()._resolve_dependencies(sym, get_hp) + mapping[sym_hidden_dim] = self._hp_hidden_dim + + return mapping + + +class _DynamicAttention(DynamicModule): + """A generic attention layer with dynamic hyperparams.""" + + def _setup(self): + # register num_attention_heads hyperparameter + hp_num_heads = TracedHp(list(range(1, self.num_attention_heads + 1))) + self._register_hparam("num_attention_heads", hp_num_heads) + + def configure_qkv_out(self, q_name: str, k_name: str, v_name: str, out_name: str) -> None: + """Utility to configure qkv and output projection matrices and their hparams.""" + # convert pytorch modules to dynamic modules in-place + q = DMRegistry.convert(self.get_submodule(q_name)) + k = DMRegistry.convert(self.get_submodule(k_name)) + v = DMRegistry.convert(self.get_submodule(v_name)) + out = DMRegistry.convert(self.get_submodule(out_name)) + + # link up embedding dim hparam manually and register with the attention module + hp_embed_dim = q.get_hparam("in_features") + self._register_hparam("embed_dim", hp_embed_dim) + k.in_features = hp_embed_dim + v.in_features = hp_embed_dim + out.out_features = hp_embed_dim + + # link up hidden dim hparam manually and register with the attention module + hp_hidden_dim = q.get_hparam("out_features") + self._register_hparam("hidden_dim", hp_hidden_dim) + k.out_features = hp_hidden_dim + v.out_features = hp_hidden_dim + out.in_features = hp_hidden_dim + + assert isinstance(out, nn.Linear) + hp_hidden_dim.register_importance(lambda: out._parameters["weight"].detach().norm(dim=0)) + + def modify( + self, *, n_heads_ratio: Optional[tuple[float, ...]] = None, n_heads_divisor: int = 1 + ) -> None: + # modify num_attention_heads + hp = self.get_hparam("num_attention_heads") + if n_heads_ratio is not None: + choices = {r * hp.original for r in n_heads_ratio} + else: + choices = set(hp.choices) + choices = {int(make_divisible(c, n_heads_divisor)) for c in choices} + hp.choices = list(set(hp.choices) & choices | {hp.original}) + + +# Provide BertSelfAttention as a parent class since we do not register to DMRegistry +class _DynamicBertSelfAttention(_DynamicAttention, BertSelfAttention): + """A ``BertSelfAttention`` layer with dynamic hyperparams.""" + + def _setup(self): + super()._setup() + self._register_dynamic_attribute( + "all_head_size", lambda mod, val: mod.num_attention_heads * mod.attention_head_size + ) + + def export(self) -> nn.Module: + self.query.export() + self.key.export() + self.value.export() + super().export() + return self + + +# TODO: can we just register BertSelfAttention and infer BertSelfOutput from tracing? +@DMRegistry.register({BertAttention: "hf.BertAttention"}) +class _DynamicBertAttention(_DynamicAttention): + """A ``BertAttention`` layer with dynamic hyperparams.""" + + def _setup(self): + # convert pytorch modules to dynamic modules + self.self = _DynamicBertSelfAttention.convert(self.self) + self._register_hparam("num_attention_heads", self.self.get_hparam("num_attention_heads")) + self.configure_qkv_out("self.query", "self.key", "self.value", "output.dense") + + # convert and link layer norm hparam to embed_dim + self.output.LayerNorm = DMRegistry.convert(self.output.LayerNorm) + self.output.LayerNorm.num_features = self.get_hparam("embed_dim") + + def export(self) -> nn.Module: + self.self.export() + self.output.dense.export() + self.output.LayerNorm.export() + super().export() + return self + + +@DMRegistry.register({GPTJAttention: "hf.GPTJAttention"}) +class _DynamicGPTJAttention(_DynamicAttention): + """A ``GPTJAttention`` layer with dynamic hyperparams.""" + + def _setup(self): + super()._setup() + self.configure_qkv_out("q_proj", "k_proj", "v_proj", "out_proj") + + def export(self) -> nn.Module: + self.q_proj.export() + self.k_proj.export() + self.v_proj.export() + self.out_proj.export() + super().export() + return self + + +def _n_heads_config(): + return {"n_heads_ratio": None, "n_heads_divisor": 1} + + +AutoNASConfig.register_default( + { + "hf.BertAttention": _n_heads_config(), + "hf.GPTJAttention": _n_heads_config(), + } +) diff --git a/modelopt/torch/nas/registry.py b/modelopt/torch/nas/registry.py new file mode 100644 index 0000000000..37cee47f17 --- /dev/null +++ b/modelopt/torch/nas/registry.py @@ -0,0 +1,23 @@ +# 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. + +"""A dynamic module registry that holds information about available DynamicModules for NAS.""" + +from modelopt.torch.opt.dynamic import _DMRegistryCls + +__all__ = ["DMRegistry"] + + +DMRegistry = _DMRegistryCls(prefix="Dynamic") # global instance for the registry diff --git a/modelopt/torch/nas/search_space.py b/modelopt/torch/nas/search_space.py new file mode 100644 index 0000000000..5985bfa902 --- /dev/null +++ b/modelopt/torch/nas/search_space.py @@ -0,0 +1,245 @@ +# 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. + +"""An extension of the DynamicSpace to handle searchable units and related utilities.""" + +import fnmatch +from typing import Any, Optional, Union + +import torch +import torch.nn as nn + +from modelopt.torch.opt import RulesDict +from modelopt.torch.opt.dynamic import DynamicModule, DynamicSpace, _DMRegistryCls +from modelopt.torch.trace import Symbol, SymMap, analyze_symbols +from modelopt.torch.utils import random + +from .registry import DMRegistry +from .traced_hp import TracedHp, TracedHpRegistry + +__all__ = ["generate_search_space", "SearchSpace"] + +SampleFunc = Union[random.FSample, dict[str, random.FSample]] + + +class SearchSpace(DynamicSpace): + """Stub for search space interface.""" + + def __init__(self, model: nn.Module, sym_map: Optional[SymMap] = None) -> None: + """Initialize the search space from the model and its associated symbolic map.""" + super().__init__(model) + self._sym_map = sym_map + if self._sym_map: + assert self._sym_map.model is self.model, "Model and sym_map's model must match." + + def _sym_to_hp(self, sym: Symbol) -> TracedHp: + assert isinstance(self._sym_map, SymMap), "Model's sym_map was not provided!" + for name, sym_ in self._sym_map.named_symbols(): + if sym is sym_: + mod, _, hp_name = name.rpartition(".") + mod = self.model.get_submodule(mod) + assert isinstance(mod, DynamicModule) + hp = mod.get_hparam(hp_name) + assert isinstance(hp, TracedHp), f"Expected TracedHp, got {type(hp)}!" + return hp + raise KeyError(f"Symbol {sym} not found in sym_map!") + + def _should_be_converted(self, mod: nn.Module) -> bool: + return self._sym_map is None or mod in self._sym_map + + def generate(self, rules: Optional[RulesDict] = None) -> None: + """Generate the search space from the model + symbolic map.""" + # sanity checks + assert isinstance(self._sym_map, SymMap), "Model's sym_map was not provided!" + + # 1. + 2. convert to dynamic module and apply rules + mods_converted = self.convert_to_dynamic(rules, DMRegistry) + + # 3. iterate through searchable/constant symbols and generate hparams for each + # dependency tree. + sym_to_hp_all: dict[Symbol, TracedHp] = {} + for name, sym in self._sym_map.named_symbols(searchable=True, constant=True): + mod_name, _, sym_name = name.rpartition(".") + mod = self.model.get_submodule(mod_name) + assert isinstance(mod, DynamicModule), f"Module {mod} must be a DynamicModule!" + + hp = TracedHpRegistry.initialize_from(sym, mod.get_hparam(sym_name)) + setattr(mod, sym_name, hp) + sym_to_hp = hp.resolve_dependencies(sym, self._sym_to_hp) + sym_to_hp_all.update(sym_to_hp) + + # 4. sanity check that we cover all symbols and hyperparameters + # NOTE: symbols and hparams of the non-leaf modules will not be included in the search space + symbols_all = dict(self._sym_map.named_symbols()) + assert set(sym_to_hp_all.keys()) == set(symbols_all.values()), "Missing symbols!" + + # 5. iterate through all dynamic modules and reassign hparams accordingly + for mod_name, mod in mods_converted.items(): + for sym_name, sym in self._sym_map[mod].items(): + # set new hp + setattr(mod, sym_name, sym_to_hp_all[sym]) + + def sample(self, sample_func: SampleFunc = random.choice) -> dict[str, Any]: + """Sample configurable hparams using the provided sample_func and return resulting config. + + Args: + model: A configurable model that contains one or more DynamicModule(s). + sample_func: A sampling function for hyperparameters. If a + dict is provided, the key indicates a Unix shell-style wildcard pattern which is + used to match the hyperparameter name. The default is random sampling corresponding + to ``{"*": random.choice}``. + + .. note:: + + Later matches take precedence over earlier ones when a dict with wilcards is + provided. + + Returns: + A dict of ``(parameter_name, choice)`` that specifies an active subnet. + """ + # wrap sample_func into a dict if it's not already + if not isinstance(sample_func, dict): + sample_func = {"*": sample_func} + + # sample configurable hparams + f_sample: Optional[random.FSample] + for hp_name, hp in self.named_hparams(configurable=True): + for pattern, f_sample in sample_func.items(): + if fnmatch.fnmatch(hp_name, pattern): + break + else: + f_sample = None + if f_sample: + hp.active = f_sample(hp.choices) + + return self.config() + + @torch.no_grad() + def sort_parameters(self, hps_to_sort: Optional[set[str]] = None) -> None: + """A graph propagation based parameter sorting algorithm. + + Args: + hps_to_sort: A set of hparam names to sort. If not provided or empty, all hparams will be sorted. + """ + # get config and set to max + config = self.config() + self.sample(sample_func=max) + + # sort configurable hparams + for name, hp in self.named_hparams(configurable=True): + hp_name = name.split(".")[-1] + if hps_to_sort and hp_name not in hps_to_sort: + continue + importance = hp.importance + # no sorting to be done (maybe because there is nothing to sort or it's unsortable) + if importance is None: + continue + # compute order from importance and enforce it + order = torch.argsort(importance, descending=True) + hp.enforce_order(order) + + # now that we have enforced an order we can force reassign all parameters/buffers! + for _, mod in self.named_dynamic_modules(): + mod.force_assign() + + # go back to old config + self.select(config) + + def export(self, dm_registry: _DMRegistryCls = DMRegistry) -> nn.Module: + """Export with default registry.""" + return super().export(dm_registry) + + def print_summary(self, skipped_hparams: list[str] = ["kernel_size"]) -> None: + """Print a summary of the search space.""" + print("\nSearch Space Summary:\n{:-^100}".format("")) + hp_visited = set() # Only highlight configurable hparams once + for name, hp in self.named_hparams(): + if not any(name.endswith(s) for s in skipped_hparams): + if len(hp.choices) > 8: + choices_to_print = hp.choices[:4] + ["..."] + hp.choices[-4:] + else: + choices_to_print = hp.choices + if hp.is_configurable and hp not in hp_visited: + hp_visited.add(hp) + prefix = "*" + else: + prefix = " " + print(f"{prefix} {name:80} {choices_to_print}") + + print("{:-^100}".format("")) + + +def generate_search_space(model: nn.Module, rules: Optional[RulesDict] = None) -> SearchSpace: + """Patch model with dynamic units and generate/return the search space. + + Args: + model: The model to patched for which we want to generate the search space. + rules: An optional dict specifying the rules used during the search space generation. + If not provided, the largest possible search space is generated. An example rule: + + .. code-block:: python + + rules = { + "nn.Sequential": {"min_depth": 2}, + "nn.Conv2d": { + "*": { + "channels_ratio": [1], + "channel_divisor": 16, + "kernel_size": [], + }, + "backbone.stages.[1-5]*.spatial_conv*": { + "channels_ratio": [0.334, 0.5, 0.667, 1], + }, + "det_blocks.*.spatial_conv*": { + "channels_ratio": [0.667, 1], + }, + }, + "nn.Linear": {}, # uses default arguments for Linear + "nn.BatchNorm2d": {}, + "nn.SyncBatchNorm": {}, + } + + .. note:: + + When rules are provided, the search space generation is restricted to the rules that + are provided. Other module types will not be included in the search space. + + .. note:: + + Rules with ``None``, e.g., ``{"nn.Sequential": None}``, will also cause the + conversion to be skipped just like when no rules are provided for the layer type. + + .. note:: + + Instead of specifying the module type as a string, the module type can also be + specified as a key directly, e.g., ``{nn.Sequential: None}``. + + Returns: + The patched model with dynamic units that describe the resulting search space. + + Generally, the search space generation follows the following outline of steps: + + 1. Analyze the model to identify the dynamic units that can be patched via modelopt.torch.trace. + 2. Patch the model with the dynamic units accordingly. + 3. Generate a consistent search space from the patched model by ensuring consistency between + dynamic units according to the symbolic map obtained from the model analysis. + + Upon return of the model, the search space can be analyzed and modified via the utilities + provided in ``modelopt.torch.dynamic.SearchSpace``. + """ + sym_map = analyze_symbols(model) + search_space = SearchSpace(model, sym_map) + search_space.generate(rules=rules) + return search_space diff --git a/modelopt/torch/nas/traced_hp.py b/modelopt/torch/nas/traced_hp.py new file mode 100644 index 0000000000..6ca72d4e39 --- /dev/null +++ b/modelopt/torch/nas/traced_hp.py @@ -0,0 +1,149 @@ +# 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. + +"""Module to specify the basic hyperparameter with tracing capabilities.""" + +from typing import Callable, Optional + +from modelopt.torch.opt.hparam import Hparam +from modelopt.torch.trace import Symbol + +__all__ = ["TracedHpRegistry", "TracedHp"] + +TracedHparamType = type["TracedHp"] + + +class TracedHpRegistry: + """A simple registry to keep track of different traced hp classes and their symbols.""" + + _registry: dict[type[Symbol], TracedHparamType] = {} # registered classes + + @classmethod + def register(cls, sym_cls: type[Symbol]) -> Callable[[TracedHparamType], TracedHparamType]: + """Use this to register a new traced hparam class for the provided symbol class. + + Usage: + + .. code-block:: python + + @TracedHpRegistry.register(MySymbol) + class MyHparam(TracedHp): + pass + """ + + def decorator(hp_class: TracedHparamType) -> TracedHparamType: + """Register hp_class with appropriate sym_class.""" + assert sym_cls not in cls._registry, f"{sym_cls} already registered!" + cls._registry[sym_cls] = hp_class + return hp_class + + return decorator + + @classmethod + def unregister(cls, sym_cls: type[Symbol]) -> None: + """Unregister a previously registered symbol class. + + It throws a KeyError if the hparam class is not registered. + """ + if sym_cls not in cls._registry: + raise KeyError(f"{sym_cls} is not registered!") + cls._registry.pop(sym_cls) + + @classmethod + def initialize_from(cls, sym: Symbol, hp: Hparam) -> "TracedHp": + """Initialize the sym-appropriate hparam from a vanilla hparam.""" + return cls._registry[type(sym)].initialize_from(hp) + + @classmethod + def get(cls, sym: Symbol) -> Optional[TracedHparamType]: + """Get Hparam type associated with symbol.""" + return cls._registry.get(type(sym)) + + +@TracedHpRegistry.register(Symbol) +class TracedHp(Hparam): + """A hparam that exhibits additional functionality required to handle tracing.""" + + @classmethod + def initialize_from(cls, hp: Hparam) -> "TracedHp": + """Initialize a new hparam from an existing vanilla Trace.""" + # sanity checks + assert type(hp) is TracedHp, f"Can only initialize from an {TracedHp} object." + # relegate implementation to child class + return cls._initialize_from(hp) + + @classmethod + def _initialize_from(cls, hp: Hparam) -> "TracedHp": + """Initialize a new hparam from an existing vanilla Hparam.""" + hp.__class__ = cls + assert isinstance(hp, cls) + return hp + + def resolve_dependencies( + self, sym: Symbol, get_hp: Callable[[Symbol], "TracedHp"] + ) -> dict[Symbol, "TracedHp"]: + """Resolve dependencies of the hparam via symbolic map. + + This method iterates through the dependency map described in sym and generates an + appropriate hparam based on the currently assigned hparam and the parent symbol. + + Args: + sym: The symbol associated with self for which we want to resolve dependencies. + get_hp: A function that returns the hparam associated with a symbol. + + Returns: + A mapping describing the hparam that should be associated with each symbol. + """ + # dependency resolution only works for searchable symbols and must be associated with self + assert sym.is_searchable or sym.is_constant, f"Symbol {sym} must be searchable or constant!" + assert get_hp(sym) is self, f"Symbol {sym} must be associated with self to resolve!" + assert TracedHpRegistry.get(sym) is type(self), "Symbol and Hparam type must match!" + + # relegate implementation to child class + return self._resolve_dependencies(sym, get_hp) + + def _resolve_dependencies( + self, sym: Symbol, get_hp: Callable[[Symbol], "TracedHp"] + ) -> dict[Symbol, "TracedHp"]: + # check sortability of sym + if not sym.is_sortable: + self._importance_estimators = None + + # process dependencies and generate hparams + to_process = set(sym._dependencies) + processed: set[Symbol] = set() + mapping: dict[Symbol, TracedHp] = {sym: self} + + while to_process: + # get a new symbol + sym_dep = to_process.pop() + processed.add(sym_dep) + + # this should never be any other symbol + assert type(sym_dep) is Symbol, f"Unexpected type {type(sym_dep)} for {sym_dep}!" + + # merge hparam into self and store mapping + self &= get_hp(sym_dep) + mapping[sym_dep] = self + + # add dependencies of sym_dep to to_process + to_process |= set(sym_dep._dependencies) - processed + + # check constant case at the end + if sym.is_constant: + self.active = self.original + self.choices = [self.original] + + return mapping diff --git a/modelopt/torch/nas/utils.py b/modelopt/torch/nas/utils.py new file mode 100644 index 0000000000..b94054cf06 --- /dev/null +++ b/modelopt/torch/nas/utils.py @@ -0,0 +1,406 @@ +# 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. + +"""Utility functions for prune-related and search-space related tasks. + +.. note:: + + Generally, methods in the :meth:`modelopt.torch.nas` module should use these utility + functions directly instead of accessing the ``SearchSpace`` class. This is to ensure that + potentially required pre- and post-processing operations are performed correctly. +""" + +import warnings +from contextlib import contextmanager +from typing import Any, Callable, Iterator, Optional, Union + +import torch +import torch.nn as nn +import torchprofile.profile as profile +from torch.autograd.grad_mode import _DecoratorContextManager + +from modelopt.torch.utils import ( + get_module_device, + is_parallel, + num2hrb, + random, + standardize_model_args, + unwrap_model, +) + +from .search_space import SampleFunc, SearchSpace + +# only certain categories of utils via __all__ +__all__ = [ + # search/profile related utils + "inference_flops", + "print_search_space_summary", + # sample/select-related utils + "get_subnet_config", + "sample", + "select", + # auto-mode related utils + "enable_modelopt_patches", + "is_modelopt_patches_enabled", + "no_modelopt_patches", + "set_modelopt_patches_enabled", + # other utils that are used frequently + "replace_forward", +] + +# we have two different numbers here since during training it might take longer to stabilize +MODELOPT_QUEUE_MAXLEN = 50 # indicates length of modelopt data queue for BN calib +MODELOPT_BN_CALIB_ITERS = ( + 100 # indicates # iters in train mode 'til we trust BN stats without calib +) + + +@contextmanager +def batch_norm_ignored_flops(): + handlers_bck = profile.handlers + handlers_new = [] + for op_names, op in profile.handlers: + if "aten::batch_norm" in op_names: + op_names_new = [] + for op_name in op_names: + if op_name != "aten::batch_norm": + op_names_new.append(op_name) + handlers_new.append((tuple(op_names_new), op)) + else: + handlers_new.append((op_names, op)) + profile.handlers = tuple(handlers_new) + yield + profile.handlers = handlers_bck + + +def inference_flops( + network: nn.Module, + dummy_input: Optional[Union[Any, tuple]] = None, + data_shape: Optional[tuple] = None, + unit: float = 1e6, + return_str: bool = False, +) -> Union[float, str]: + """Get the inference FLOPs of a PyTorch model. + + Args: + network: The PyTorch model. + args: The dummy input as defined in + :meth:`mtn.convert() `. + data_shape: The shape of the dummy input if the dummy input is a single tensor. If provided, + ``args`` must be ``None``. + unit: The unit to return the number of parameters in. Default is 1e6 (million). + return_str: Whether to return the number of FLOPs as a string. + + Returns: + The number of inference FLOPs in the given unit as either string or float. + """ + if is_parallel(network): + network = network.module + if data_shape is not None and dummy_input is not None: + raise ValueError("Please provide either data_shape or args tuple.") + if data_shape is not None: + dummy_input = (torch.zeros(data_shape, device=get_module_device(network)),) + dummy_input = standardize_model_args(network, dummy_input) + is_training = network.training + network.eval() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with batch_norm_ignored_flops(): + flops = profile.profile_macs(network, args=dummy_input) + network.train(is_training) + if return_str: + return num2hrb(flops) + return flops / unit + + +class _SearchSpaceUnwrapped(SearchSpace): + """A wrapper for the SearchSpace class to handle unwrapping of model wrappers like DDP. + + This is useful to ensure that configurations are valid for both vanilla models and wrapped + models, see :meth:`unwrap_models` for supported wrappers. + """ + + def __init__(self, model: nn.Module) -> None: + super().__init__(unwrap_model(model)) + + +def sample(model: nn.Module, sample_func: SampleFunc = random.choice) -> dict[str, Any]: + """Sample searchable hparams using the provided sample_func and return resulting config. + + Args: + model: A searchable model that contains one or more DynamicModule(s). + sample_func: A sampling function for hyperparameters. Default: random sampling. + + Returns: + A dict of ``(parameter_name, choice)`` that specifies an active subnet. + """ + _reset_before_sample(model) + return _SearchSpaceUnwrapped(model).sample(sample_func) + + +def select(model: nn.Module, config: dict[str, Any], strict: bool = True) -> None: + """Select the sub-net according to the provided config dict. + + Args: + model: A model that contains DynamicModule(s). + config: Config of the target subnet as returned by + :meth:`mtn.config()` + and :meth:`mtn.search()`. + strict: Raise an error when the config does not contain all necessary keys. + """ + # select new sub-net configuration + _reset_before_sample(model) + _SearchSpaceUnwrapped(model).select(config, strict) + + +def get_subnet_config(model: nn.Module, configurable: Optional[bool] = None) -> dict[str, Any]: + """Return the config dict of all hyperparameters. + + Args: + model: A model that contains DynamicModule(s). + configurable: None -> all hps, True -> configurable hps without duplicates + + Returns: + A dict of ``(parameter_name, choice)`` that specifies an active subnet. + """ + return _SearchSpaceUnwrapped(model).config(configurable) + + +def _reset_before_sample(model: nn.Module): + """Call pre-sample reset hook from patch manager.""" + from ._patch import PatchManager + + if PatchManager.is_patched(model): + PatchManager.get_manager(model).reset_before_sample() + + +def sort_parameters(model: nn.Module, hps_to_sort: Optional[set[str]] = None) -> None: + """Sort the parameters of the model according to the stored importances. + + Args: + hps_to_sort: A set of hparam names to sort. If not provided or empty, all hparams will be sorted. + """ + _SearchSpaceUnwrapped(model).sort_parameters(hps_to_sort) + + +def print_search_space_summary( + model: nn.Module, skipped_hparams: list[str] = ["kernel_size"] +) -> None: + """Print the search space summary. + + Args: + model: A model that contains DynamicModule(s). + """ + _SearchSpaceUnwrapped(model).print_summary(skipped_hparams) + + +class _ModeloptOpsState: + """Global flag to enable/disable modelopt patches - defaults to True.""" + + _instance: Optional["_ModeloptOpsState"] = None + _auto_enabled = True + + def __new__(cls) -> "_ModeloptOpsState": + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def set_modelopt_mode_enabled(self, enabled: bool): + self._auto_enabled = enabled + + @property + def enabled(self) -> bool: + return self._auto_enabled + + +def is_modelopt_patches_enabled() -> bool: + """Check if modelopt patches for model are enabled.""" + return _ModeloptOpsState().enabled + + +class no_modelopt_patches(_DecoratorContextManager): # noqa: N801 + """Context manager to disable ``modelopt`` patches to the model. + + Disabling ``modelopt`` patches is useful when you want to use the model's original behavior + For example, you can use this to perform a forward pass without NAS operations. + + It can also be used as a decorator (make sure to instantiate with parenthesis). + + For example: + + .. code-block:: python + + modelopt_model.train() + modelopt_model(inputs) # architecture changes + + with mtn.no_modelopt(): + modelopt_model(inputs) # architecture does not change + + + @mtn.no_modelopt() + def forward(model, inputs): + return model(inputs) + + + forward(modelopt_model, inputs) # architecture does not change + """ + + def __init__(self): + """Constructor.""" + self.prev = False + + def __enter__(self): + self.prev = is_modelopt_patches_enabled() + set_modelopt_patches_enabled(False) + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any): + set_modelopt_patches_enabled(self.prev) + + +class enable_modelopt_patches(_DecoratorContextManager): # noqa: N801 + """Context manager to enable ``modelopt`` patches such as those for autonas/fastnas. + + It can also be used as a decorator (make sure to instantiate with parenthesis). + + For example: + + .. code-block:: python + + modelopt_model.train() + modelopt_model(inputs) # architecture changes + + with mtn.no_modelopt(): + with mtn.enable_modelopt(): + modelopt_model(inputs) # architecture changes + + + @mtn.enable_modelopt() + def forward(model, inputs): + return model(inputs) + + + with mtn.no_modelopt(): + forward(modelopt_model, inputs) # architecture changes because of decorator on forward + """ + + def __init__(self): + """Constructor.""" + self.prev = True + + def __enter__(self): + self.prev = is_modelopt_patches_enabled() + set_modelopt_patches_enabled(True) + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any): + set_modelopt_patches_enabled(self.prev) + + +class set_modelopt_patches_enabled(_DecoratorContextManager): # noqa: N801 + """Context manager that sets patches to on or off. + + It can be used as context manager or as a function. If used as function, operations + are disabled globally (thread local). + + Args: + enabled: whether to enable (``True``) or disable (``False``) patched methods. + + For example: + + .. code-block:: python + + modelopt_model.train() + modelopt_model(inputs) # architecture changes + + mtn.set_modelopt_enabled(False) + modelopt_model(inputs) # architecture does not change + + with mtn.set_modelopt_enabled(True): + modelopt_model(inputs) # architecture changes + + modelopt_model(inputs) # architecture does not change + """ + + def __init__(self, enabled: bool): + """Constructor.""" + self.prev = is_modelopt_patches_enabled() + _ModeloptOpsState().set_modelopt_mode_enabled(enabled) + self.enabled = enabled + + def __enter__(self): + pass + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any): + _ModeloptOpsState().set_modelopt_mode_enabled(self.prev) + + def clone(self): + """Clone the context manager.""" + return self.__class__(self.enabled) + + +@contextmanager +def replace_forward(model: nn.Module, new_forward: Callable) -> Iterator[None]: + """Context manager to temporarily replace the forward method of the underlying type of a model. + + The original forward function is temporarily accessible via ``model.forward_original``. + + Args: + model: The model whose type's forward method is to be temporarily replaced. + new_forward: The new forward method. The forward method should either be a bound method to + the model instance or take the model (``self``) as the first argument. + + For example: + + .. code-block:: python + + fake_forward = lambda _: None + + with replace_forward(model, fake_forward): + out = model(inputs) # this output is None + + out_original = model(inputs) # this output is the original output + """ + # This is functional context manager using the @contextmanager decorator: + # https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager + type(model).forward_original = type(model).forward + type(model).forward = getattr(new_forward, "__func__", new_forward) + try: + yield + finally: + type(model).forward = type(model).forward_original + del type(model).forward_original + + +@contextmanager +def sample_and_reset(model: nn.Module, sample_func: SampleFunc = random.choice) -> Iterator[None]: + """Context manager to temporarily sample a subnet based on a given sample_func. + + Args: + model: A searchable model that contains one or more DynamicModule(s). + sample_func: A sampling function for hyperparameters. Default: random sampling. + + For example: + + .. code-block:: python + + with sample_and_reset(model, random.original): + out = model(inputs) # uses the original subnet + + """ + config = get_subnet_config(model) + try: + sample(model, sample_func) + yield + finally: + select(model, config) diff --git a/modelopt/torch/opt/__init__.py b/modelopt/torch/opt/__init__.py new file mode 100644 index 0000000000..2649082fc7 --- /dev/null +++ b/modelopt/torch/opt/__init__.py @@ -0,0 +1,43 @@ +# 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. + +"""Module for general-purpose model optimization infrastructure. + +``modelopt.torch.opt`` contains tooling to: + +* ingest a user-provided model and set it up for optimization; +* define and implement search and optimization procedures; +* export a model back to a regular pytorch model after optimization; +* save, restore, and manage checkpoints from which the model modifications can be resumed. + +Please refer to each individual sub-module to learn more about the various concepts wihin +``modelopt.torch.opt`` and how to use them to implement a model optimization algorithm. +""" + +try: + from . import _hooks +except ImportError: + pass + +try: + from .plugins.huggingface import * +except ImportError: + pass + +from . import plugins, utils +from .config import * +from .conversion import * +from .mode import * +from .searcher import * diff --git a/modelopt/torch/opt/_hooks.py b/modelopt/torch/opt/_hooks.py new file mode 100644 index 0000000000..0b52a7f631 --- /dev/null +++ b/modelopt/torch/opt/_hooks.py @@ -0,0 +1,66 @@ +# 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. + +"""Hooks for Pytorch to ensure compatibility with dynamic modules.""" + +from contextlib import ExitStack, nullcontext +from typing import Union + +import torch +import torch.nn as nn + +# Older versions have torch.distributed.fsdp.flat_param module but without `_safe_setattr_tensor_or_param` +from torch.distributed.fsdp import _flat_param +from torch.distributed.fsdp._flat_param import FlatParamHandle + +from .dynamic import DynamicModule + + +def _safe_setattr_tensor_or_param_with_dm_check( + module: nn.Module, param_name: str, tensor_or_param: Union[torch.Tensor, nn.Parameter] +): + """A batched version of _safe_setattr_tensor_or_param ensuring compatibility with DMs. + + This function is called during the creation of the flat param in the FSDP module. We intercept + this call to ensure that dynamic attributes are correctly preserved without interfering with + FSDP's management, retrieval, and updates of parameters that might also be dynamic attributes. + """ + with module.reset_dynamic_attributes() if isinstance(module, DynamicModule) else nullcontext(): + return _flat_param._safe_setattr_tensor_or_param_original( + module, param_name, tensor_or_param + ) + + +# set the new function as the replacement +_flat_param._safe_setattr_tensor_or_param_original = _flat_param._safe_setattr_tensor_or_param +_flat_param._safe_setattr_tensor_or_param = _safe_setattr_tensor_or_param_with_dm_check + + +def _writeback_orig_param(self: FlatParamHandle): + """A batched version of FlatParamHandle._writeback_orig_params ensuring compatibility with DMs. + + This function is called when we use_orig_params=True. We intercept this call to ensure that + dynamic attributes are correctly preserved without interfering with FSDP's management, + retrieval, and updates of parameters that might also be dynamic attributes. + """ + with ExitStack() as stack: + for _, module, _ in self.flat_param._param_infos: + if isinstance(module, DynamicModule): + stack.enter_context(module.reset_dynamic_attributes()) + return self._writeback_orig_params_original() + + +FlatParamHandle._writeback_orig_params_original = FlatParamHandle._writeback_orig_params +FlatParamHandle._writeback_orig_params = _writeback_orig_param diff --git a/modelopt/torch/opt/config.py b/modelopt/torch/opt/config.py new file mode 100644 index 0000000000..2d1b4eedd6 --- /dev/null +++ b/modelopt/torch/opt/config.py @@ -0,0 +1,396 @@ +# 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. + +"""Modelopt's pydantic BaseModel used for any type of configuration in algorithms and mode.""" + +import fnmatch +import json +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ItemsView, + Iterator, + KeysView, + Optional, + Union, + ValuesView, +) + +import pydantic +from pydantic import ( + BaseModel, + Field, + TypeAdapter, + ValidationError, + ValidationInfo, + field_validator, + model_validator, +) +from pydantic_core import PydanticUndefined + +if TYPE_CHECKING: # Only executed during type checking, e.g. by mypy + # typing.TypeAlias is only available for python 3.10+. + from typing_extensions import TypeAlias + +# A simple type alias for a config dictionary that is used as input to initialize a ModeloptBaseConfig. +ConfigDict = dict[str, Any] # config dict for one mode + +# These are a static type annotations for rules that can be used to annotate functions that accept +# various rule dictionaries. Proper type checking is done by the respective rule config class, see +# below. +# NOTE: the ``WrappedRule`` and ``Rule`` types are further refined via +# ``ModeloptBaseRule.get_rule_type(wrapped_only=True)`` and +# ``ModeloptBaseRule.get_rule_type(wrapped_only=False)``, respectively. +SimpleRule = dict[str, Any] # a simple rule type +WrappedRule = dict[str, Optional[SimpleRule]] # a wrapped rule type with glob patterns +Rule = Union[Optional[SimpleRule], WrappedRule] # a rule type that can be simple, wrapped, or None +RulesDict = dict[str, Rule] # a dictionaries containing rules for different types + + +def ModeloptField(default: Any = PydanticUndefined, **kwargs): # noqa: N802 + """A pydantic.Field that enforces setting a default value.""" + assert default is not PydanticUndefined, "A default value must be set for ModeloptField." + return Field(default=default, **kwargs) + + +# TODO: expand config classes to searcher + + +class ModeloptBaseConfig(BaseModel): + """Our config base class for mode configuration. + + The base class extends the capabilities of pydantic's BaseModel to provide additional methods + and properties for easier access and manipulation of the configuration. + """ + + model_config = pydantic.ConfigDict(extra="forbid", validate_assignment=True) + + def model_dump(self, **kwargs): + """Dump the config to a dictionary with aliases and no warnings by default.""" + kwargs = {"by_alias": True, "warnings": False, **kwargs} + return super().model_dump(**kwargs) + + def model_dump_json(self, **kwargs): + """Dump the config to a json with aliases and no warnings by default.""" + kwargs = {"by_alias": True, "warnings": False, **kwargs} + return super().model_dump_json(**kwargs) + + @property + def _iterable_model_extra(self) -> dict[str, Any]: + """Return model_extra or empty dict if None so we can iterate over it.""" + return self.model_extra or {} + + def get_field_name_from_key(self, key: str) -> str: + """Get the field name from the given key (can be name or alias of field).""" + assert isinstance(key, str), f"key must be a string, got {type(key)}" + + if key in self.model_fields or key in self._iterable_model_extra: + return key + else: + for name, field_info in self.model_fields.items(): + if field_info.alias == key: + return name + raise AttributeError(f"Key {key} not found in the config.") + + def __contains__(self, key: str) -> bool: + """Check if the given key is present in the config by its actual name or alias.""" + try: + self.get_field_name_from_key(key) + return True + except AttributeError: + return False + + def __getitem__(self, key: str) -> Any: + """Get the value for the given key (can be name or alias of field).""" + return getattr(self, self.get_field_name_from_key(key)) + + def __setitem__(self, key: str, value: Any) -> None: + """Set the value for the given key (can be name or alias of field).""" + setattr(self, self.get_field_name_from_key(key), value) + + def get(self, key: str, default: Any = None) -> Any: + """Get the value for the given key (can be name or alias) or default if not found.""" + try: + return self[key] + except AttributeError: + return default + + def __len__(self) -> int: + """Return the length of the config.""" + return len(self.model_fields) + len(self._iterable_model_extra) + + def __iter__(self) -> Iterator[str]: + """Iterate over aliases (or name if alias is not defined) of fields.""" + for field_name, field_info in self.model_fields.items(): + yield field_info.alias or field_name + for key in self._iterable_model_extra: + yield key + + def _get_kv_dict(self) -> dict[str, Any]: + """Return a dictionary with keys as aliases if possible.""" + return {k: self[k] for k in self} + + def keys(self) -> KeysView[str]: + """Return the keys (aliases prioritized over names) of the config.""" + return self._get_kv_dict().keys() + + def values(self) -> ValuesView[Any]: + """Return the values of the config.""" + return self._get_kv_dict().values() + + def items(self) -> ItemsView[str, Any]: + """Return the items of the config with keys as aliases if possible.""" + return self._get_kv_dict().items() + + def update(self, config: ConfigDict) -> None: + """Update the config with the given config dictionary.""" + for key, value in config.items(): + self[key] = value + + +class ModeloptBaseRule(ModeloptBaseConfig): + """Our base config class for rule-based config classes. + + Rules are what governs the configuration for modifying dynamic module classes. + """ + + @classmethod + def get_rule_type(cls, wrapped_only: bool = False) -> "TypeAlias": + """Get the rule type for the given ModeloptBaseConfig.""" + optional_rule_type: "TypeAlias" = Optional[cls] # type: ignore[valid-type] + rule_type = dict[str, optional_rule_type] + if wrapped_only: + return rule_type + return Union[optional_rule_type, rule_type] + + @classmethod + def validate_rule(cls, rule: Rule) -> WrappedRule: + """Validate a rule with the current cls rule. + + We will check the full rule type (wrapped and unwrapped) and then return the wrapped type. + """ + # validate the rule with full rule tupe + TypeAdapter(cls.get_rule_type()).validate_python(rule) + + # we want to reduce it to the wrapped type + wrapped_rule: WrappedRule + if rule is None: + wrapped_rule = {"*": rule} + else: + try: + # if this call doesn't fail, it means that field is a valid rule_tpye and we + # need to wrap it. + cls.model_validate(rule) + wrapped_rule = {"*": rule} + except ValidationError: + wrapped_rule = rule + + # now check wrapped type only + TypeAdapter(cls.get_rule_type(wrapped_only=True)).validate_python(wrapped_rule) + + return wrapped_rule + + @classmethod + def customize_rule(cls, rule: Rule, key: str) -> Optional[SimpleRule]: + """Construct custom rule according to the provided key which is matched.""" + # validate rule first + rule = cls.validate_rule(rule) + + # create custom rule now + rule_custom: Optional[SimpleRule] = None + for pattern, subrule in rule.items(): + if fnmatch.fnmatch(key, pattern): + if rule_custom is None or subrule is None: + rule_custom = subrule + else: + rule_custom = {**rule_custom, **subrule} + return rule_custom + + +class ModeloptBaseRuleConfig(ModeloptBaseConfig): + """Our config base class for mode configuration that are purely made from rules. + + The base class extends the capabilities of pydantic's BaseModel to provide additional methods + and properties for easier access and manipulation of the configuration. + """ + + model_config = pydantic.ConfigDict(extra="allow") + + @classmethod + def __init_subclass__(cls, *args, registry, **kwargs): + cls._registry = registry + extra_default: dict[str, WrappedRule] = {} + cls._extra_default = extra_default + super().__init_subclass__(*args, **kwargs) + + @classmethod + def register_default(cls, extra_default: dict[str, WrappedRule]) -> None: + """Register a new default value for the given key.""" + # we don't allow overwriting existing field defaults + all_keys = cls.model_fields.keys() | {v.alias or k for k, v in cls.model_fields.items()} + overlap = all_keys & extra_default.keys() + assert not overlap, f"Updating default values for regular fields not allowed: {overlap}" + + # iterate through provided defaults + for k, v in extra_default.items(): + # run the rule validator and store + cls._extra_default[k] = cls._registry.get_rule_class(k).validate_rule(v) + + @classmethod + def unregister_default(cls, key: str) -> None: + """Unregister the default value for the given key.""" + cls._extra_default.pop(key) + + @model_validator(mode="after") + def _check_for_extra_fields(self): + """Check for extra fields and either apply their field validator or raise an error.""" + _check_field = "_updating_extra_fields" + if hasattr(self, _check_field): + # avoid recursion + return self + setattr(self, _check_field, True) + for k, v in {**self._extra_default, **self._iterable_model_extra}.items(): + # try validating and updating the field from registry + setattr(self, k, self._registry.get_rule_class(k).validate_rule(v)) + delattr(self, _check_field) + return self + + +def _get_default_description( + prefix: str, alias: str, rule_cls: type[ModeloptBaseRule], default: Rule +) -> str: + prefix = prefix.lower() + + newline_with_indent = "\n" + " " * 4 + + default_config = newline_with_indent.join(json.dumps(default, indent=2).splitlines()) + default_config = " " * 4 + default_config + return f"""Configuration for {prefix} {alias} module. + +If the ``"{alias}"`` key is not specified, the default configuration (shown in JSON) will be used: + +.. code-block:: json + +{default_config} + +To deactivate any {prefix} {alias} module, use ``None`` instead of providing a dictionary ``{{}}``. + +To specify layer-specific configurations, you can specify a config for each submodule with the key +specifying a glob pattern that matches the submodule name. For example, to convert to a {prefix} +module for all ``{alias}`` layers except for those in the ``"lm_head"`` submodule use: + +.. code-block:: python + + {{ + "*": {{...}}, + "*lm_head*": None, + }} + +Note that glob expressions are processed sequentially in the order they are specified. Later keys in +the config will overwrite earlier keys if they match the same submodule name. + +If you want to specify the same configuration for all submodules, you can provide an unnested +dictionary as well: + +.. code-block:: python + + {{...}} + +which is short for + +.. code-block:: python + + {{ + "*": {{...}}, + }} +""" + + +def _get_field_validator(alias: str) -> Callable: + def _validate_rule_with_correct_signature(cls, field: Any, info: ValidationInfo) -> Any: + return cls._registry.get_rule_class(alias).validate_rule(field) + + return field_validator(_get_field_name(alias))(_validate_rule_with_correct_signature) + + +def _get_field_name(alias: str) -> str: + return alias.replace(".", "_").lower() + + +# TODO: ideally we annotate the correct type for registry but this would require a refactor since +# we don't wanna introduce the `modelopt.torch.dynamic` dependency here. Leaving this to future work. +def get_kwargs_for_create_model_with_rules( + registry: Any, default_rules: RulesDict, doc: str +) -> dict[str, Any]: + """Generate the kwargs for ``pydantic.create_model`` to auto-generate a rule config class. + + Args: + registry: The dynamic module registry that contains all relevant dynamic modules. + rule_fields: The fields that the rule-based config class should have. + doc: The docstring for the rule-based config class. + + A rule-based config class is a config class that purely consists of fields that pertain to + rules. We can procedurally generate these rule config classes by using + + .. code-block:: python + + from pydantic import create_model + + MyRuleConfigs = create_model( + "MyRuleConfigs", **get_create_model_kwargs_for_rule_model(registry, rule_fields) + ) + + For more info and example usage, you can take a look at + :meth:`SparseMagnitudeConfig`. + + .. note:: + + We have this convenience function in place since autodocs only get generated when + ``create_model`` is *explicitly* called in the respective config file. So this function is a + workaround to at least lower the burden of correctly calling ``create_model``. + + + """ + # generate fields + field_specs = {} + for alias, default in default_rules.items(): + rule_class: type[ModeloptBaseRule] = registry.get_rule_class(alias) + default_validated = rule_class.validate_rule(default) + field_specs[_get_field_name(alias)] = ( + rule_class.get_rule_type(), + ModeloptField( + alias=alias, + default=default_validated, + title=f"{alias} config", + description=_get_default_description( + registry.prefix, alias, rule_class, default_validated + ), + ), + ) + + # generate validators + field_validators = { + f"_validate_for_{_get_field_name(alias)}": _get_field_validator(alias) + for alias in default_rules + } + + return { + "__base__": ModeloptBaseRuleConfig, + "__validators__": field_validators, + "__doc__": doc, + "__cls_kwargs__": {"registry": registry}, + **field_specs, + } diff --git a/modelopt/torch/opt/conversion.py b/modelopt/torch/opt/conversion.py new file mode 100644 index 0000000000..3ffc25d279 --- /dev/null +++ b/modelopt/torch/opt/conversion.py @@ -0,0 +1,588 @@ +# 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. + +"""Module to handle model converting and restoring for optimization methods. + +When applying a model optimization algorithm, we usually need to modify the model in each step +(mode) of the algorithm. This module provides the state manager, which is a standardized interface +(class) to record and store state information in the model. + +Op top of the state manager, this module provides utilities to save a history of these modifications +("modelopt state dict") and restoring a unmodified model to the state indicated in the state dict. +""" + +import copy +import os +import warnings +from collections import deque +from typing import Any, BinaryIO, Deque, Iterator, Optional, Union + +import torch +import torch.nn as nn + +from modelopt import __version__ +from modelopt.torch.utils import ModelLike, init_model_from_model_like, unwrap_model + +from .config import ConfigDict, ModeloptBaseConfig +from .mode import ( + MetadataDict, + ModeLike, + ModeState, + ModeType, + _ModeDescriptor, + _ModeRegistryCls, + get_mode_config, +) + +__all__ = [ + "ModeloptStateManager", + "apply_mode", + "modelopt_state", + "save", + "restore_from_modelopt_state", + "restore", +] + +ModeloptStateList = list[tuple[str, ModeState]] # state data structure for multiple modes + + +class ModeloptStateManager: + """A class to handle the modelopt state stored for each mode correspondig to a task/mode.""" + + _state_key = "_modelopt_state" + _state_version_key = "_modelopt_state_version" + + def __init__(self, model: Optional[nn.Module] = None, init_state: bool = False) -> None: + """Initialize state manager. + + Args: + model: Module that has modelopt_state stored. If None, a fake module is created to store + any state that might be added with the manager. + init_state: Whether to initialize the modelopt state for the model if it does not exist. + """ + # just assume fake module for easy implementation below + if not model: + model = nn.Module() + init_state = True # always initialize fake module + + # initialize modelopt state if desired. Note that by default we don't do that to avoid + # accidentally modifying a user-provided model. + if init_state: + assert not hasattr(model, self._state_key), "Model already has modelopt state!" + setattr(model, self._state_key, []) + setattr(model, self._state_version_key, [__version__]) + + # sanity check that root module has modelopt state now + assert self.is_converted(model, is_root=True), "Model must have modelopt state!" + + # store reference to state + self._state: ModeloptStateList = getattr(model, self._state_key) + self._state_version = getattr(model, self._state_version_key) + + @property + def has_state(self) -> bool: + """Return whether the model has a non-trivial modelopt state.""" + return bool(self._state) + + @classmethod + def is_converted(cls, model: nn.Module, is_root: bool = False) -> bool: + """Check if model is converted. + + Args: + model: A model to be checked for state/metadata from the convert process. + is_root: Additionally check whether the module with state is the root module. + + Returns: + True if the model contains modelopt state indicating that it has been converted. + + This method raises an assertion when multiple modelopt_states are detected or when is_root is + set to True but the module with state is not the root module. + """ + # check for submodules with state + mods_with_state = [name for name, m in model.named_modules() if hasattr(m, cls._state_key)] + # check if there is multiple submodules with state + assert len(mods_with_state) <= 1, "Model has multiple modelopt states!" + is_converted = bool(mods_with_state) + + # check if mod with state is root module if desired + if is_converted: + assert ( + not is_root or mods_with_state[0] == "" + ), "Model has modelopt state but not the root!" + + return is_converted + + # TODO: consider renaming state_dict??? + def state_dict(self) -> ModeloptStateList: + """Return the metadata of the model.""" + return self._state + + @property + def state_version(self) -> str: + """Return the version of modelopt state contained in the state manager.""" + return self._state_version[0] + + def load_state_dict(self, state_dict: ModeloptStateList, version: str) -> None: + """Load the provided ``state_dict`` to the modelopt_state.""" + assert not self.has_state, "Cannot load state_dict if there is already one." + + # Alert if the first 2 version numbers do not match, e.g., 0.3.2 vs 0.4.0. + if tuple(version.split(".")[:2]) != tuple(__version__.split(".")[:2]): + warnings.warn( + f"The checkpoint is stored with version {version}, but current version is" + f" {__version__}. Compatibility of checkpoint with current version is not guaranteed!" + ) + + # make sure we operate on deepcopy + state_dict = copy.deepcopy(state_dict) + # add modes one-by-one + for m_str, m_state in state_dict: + # adds config and metadata with sanity checks + config = self.get_config_class(m_str, m_state["config"]) + self.add_mode(m_str, config, m_state["metadata"]) + + # overwrite state manually afterwards to ensure exact consistency with provided state_dict + self._state.clear() + self._state.extend(state_dict) + self._state_version[0] = version + + @classmethod + def transfer_state_dict(cls, model_from: nn.Module, model_to: nn.Module) -> None: + """Transfer the state (same instance) from one model to another.""" + manager_from = ModeloptStateManager(model_from, init_state=False) # state must exist + manager_to = ModeloptStateManager(model_to, init_state=True) # state must NOT exist + + # transfer state_dict (this uses sanity checks + deepcopy) + manager_to.load_state_dict(manager_from.state_dict(), manager_from.state_version) + + # manually set the state dict to be the exact same instance + setattr(model_to, cls._state_key, manager_from.state_dict()) + manager_to = ModeloptStateManager(model_to, init_state=False) # state must exist now + + # remove state from model_from + delattr(model_from, cls._state_key) + + def modes_with_states( + self, + ) -> Iterator[tuple[_ModeDescriptor, ModeloptBaseConfig, MetadataDict]]: + """Yield the mode together with the full config and metadata from the state.""" + for m_str, m_state in self._state: + config = self.get_config_class(m_str, m_state["config"]) + yield _ModeRegistryCls.get_from_any(m_str), config, m_state["metadata"] + + @property + def last_mode(self) -> Optional[_ModeDescriptor]: + """Return the last mode applied to the model (last stored mode).""" + return _ModeRegistryCls.get_from_any(self._state[-1][0]) if self._state else None + + @property + def _last_metadata(self) -> MetadataDict: + """Return the metadata of the last mode applied to the model (must exist!).""" + return self._state[-1][1]["metadata"] + + @property + def _last_config(self) -> ModeloptBaseConfig: + """Return the config of the last mode applied to the model (must exist!).""" + return self.get_config_class(self._state[-1][0], self._state[-1][1]["config"]) + + @_last_config.setter + def _last_config(self, config: ModeloptBaseConfig) -> None: + """Set the config of the last mode applied to the model (must exist!).""" + self._state[-1][1]["config"] = config.model_dump() + + @property + def _export_stack(self) -> Deque[tuple[str, str]]: + """Infer the stack of export modes that still must be applied from existing modes. + + Returns: + A deque of tuples of the form ``(mode_str, export_mode_str)`` representing the mode + which requires an export mode and the export mode itself. + """ + stack = deque() + for m, _, _ in self.modes_with_states(): + if m.export_mode: + stack.append((str(m), m.export_mode)) + elif m.is_export_mode: + assert str(m) == stack.pop()[1], "Inconsistent export stack!" + return stack + + @staticmethod + def get_config_class(mode: ModeType, config: ConfigDict) -> ModeloptBaseConfig: + """Standardize the provided config to the corresponding config class.""" + # validate and standardize to the config class and return + return _ModeRegistryCls.get_from_any(mode).config_class(**config) + + def check_mode(self, mode: ModeType) -> None: + """Check if the proposed mode is compatible with the current state.""" + # standardize mode to descriptor + mode_d = _ModeRegistryCls.get_from_any(mode) + + # check for export mode compatibility + export_stack = self._export_stack + if mode_d.is_export_mode: + assert ( + export_stack and str(mode_d) == export_stack[-1][1] + ), f"Cannot add {mode_d} according to the current export stack: {export_stack}." + + # sanity checks for next mode incompatibilities according to the current last mode + last_mode = self.last_mode + if last_mode: + assert last_mode.next_modes is None or str(mode_d) in last_mode.next_modes, ( + f"Cannot add {mode_d} after {last_mode}! Next modes of {last_mode} are" + f" {last_mode.next_modes}." + ) + + # sanity checks for next mode incompatible with last mode in the stack. These + # incompatibilities still apply since we did not apply corresponding export mode yet. + if export_stack: + # last_m_stack := last mode on the stack that has a corresponding export mode + last_m_stack = _ModeRegistryCls.get_from_any(export_stack[-1][0]) + assert last_m_stack.next_modes is None or str(mode_d) in last_m_stack.next_modes, ( + f"Cannot add {mode_d} after {last_m_stack}! Next modes of {last_m_stack} are" + f" {last_m_stack.next_modes}." + ) + + def add_mode(self, mode: ModeType, config: ModeloptBaseConfig, metadata: MetadataDict) -> None: + """Add mode and update state in-place. + + Note that self._state is a list (preserves insertion order of keys) and we can therefore + recall the order of modes! + """ + # standardize mode to descriptor + mode_d = _ModeRegistryCls.get_from_any(mode) + + # sanity checks for mode incompatibilities + self.check_mode(mode_d) + + # store mode information + m_state: ModeState = {"config": config.model_dump(), "metadata": metadata} + + self._state.append((str(mode_d), m_state)) + + def update_last_state_before_new_mode(self, model: nn.Module) -> None: + """Update the metadata and config of the last mode applied to the model.""" + last_mode = self.last_mode + if last_mode is not None: + last_config = self._last_config + last_mode.update_for_new_mode(model, last_config, self._last_metadata) + self._last_config = last_config + + def update_last_state_before_save(self, model: nn.Module) -> None: + """Update the metadata and config of the last mode applied to the model.""" + last_mode = self.last_mode + if last_mode is not None: + last_config = self._last_config + last_mode.update_for_save(model, last_config, self._last_metadata) + self._last_config = last_config + + +class ApplyModeError(RuntimeError): + """Error raised when applying a mode to a model fails.""" + + +class ModelLikeModule(nn.Module): + """Just a temp module type to store the initialization recipe for the actual model.""" + + def __init__(self, modellike: ModelLike) -> None: + super().__init__() + assert not isinstance(modellike, nn.Module), "modellike should not be a nn.Module!" + self.modellike = modellike + + def init_modellike(self) -> nn.Module: + """Initialize the modellike to be an actual model.""" + model = init_model_from_model_like(self.modellike) + ModeloptStateManager.transfer_state_dict(self, model) + return model + + +def _check_init_modellike(model: nn.Module, mode: _ModeDescriptor) -> nn.Module: + """Utility to initialize a ModelLikeModule if needed according to the mode.""" + if mode.require_model_like: + assert isinstance(model, ModelLikeModule), "Model must be a ModelLikeModule!" + elif isinstance(model, ModelLikeModule): + model = model.init_modellike() + return model + + +def apply_mode( + model: ModelLike, + mode: ModeLike, + registry: Optional[_ModeRegistryCls] = None, + init_state: Optional[bool] = None, +) -> nn.Module: + """Apply the provided modes the model, record the changes, and return the model. + + Args: + model: A model-like object. Can be an nn.Module, a model class type, or a tuple. + tuple must be of the form ``(model_cls,)`` or ``(model_cls, args)`` or + ``(model_cls, args, kwargs)``. Model will be initialized as + ``model_cls(*args, **kwargs)``. + mode: A mode, a list of modes or a list of tuples containing the mode and its config. The + mode may be specified as a string or as the actual + :class:`_ModeDescriptor` class such as + :class:`QuantizeModeDescriptor` class. + registry: An optional mode registry from which to retrieve the mode. If not provided, all + registries will be searched. + init_state: Flag indicating whether we should initialize the state manager for the model. If + not provided, it will be inferred from the model. This flag can be used to enforce a + certain behavior. For example, for ``init_state=True`` the state manager will raise an + error if the model already contains state. + + Returns: + The converted model after applying the desired modes. + """ + # initialize ModelLikeModule if needed. + model = model if isinstance(model, nn.Module) else ModelLikeModule(model) + + # vanilla case (just initialize+return) + if not mode: + return model.init_modellike() if isinstance(model, ModelLikeModule) else model + + # check if the model is in a wrapper + model = unwrap_model(model, raise_error=True) + + # standardize mode to ModeConfigList + mode_and_config = get_mode_config(mode) + + # get or initialize the state manager for the model + manager = ModeloptStateManager( + model, + init_state=( + not ModeloptStateManager.is_converted(model) if init_state is None else init_state + ), + ) + + # get mode function based on registry argument + get_mode = registry.__getitem__ if registry else _ModeRegistryCls.get_from_any + + # update metadata of currently last mode before adding new modes + manager.update_last_state_before_new_mode(model) + + # check whether a ModelLike should be initialized + model = _check_init_modellike(model, get_mode(mode_and_config[0][0])) + + # loop through modes and call convert entrypoint for each mode and record data in manager. + for m, config in mode_and_config: + manager.check_mode(m) + config = manager.get_config_class(m, config) + model, metadata = get_mode(m).convert(model, config) + manager.add_mode(m, config, metadata) + + # If the existing mode is empty, create an model instance from ModelLikeModule. + if not manager.has_state and isinstance(model, ModelLikeModule): + model = model.init_modellike() + + # it cannot be a ModelLikeModule anymore at the end + assert not isinstance(model, ModelLikeModule), "Model must be a regular Module now!" + + # return model with state recorded + return model + + +def get_mode(model: nn.Module) -> Optional[_ModeDescriptor]: + """Get mode of converted network. + + model: A model that contains modelopt_state + + The mode of the model is defined as the last mode activated during the convert process. + """ + if ModeloptStateManager.is_converted(model): + return ModeloptStateManager(model).last_mode + return None + + +def modelopt_state(model: nn.Module) -> dict[str, Any]: + """Return the modelopt state dict describing the modifications to the model. + + Note that the returned ``modelopt_state`` does not contain the model parameters such as weights and biases. + ``modelopt_state`` is useful for saving and loading various modelopt optimization states separately from the + model parameters. For example: + + .. code-block:: + + import modelopt.torch.opt as mto + + # Save the modelopt state and model weights separately + torch.save(mto.modelopt_state(model), "modelopt_state.pt") # Save the modelopt state + torch.save(model.state_dict(), "model_weights.pt") # Save the model weights + + If you want to save the model weights and the modelopt state together, please use + :meth:`mto.save()`. + + Args: + model: the modelopt-modified model. + + Returns: + An modelopt state dictionary describing the modifications to the model. + """ + # unwrap model + model = unwrap_model(model, force_unwrap=True) + + # retrieve state manager + manager = ModeloptStateManager( + model=model if ModeloptStateManager.is_converted(model) else None + ) + + # update metadata of current mode as needed + manager.update_last_state_before_save(model) + + # construct state dict and return it + objs = { + "modelopt_state_dict": ( + manager.state_dict() + ), # empty state_dict is okay (saving regular models) + "modelopt_version": __version__, + } + return objs + + +def save(model: nn.Module, f: Union[str, os.PathLike, BinaryIO], **kwargs) -> None: + """Save a model's state dict together with the modelopt state dict to restore its architecture. + + Args: + model: Any model. + f: Target file location. + **kwargs: additional args for ``torch.save()``. + + .. note:: + + If model is a wrapper such as DistributedDataParallel, it will be unwrapped for saving. + """ + # unwrap model + model = unwrap_model(model, warn=True) + + # store ckpt + ckpt_dict = { + "modelopt_state": modelopt_state(model), + "model_state_dict": model.state_dict(), + } + + # store object + torch.save(ckpt_dict, f, **kwargs) + + +def restore_from_modelopt_state(model: ModelLike, modelopt_state: dict[str, Any]) -> nn.Module: + """Restore the model architecture from the modelopt state dictionary based on the user-provided model. + + This method does not restore the model parameters such as weights, biases and quantization scales. + Please load the weights and biases with the original checkpoint loading method after restoring + modelopt states with `restore_from_modelopt_state`. For example: + + .. code-block:: python + + import modelopt.torch.opt as mto + + model = ... # Create the model-like object + + # Restore the previously saved modelopt state followed by model weights + mto.restore_from_modelopt_state( + model, torch.load("modelopt_state.pt") + ) # Restore modelopt state + model.load_state_dict(torch.load("model_weights.pt"), ...) # Load the model weights + + If you want to restore the model weights and the modelopt state with saved scales, please use + :meth:`mto.restore()`. + + Args: + model: A model-like object. Can be an nn.Module, a model class type, or a tuple. + tuple must be of the form ``(model_cls,)`` or ``(model_cls, args)`` or + ``(model_cls, args, kwargs)``. Model will be initialized as + ``model_cls(*args, **kwargs)``. + modelopt_state: The modelopt state dict describing the modelopt modifications to the model. The + ``modelopt_state`` can be generated via + :meth:`mto.modelopt_state()`. + + Returns: + A modified model architecture based on the restored modifications with the unmodified + weights as stored in the provided ``model`` argument. + + .. note:: + + Note that wrappers such as DistributedDataParallel are `not` supported during the restore + process. Please wrap the model after the restore process. + """ + # initialize ModelLikeModule if needed. + model = model if isinstance(model, nn.Module) else ModelLikeModule(model) + + # initialize state manager and load state + manager = ModeloptStateManager(model=model, init_state=True) + manager.load_state_dict( + modelopt_state["modelopt_state_dict"], modelopt_state["modelopt_version"] + ) + + # apply restore entrypoints for each of the modes + for i, (m, config, metadata) in enumerate(manager.modes_with_states()): + if i == 0: + model = _check_init_modellike(model, m) + model = m.restore(model, config, metadata) + + # If the existing mode is empty, create an model instance from ModelLikeModule. + if not manager.has_state and isinstance(model, ModelLikeModule): + model = model.init_modellike() + + # it cannot be a ModelLikeModule anymore at the end + assert not isinstance(model, ModelLikeModule), "Model must be a regular Module now!" + + return model + + +def restore(model: ModelLike, f: Union[str, os.PathLike, BinaryIO], **kwargs) -> nn.Module: + """Load the checkpoint, restore the modelopt model modifications, and load the model's weights. + + Args: + model: A model-like object. Can be an nn.Module, a model class type, or a tuple. + tuple must be of the form ``(model_cls,)`` or ``(model_cls, args)`` or ``(model_cls, args, kwargs)``. + Model will be initialized as ``model_cls(*args, **kwargs)``. + f: Target file location generated by :meth:`mto.save()`. + **kwargs: additional args for ``torch.load()``. + + Returns: + The model with original weights and stored architecture. + + .. note:: + + Note that wrappers such as DistributedDataParallel are `not` supported during the restore + process. Please wrap the model after the restore process. + """ + # initialize ModelLikeModule if needed. + model = model if isinstance(model, nn.Module) else ModelLikeModule(model) + + # load checkpoint + kwargs.setdefault("map_location", "cpu") + objs = torch.load(f, **kwargs) + + # restore model architecture + model_restored = restore_from_modelopt_state(model, objs["modelopt_state"]) + + # load weights from checkpoint + model_restored.load_state_dict(objs["model_state_dict"]) + + # it cannot be a ModelLikeModule anymore at the end + assert not isinstance(model_restored, ModelLikeModule), "Model must be a regular Module now!" + + return model_restored + + +# TODO: add a generic export function that will apply all remaining export modes. +def export(model: nn.Module) -> nn.Module: + """Fully export the model to a regular model and finalize any model modifications.""" + raise NotImplementedError + + +def has_nas_modelopt_state(modelopt_state: dict[str, Any]) -> bool: + """Check if the modelopt state contains NAS state.""" + from modelopt.torch.nas.mode import NASModeRegistry + + return any(m_str in NASModeRegistry for m_str, _ in modelopt_state["modelopt_state_dict"]) diff --git a/modelopt/torch/opt/dynamic.py b/modelopt/torch/opt/dynamic.py new file mode 100644 index 0000000000..1fb800e99b --- /dev/null +++ b/modelopt/torch/opt/dynamic.py @@ -0,0 +1,1295 @@ +# 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. + +"""Basic dynamic module class and hparam class.""" + +# cython: annotation_typing = False + +import inspect +import types +import warnings +from contextlib import contextmanager +from itertools import chain +from typing import Any, Callable, Iterable, Iterator, Optional, TypeVar, Union + +import torch +import torch.nn as nn +from pydantic import create_model +from torch.fx.proxy import Proxy +from torch.nn.parameter import Parameter + +from modelopt.torch.utils import get_unwrapped_name, is_channels_last, unwrap_model + +from .config import ModeloptBaseRule, RulesDict +from .hparam import Hparam + +__all__ = ["DynamicModule", "_DMRegistryCls", "DynamicSpace"] + +_pytorch_managed = type("_pytorch_managed", (), {}) # pylint: disable=invalid-name +_da_val_default = type("_da_val_default", (), {}) # pylint: disable=invalid-name + + +DynamicAttributeCallback = Callable[["DynamicModule", Any], Any] +SetAttrHook = Callable[["DynamicModule", str, Any], None] +DelAttrHook = Callable[["DynamicModule", str], None] + + +class _FoldedCallback: + """A callback functor that enables folding additional callbacks into an existing callback.""" + + def __init__(self, *callbacks: DynamicAttributeCallback): + self._callbacks: list[DynamicAttributeCallback] = [] + self.extend(callbacks) + + @property + def callback(self) -> DynamicAttributeCallback: + """Return the main callback.""" + return self._callbacks[0] + + @callback.setter + def callback(self, cb: DynamicAttributeCallback): + """Set the main callback.""" + self._callbacks[0] = cb + + def __iter__(self) -> Iterator[DynamicAttributeCallback]: + """Iterate over all callbacks.""" + yield from self._callbacks + + def __call__(self, self_module: "DynamicModule", val: Any) -> Any: + """Call the callback and all other callbacks.""" + if isinstance(val, Proxy): + return val + for cb in self: + val = cb(self_module, val) + return val + + def __len__(self) -> int: + """Return the number of callbacks.""" + return len(self._callbacks) + + def extend(self, cbs_other: Iterable[DynamicAttributeCallback]): + """Extend the list of other callbacks.""" + self._callbacks.extend(cbs_other) + + +class _DMAttributeManager: + """A class to manage the special attributes for dynamic modules in a hierarchical fashion. + + Abbreviations: + - og_cls: original class + - hp: hparam + - da: dynamic attribute + - cb: callback + - attr: generic temporary attributes + + This class handles all types of special attributes for dynamic modules. It is used to store + them, manage them, and provide access to them. Moreover, this class stores the hierarchy of + dynamic module conversions that were applied to the module and can associate the special + attributes with the correct level of inheritance. + + Specifically, we handle the following types of special attributes: + + 1. **Original classes**: + We maintain a stack of original classes as we stack/convert dynamic modules. As we stack + dynamic modules on top or remove/export them, we maintain the stack. + + 2. **Hparams**: + We store hparams in a list of dictionaries. This way we can associate each hparam with + the corresponding level of hierarchy/inheritance. For hparams we do NOT allow adding or + overwriting hparams that exist in a different level of the stack, i.e., hparams must be + _unique_ across all levels of inheritance. Note that you can still access hparams from + any level of inheritance. + + 3. **Dynamic attributes**: + We store dynamic attributes in a dictionary. Additionally, a dynamic attribute can be + associated with a callback in an _arbitrary_ number of levels in the hierarchy. The + callbacks are stored in a list of dictionaries. This way we can associate each callback + with the corresponding level of hierarchy/inheritance and still allow for multiple + callbacks for the same dynamic attribute. + + 4. **Generic attributes**: + These are generic attributes that are not hparams or dynamic attributes. They are kept + track of so that we can remove them again upon the final export. They must be unique + and cannot correspond to an hparam or dynamic attribute. + """ + + def __init__(self): + self._og_cls_all: list[type[nn.Module]] = [] + self._hp_all: list[dict[str, Hparam]] = [] + self._da: dict[str, Any] = {} + self._da_cb_all: list[dict[str, _FoldedCallback]] = [] + self._attr: dict[str, tuple[SetAttrHook, DelAttrHook]] = {} + + def __bool__(self) -> bool: + all_data = [self._og_cls_all, self._da, self._da_cb_all, self._hp_all, self._attr] + has_state = [bool(x) for x in all_data] + if not has_state[0]: + # if we have no original class, we should have no other state either! + assert not any(has_state), "Inconsistent state for DM attributes!" + return has_state[0] + + def _get_lookup(self, name: str, lookup_all: list[dict[str, Any]]) -> dict[str, Any]: + """Return the lookup for the given name.""" + for lookup in lookup_all: + if name in lookup: + return lookup + raise KeyError(f"{name} not found!") + + @property + def og_cls(self) -> type[nn.Module]: + """Return the original class of the dynamic module.""" + return self._og_cls_all[-1] + + @property + def level(self) -> int: + """Return the current level of stacked inheritance.""" + return len(self._og_cls_all) + + def append_level(self, original_cls: type[nn.Module]): + """Initialize a new level of special atttributes.""" + self._og_cls_all.append(original_cls) + self._hp_all.append({}) + self._da_cb_all.append({}) + + def pop_level(self) -> type[nn.Module]: + """Remove the last level of special attributes and return the original class.""" + # sanity checks on hparams + hp_lookup = self._hp_all[-1] + assert not hp_lookup, "Some hparams were not removed properly!" + + # sanity checks on dynamic attributes and corresponding callbacks + cb_lookup = self._da_cb_all[-1] + assert not cb_lookup, "Some dynamic attribute callbacks were not removed properly!" + + da_keys_remaining = set(chain(*self._da_cb_all[:-1])) + assert da_keys_remaining == self.da_keys(), "Dynamic attributes were not removed properly!" + + # sanity checks on temporary attributes on last level + if self.level <= 1: + assert not self._attr, "Some attributes were not removed properly!" + + # after clean-up, we can pop and return the original class + self._hp_all.pop() + self._da_cb_all.pop() + return self._og_cls_all.pop() + + def hp_keys(self, all: bool = True) -> set[str]: + """Return the keys of current or all hparams.""" + idx_start = 0 if all else -1 + return set(chain(*self._hp_all[idx_start:])) + + def get_hp(self, name: str) -> Hparam: + """Return the hparam with the given name.""" + hp_lookup = self._get_lookup(name, self._hp_all) + return hp_lookup[name] + + def named_hps(self, all: bool = True) -> Iterator[tuple[str, Hparam]]: + """Return the name and hparam for all or current hparams.""" + idx_start = 0 if all else -1 + for hp_lookup in self._hp_all[idx_start:]: + for name, hp in hp_lookup.items(): + yield name, hp + + def set_hp(self, name: str, hparam: Hparam): + """Store hparam by the provided name in the right location.""" + # we do not allow the same hparam in different levels of inheritance + if name in self.hp_keys(): + assert name in self.hp_keys(all=False), "Hparam already exists in the base cls!" + self._hp_all[-1][name] = hparam + + def pop_hp(self, name: str) -> Any: + """Pop hparam by the provided name.""" + hp_lookup = self._get_lookup(name, self._hp_all) + return hp_lookup.pop(name) + + def da_keys(self) -> set[str]: + """Return the keys of all dynamic attributes.""" + return set(self._da) + + def set_da( + self, + name: str, + val: Any = _da_val_default, + cb: Optional[DynamicAttributeCallback] = None, + ): + """Store a dynamic attribute together with its callback method.""" + # sanity checks + if name in self.da_keys(): + val = self.get_da_value(name) if val is _da_val_default else val + else: + assert val is not _da_val_default, "Value must be provided for new dynamic attribute!" + assert cb is not None, "Callback must be provided for new dynamic attribute!" + + # store value + self._da[name] = val + + # store wrapped callback + # NOTE: we only allow updating the original callback of this level of inheritance. Updating + # folded callbacks is not allowed! + if cb is not None: + if name in self._da_cb_all[-1]: + self._da_cb_all[-1][name].callback = cb + else: + self._da_cb_all[-1][name] = _FoldedCallback(cb) + + def pop_da(self, name: str) -> Any: + """Pop dynamic attribute by the provided name.""" + val = self._da.pop(name) + for cb_lookup in self._da_cb_all: + cb_lookup.pop(name, None) + return val + + def get_da_value(self, name: str) -> Any: + """Return the raw value of the dynamic attribute.""" + return self._da[name] + + def get_da_cb(self, name: str) -> _FoldedCallback: + """Return the full callback for the dynamic attribute.""" + cbs_all = list(chain.from_iterable(cb_d[name] for cb_d in self._da_cb_all if name in cb_d)) + return _FoldedCallback(*cbs_all) + + @contextmanager + def retain_cbs(self): + """Context manager to retain all callback information for dynamic attributes. + + Any changes that are made to the callbacks within this context manager are not retained upon + exiting the context manager. However, changes to the value of dynamic attributes as well + as well deletions and additions of dynamic attributes are retained. + """ + # retain original callbacks while exposing a list of empty dictionaries + da_cb_all = self._da_cb_all + da_keys_before = self.da_keys() + self._da_cb_all = [{} for _ in da_cb_all] + try: + yield None + finally: + # 1. We will restore callbacks to the original state for da's that still exist + # 2. We will maintain callbacks for newly added da's + # 3. We will throw away callbacks for da's that were removed + da_keys_removed = da_keys_before - self.da_keys() + da_keys_new = self.da_keys() - da_keys_before + da_keys_same = da_keys_before & self.da_keys() + + da_cb_all_final = [] + for cb_lookup_before, cb_lookup_now in zip(da_cb_all, self._da_cb_all): + cb_lookup_new = {k: v for k, v in cb_lookup_now.items() if k in da_keys_new} + cb_lookup_same = {k: v for k, v in cb_lookup_before.items() if k in da_keys_same} + cb_lookup_final = {**cb_lookup_new, **cb_lookup_same} + assert not (cb_lookup_final.keys() & da_keys_removed) # final sanity check + da_cb_all_final.append({**cb_lookup_same, **cb_lookup_new}) + self._da_cb_all = da_cb_all_final + + def fold_cbs(self) -> set[str]: + """Fold all callbacks that appear in lower levels into the corresponding lower level. + + After this call, only dynamic attributes that appear solely in the top-level remain in the + top-level. + + Returns: + The set of dynamic attributes that remain in the top-level. + """ + cb_lookup = self._da_cb_all[-1] + + # check if we can fold the callbacks into a lower level + for cb_lookup_other in reversed(self._da_cb_all[:-1]): + for k_other, cb_other in cb_lookup_other.items(): + if k_other in cb_lookup: + cb_other.extend(cb_lookup.pop(k_other)) + + # sanity check: at this point there should be no overlap between top-level and other levels + assert not any(k in cb_lookup for k in chain(*self._da_cb_all[:-1])), "Callback overlap!" + + return set(cb_lookup) + + def attr_keys(self) -> set[str]: + """Return the keys of all or current other attributes.""" + return set(self._attr) + + def get_attr_set_hook(self, name: str) -> SetAttrHook: + """Return the setter for the given name.""" + return self._attr[name][0] + + def set_attr(self, name: str, set_hook: Optional[SetAttrHook], del_hook: Optional[DelAttrHook]): + """Store the name of an attribute in the other list.""" + assert name not in self.attr_keys(), "Attribute already exists!" + + set_hook = set_hook or (lambda m, n, v: None) + del_hook = del_hook or (lambda m, n: None) + self._attr[name] = (set_hook, del_hook) + + def pop_attr(self, name: str) -> DelAttrHook: + """Pop attr by the provided name and return hook for delattr.""" + return self._attr.pop(name)[1] + + +class DynamicModule(nn.Module): + """Base class for dynamic modules. + + Dynamic modules are usually extended from ``torch.nn.Module``'s. They + dynamically support a family of `torch.nn.Module`'s with different architectural + parameters, such as input/output channel numbers. + + Dynamic modules can also be used to construct the basic searchable unit in a search space with + the option to select and sample a candidate unit. + + Candidate units are usually described using ``Hparam`` objects and dynamic attributes. Each + hparam describes a basic searchable unit (e.g. number of output channels in ``DynamicConv2d``). + Dynamic attributes are callbacks to dynamically construct the attribute depending on the current + value of the hparam(s), e.g., the ``weight`` tensor in ``DynamicConv2d`` that depends on the + hparam ``out_channels``. + + In addition, dynamic modules also support registering general attributes that are removed upon + export the module. This is useful for storing temporary attributes that are not hparams or + existing attributes that are converted to dynamic attributes. + + For a ``DynamicModule`` class that contains other dynamic modules, the class implementation + should ensure only to expose ``hparams`` in the outermost class and handle other ``hparams`` + internally including ``hparams`` of child modules that are exposed on their own usually + (e.g. block module implementations containing DynamicLinear). + """ + + # this is needed to store the special attributes for dynamic modules + _dm_attribute_manager: _DMAttributeManager + + def __init__(self, *args, **kwargs): + """Initializing a dynamic module is not allowed!""" + raise RuntimeError("DynamicModule cannot be initialized directly; use convert instead!") + + def _get_dm_attribute_manager(self, use_default: bool = False) -> _DMAttributeManager: + """Return the attribute manager or a default one if not available.""" + if "_dm_attribute_manager" not in self.__dict__: + if use_default: + return _DMAttributeManager() + raise AttributeError("DynamicModule.convert() must be called before using the module!") + return self._dm_attribute_manager + + def _register_hparam(self, name: str, hparam: Hparam): + """Register hparam by the provided name and do some sanity checks during registering. + + Args: + name: The name of the hparam. Note that an hparam can be a new attribute or an existing + attribute that is not already registered as temporary attribute or dynamic + attribute. In addition, an hparam cannot be registered for tensor and therefore + neither for parameters nor buffers. Finally, an hparam cannot be registered when an + hparam under the same exists in one of the parent classes. However, one can re- + register an hparam of itself. + hparam: The hparam object to be registered. + + An hparam is useful to configure the dynamic module on the fly. The collection of hparams + define the configurable space of the dynamic module. + """ + # retrieve manager + manager = self._get_dm_attribute_manager() + + # various sanity checks + if not isinstance(hparam, Hparam): + raise TypeError(f"Cannot assign {type(hparam)} to {name}. Hparam required!") + elif name in manager.attr_keys(): + raise KeyError(f"Cannot use the reserved name {name} to assign an hparam!") + elif name in manager.da_keys(): + raise RuntimeError(f"Cannot register attribute {name} that is dynamic as hparam.") + elif name in manager.hp_keys(all=True) and name not in manager.hp_keys(all=False): + raise RuntimeError(f"Cannot overwrite hparam {name} that is an hparam in the base cls.") + elif "." in name: + raise KeyError('hparam name can\'t contain "."') + elif name == "": + raise KeyError('hparam name can\'t be empty string ""') + elif name in ["_parameters", "_buffers", "_modules"]: + raise KeyError(f"Cannot use reserved name {name} to assign hparam.") + # delete name if it already exists + elif hasattr(self, name): + if isinstance(getattr(self, name), (torch.Tensor, torch.nn.Module)): + raise RuntimeError(f"Cannot register tensor/module attribute {name} as hparam.") + delattr(self, name) + + # store new hparam + manager.set_hp(name, hparam) + + def _register_dynamic_attribute(self, name: str, callback: DynamicAttributeCallback): + """Register a dynamic attribute together with its callback method. + + Args: + name: The name of the dynamic attribute. + callback: The callback method that is executed when the dynamic attribute is accessed. + Generally, this callback can access the whole module (``self``) as well as expect + to have access to the original attribute without any modifications. + + Note that a dynamic attribute should be an existing attribute that is dynamically affected + when the value of one or more hparams changes. It is not allowed to register a dynamic + attribute for a non-existing attribute. However, specifically tensors, parameters, and + buffers are supported. Moreover, a dynamic attribute can also be registered for an existing + dynamic attribute in the parent class. In this case, the ``callback`` will be executed + *after* the callback of the parent class. + """ + # retrieve manager + manager = self._get_dm_attribute_manager() + + # various sanity checks + if not hasattr(self, name): + raise AttributeError(f"{name} is not a valid attribute.") + elif name in manager.attr_keys(): + raise KeyError(f"Cannot use the reserved name {name} to assign a dynamic attribute!") + elif name in manager.hp_keys(): + raise RuntimeError(f"Cannot register attribute {name} that is an hparam as dynamic.") + elif "." in name: + raise KeyError('dynamic attribute name can\'t contain "."') + elif name == "": + raise KeyError('dynamic attribute name can\'t be empty string ""') + + # however we have to handle regular attributes and params/buffers separately. Specifically, + # params/buffers won't appear in self.__dict__ since they are already managed by pytorch + # internally. + if name in self.__dict__["_parameters"] or name in self.__dict__["_buffers"]: + # here the attribute is already managed, so simply storing a fake reference. + value = _pytorch_managed + elif name in self.__dict__: + # here, we retrieve the value and delete the attribute from the regular __dict__ so we + # manage it. + value = getattr(self, name) + delattr(self, name) + elif name in manager.da_keys(): + # here we previously registered the attribute as dynamic + value = manager.get_da_value(name) + else: + raise RuntimeError(f"The value of attribute {name} cannot be retrieved correctly!") + + # store new dynamic attribute + manager.set_da(name, value, callback) + + def _register_temp_attribute( + self, + name: str, + val: Any, + set_hook: Optional[SetAttrHook] = None, + del_hook: Optional[DelAttrHook] = None, + ): + """Register a temporary attribute to the instance that is deleted upon the final export. + + Args: + name: The name of the attribute. + val: The value of the attribute. + set_hook: A hook that is executed before the attribute is set (Defaults to + ``lambda m, n, v: None``). + del_hook: A hook that is executed before the attribute is deleted (Defaults to + ``lambda m, n: None``). + + Any attribute that is set like this will be removed upon the final export, i.e., only when + the module is not a dynamic module anymore after export. This is helpful if you want to + store extra attributes in the class that are not dynamic attributes or hparams but should + still be removed upon the final export, e.g., attributes or buffers that may access during + the dynamic attribute callbacks. + + Unlike dynamic attributes, this can be used to set attributes that did not exist before and + *cannot* be used to overwrite existing attributes. + """ + # retrieve manager + manager = self._get_dm_attribute_manager() + + # sanity checks + if hasattr(self, name): + raise KeyError(f"Attribute {name} already exists!") + elif isinstance(val, Hparam): + raise RuntimeError("Please register hparam directly via _register_hparam()!") + + # store in manager and set new attribute + manager.set_attr(name, set_hook, del_hook) + setattr(self, name, val) + + @torch.no_grad() + def export(self) -> nn.Module: + """Export self (a dynamic module) **in-place** and return the exported module. + + The export process will remove the top-level dynamic module and replace it with the original + module class. Note that the original class may be either another type of dynamic module or + a vanilla nn.Module. Consequently, any methods (including properties) that are implemented + in the child class will be removed. Hparams as well as dynamic and temporary attributes are + handled in a special fashion, see below. + + In order to ensure that the exported module is still consistent there a several mechanisms + in place to handle hparams, dynamic attributes, and temporary attributes: + + * **Hparams** of the current type are replaced with their currently active value. + Note that we do not need to explicitly handle hparams of the parent class as they are + mutually-exlusive, i.e., hparams are unique across all levels of inheritance. + * **Dynamic Attributes** are handled depending on whether they exist in a parent class: + 1. The same dynamic attribute exists in a parent class. In this case, the callback is + folded into ("appended to") the callback for the same dyanmic attribute of the parent + class. This way we ensure that the final value of the attribute remains consistent. + 2. The dynamic attribute does not exist in a parent class. In this case, the attribute + is not dynamic anymore as there are no more callbacks that could affect the value. + Therefore, we simply overwrite the underlying original object with the current value + and revert it to a regular attribute. + * **Temporary Attributes** are kept until the final export, i.e., until the resulting class + is not a dynamic module anymore. This is to ensure that folded callbacks that may need + access to these attributes can still access them. + """ + manager = self._get_dm_attribute_manager() + + # dissolve dynamic attributes when their corresponding callbacks cannot be folded into the + # parent class callbacks anymore. + da_removable = manager.fold_cbs() + for k in da_removable: + val = getattr(self, k) + if isinstance(val, torch.Tensor): + val = val.detach().clone() + if k in self._parameters: + val = Parameter(val) + manager.pop_da(k) # remove dynamic attribute from manager + setattr(self, k, val) # now we set it as regular attribute/parameter/buffer + + # replace hparams with active values and delete removable attributes in last level + with self._dict_with_special(): # this way we avoid recursion issues + for k in manager.hp_keys(all=False): + val = getattr(self, k) + manager.pop_hp(k) + setattr(self, k, val) + + # remove attributes in the last level + if manager.level <= 1: + for k in manager.attr_keys(): + delattr(self, k) + + # pop original class and the whole level of inheritance + self.__class__ = manager.pop_level() + + # double-check that any remaining dynamic attributes still work + for k in manager.da_keys(): + try: + getattr(self, k) + except Exception as e: + raise RuntimeError(f"Dynamic attribute {k} unretrievable after export!") from e + + # remove manager if not needed anymore and check whether it should still be a DynamicModule + is_dynamic = isinstance(self, DynamicModule) + if manager: + assert is_dynamic, "Exported module should still be a DynamicModule!" + else: + assert not is_dynamic, "Exported module must not be a DynamicModule anymore!" + delattr(self, "_dm_attribute_manager") + + return self + + @torch.no_grad() + def force_assign(self): + """Force re-assign all dynamic attributes to their current values. + + .. warning:: + + Note that this method overwrittes the actual buffers and parameters! Only use in + specific circumstances!! + """ + # force-reassign all dynamic attributes + for name in self._get_dm_attribute_manager().da_keys(): + val = getattr(self, name) + if isinstance(val, torch.Tensor): + val = val.detach().clone() + if name in self._parameters: + val = val if val is None else Parameter(val) + self.register_parameter(name, val) + elif name in self._buffers: + self.register_buffer(name, val) + else: + setattr(self, name, val) + + @classmethod + @torch.no_grad() + def convert(cls, module: nn.Module) -> "DynamicModule": + """Converts a module in-place into its dynamic counterpart by patching its class. + + Args: + module: The module to be converted into a dynamic module. + + Returns: + The converted dynamic module. + + This should generally be a *final* method and child classes should inherit ``_setup()`` + instead to customize the conversion process. + + Patching is achieved by updating the ``__class__`` attribute of the module to its dynamic + counterpart. The dynamic counterpart is a subclass of the original class, hence, we ensure + the module is fully compatible with the original class. Simultaneously, we can inject the + corresponding dynamic behavior in a standardized and rigoruos fashion. + """ + + def _remove_and_reattach_accelerate_monkey_patching(self): + # accelerate monkey patches the forward method of modules + # This is often used to pipeline parallelize a model or to perform forward with parameters that are + # cpu offloaded, but on-demand moved to GPU + # For example, see https://github.com/huggingface/accelerate/blob/2d74c0c077d41fdfee2bafca56527eda1db1228b/src/accelerate/hooks.py#L529 + + # Monkey patched modules will not call the correct forward path according to the MRO after calling convert + # To fix this, we need to remove the monkey patching, update the forward method and add the hook back + + if not hasattr(self.forward, "__func__") or ( + self.forward.__func__ is not self.__class__.forward + ): + if hasattr(self, "_hf_hook"): + # The forward of this module has been monkey patched by HF accelerate + # So it will not call the correct forward function after converting the module to a DynamicModule + from accelerate.hooks import add_hook_to_module, remove_hook_from_module + + assert hasattr(self, "_old_forward") + + hook = self._hf_hook + remove_hook_from_module(self) + self.forward = types.MethodType(self.__class__.forward, self) + add_hook_to_module(self, hook) + else: + warnings.warn( + "Received a module with monkey patched forward method. Dynamic converted module" + " might not work." + ) + + # update class + original_cls = type(module) + module.__class__ = cls + assert isinstance(module, cls), f"Failed to convert {original_cls} to {cls}!" # for mypy + + # setup/update the attribute manager + if issubclass(original_cls, DynamicModule): + assert hasattr(module, "_dm_attribute_manager"), "Attribute manager not found!" + else: + assert not hasattr(module, "_dm_attribute_manager"), "Attribute manager found!" + module._dm_attribute_manager = _DMAttributeManager() + module._dm_attribute_manager.append_level(original_cls) + + _remove_and_reattach_accelerate_monkey_patching(module) + + # setup new hparams and dynamic attributes + module._setup() + + return module + + def _setup(self): + """Setup dynamic attributes and hparams after the convert call. + + This method should be overriden by the child class! + """ + raise NotImplementedError("_setup() must be implemented by child class!") + + def modify(self, *args, **kwargs): + """Modify the module's dynamic choices in a standardized & scalable fashion. + + This method can be overriden by the child class! While users can also directly modify + the choices of individual hparams, this method should provide a way to modify a batch of + dynamic modules with the same arguments, e.g., ``out_features_ratio`` for ``DynamicLinear``. + + Note that arguments of the modify method that are exposed to the user via the rule system + should be specified as **keyword-only arguments**. When they are exposed as keyword-only + arguments, the `_DMRegistryCls` can automatically generate the corresponding config class + on the fly that lets user provide configs and then they are automatically validated before + being passed to the ``modify`` method. + + If possible, modify()'s keyword arguments should have default values that leave the hparams + intact if not provided, e.g., one might call ``some_dynamic_module.modify()`` without any + arguments and the module will remain unchanged. + """ + pass + + def freeze(self): + """Restrict the hparams of the dynamic module to the orginal choices. + + This is useful to enforce the behavior of the parent class. + + .. note:: + + After this call, the module's hparams can no longer be modified although the underlying + type is still a dynamic module. + """ + for _, hp in self.named_hparams(configurable=True): + hp.active = hp.original + hp.choices = [hp.original] + + @contextmanager + def reset_dynamic_attributes(self): + """Context manager to temporarily remove any dynamic attributes and re-register values. + + This context manager is intended to be used when we want to access a dynamic attribute in + its original unmodified version, i.e., without this class interfering with its original + value and its corresponding getattr/setattr/delattr behavior. + + Upon exiting the context manager, the dynamic attributes are re-registered and the same + callbacks are re-registered together with the new value. + """ + manager = self._get_dm_attribute_manager() + + with manager.retain_cbs(): + # pop all dynamic attributes and make sure they are accessible in their original form + da_keys = manager.da_keys() + for k in da_keys: + val = manager.pop_da(k) + if not hasattr(self, k): + setattr(self, k, val) + assert not manager.da_keys(), "Dynamic attributes were not removed properly!" + try: + yield None + finally: + # re-register dynamic attributes with current values + dummy callback (--> will be + # thrown away!) + # NOTE that we respect if attributes were deleted entirely. + for k in da_keys: + if hasattr(self, k): + self._register_dynamic_attribute(k, lambda _, v: v) + + @contextmanager + def _dict_with_special(self): + """Context manager that checks that __dict__ contains _modules, _parameters, _buffers.""" + # build up set of special keys that we temporarily add to __dict__ + manager = self._get_dm_attribute_manager(use_default=True) + nn_special = set() + for key in ["_modules", "_parameters", "_buffers"]: + if key not in self.__dict__ and (key in manager.da_keys() or key in manager.hp_keys()): + nn_special.add(key) + + # temporarily add to __dict__ + for key in nn_special: + self.__dict__[key] = getattr(self, key) + try: + yield None + finally: + # remove from __dict__ + for key in nn_special: + del self.__dict__[key] + + def __setattr__(self, name: str, value: Any): + """Set attr and specifically handle hparams as well as dynamic & temporary attributes.""" + # retrieve manager + manager = self._get_dm_attribute_manager(use_default=True) + + if isinstance(value, Hparam): + self._register_hparam(name, value) + elif name in manager.hp_keys(): + manager.get_hp(name).active = value + elif name in manager.da_keys() and manager.get_da_value(name) is not _pytorch_managed: + manager.set_da(name, value) + else: + if name in manager.attr_keys(): + manager.get_attr_set_hook(name)(self, name, value) + with self._dict_with_special(): + super().__setattr__(name, value) + + def __getattr__(self, name: str) -> Union[torch.Tensor, torch.nn.Module]: + """Get attr and specifically handle hparams as well as dynamic & temporary attributes.""" + # retrieve manager + manager = self._get_dm_attribute_manager(use_default=True) + + # check if we can get value from our hparams + if name in manager.hp_keys(): + return manager.get_hp(name).active + + # check for dynamic attributes + if name in manager.da_keys(): + value = manager.get_da_value(name) + # we might also need to grab value from super call (pytorch managed) + value = super().__getattr__(name) if value is _pytorch_managed else value + # apply all callbacks in order + return manager.get_da_cb(name)(self, value) + + # regular case + with self._dict_with_special(): + attr = super().__getattr__(name) + return attr + + def __delattr__(self, name: str): + """Del an attr and specifically handle hparams as well as dynamic & temporary attributes.""" + manager = self._get_dm_attribute_manager(use_default=True) + if name in manager.da_keys(): + manager.pop_da(name) + # check if it still exists + if hasattr(self, name): + delattr(self, name) + elif name in manager.hp_keys(): + manager.pop_hp(name) + else: + if name in manager.attr_keys(): + del_hook = manager.pop_attr(name) + del_hook(self, name) + return super().__delattr__(name) + + def get_hparam(self, target: str) -> Hparam: + """Look up and return hparam (like "torch.nn.Module.get_parameter()" but for hparam).""" + return self._get_dm_attribute_manager().get_hp(target) + + def named_hparams(self, configurable: Optional[bool] = None) -> Iterator[tuple[str, Hparam]]: + """Return an iterator over all hparams of the module. + + Args: + configurable: Whether to include configurable hparams. + + Yields: + (name, Hparam): tuple containing the name and hparam. + + Default behavior is to iterate over configurable and non-configurable hparams. Set + ``configurable`` accordingly to only iterate over either. If ``configurable`` is set to + ``True``, only configurable symbols are iterated over. If ``configurable`` is set to + ``False``, configurable symbols are skipped over (only non-configurable symbols). + """ + for hp_name, hp in self._get_dm_attribute_manager().named_hps(): + if configurable is None or hp.is_configurable == configurable: + yield hp_name, hp + + def extra_repr(self): + """Generate extra_repr making sure all dynamic keys exist in self.__dict__. + + Pytorch heavily uses self.__dict__ to generate extra_repr. However, we remove certain + attributes from self.__dict__ so we can manage them dynamically. Temporarily, adding them + back in here and removing them again afterwards. + """ + added_keys = set() + manager = self._get_dm_attribute_manager() + for name in chain(manager.hp_keys(), manager.da_keys()): + val = getattr(self, name) + if name in self.__dict__ or val is None: + continue + added_keys.add(name) + self.__dict__[name] = getattr(self, name) + + # make super call + extra_repr = super().extra_repr() + + # remove keys again + for key in added_keys: + del self.__dict__[key] + + return extra_repr + + @property + def original_cls(self) -> type[nn.Module]: + """Return the original class of the dynamic module.""" + return self._get_dm_attribute_manager().og_cls + + +class _DMRegistryCls: + """A registry to keep track of available dynamic modules. + + The registry can also dynamically generate new entries when we have a class that inherits from + the registered nn.Module. + """ + + T = TypeVar("T", bound=DynamicModule) + + def __init__(self, prefix: str): + super().__init__() + + self._prefix = prefix # global prefix to use to register dynamic classes + + self._registry: dict[type[nn.Module], type[DynamicModule]] = {} # registered classes + self._key_registry: dict[type[nn.Module], str] = {} # registered str-keys for classes + + self._dynamic_classes: dict[type[nn.Module], type[DynamicModule]] = {} # generated DMs + # generated ModeloptBaseRule classes for the dynamic classes + self._rule_classes: dict[type[DynamicModule], type[ModeloptBaseRule]] = {} + + def _generate_rule_class(self, dm_cls: type[DynamicModule]) -> type[ModeloptBaseRule]: + """Generate a ModeloptBaseRule type for the given dynamic module class.""" + spec = inspect.getfullargspec(dm_cls.modify) + if spec.kwonlyargs: + assert spec.kwonlydefaults is not None, "modify method should have default arguments!" + fields = {n: (spec.annotations[n], spec.kwonlydefaults[n]) for n in spec.kwonlyargs} + else: + fields = {} + return create_model(f"{dm_cls.__name__}Config", __base__=ModeloptBaseRule, **fields) + + def _get_dynamic_class_name(self, nn_cls: type[nn.Module]) -> str: + """Generate a name for the dynamic class.""" + + def _is_occupied(name: str) -> bool: + """Check if a name is already occupied by a dynamic class.""" + return any(name == dm.__name__ for dm in self._dynamic_classes.values()) + + # first try: simply prepend self._prefix to the nn_cls name but only if not occupied already + name = f"{self._prefix}{nn_cls.__name__}" + if not _is_occupied(name): + return name + + # 2nd try: add self._prefix and module name to the nn_cls name (should never be occupied!) + name = f"{nn_cls.__module__.replace('.', '_')}_{self._prefix}{nn_cls.__name__}" + assert not _is_occupied(name), f"Dynamic class name {name} should not be occupied!" + return name + + def _get_registered_nn_class(self, nn_cls: type[nn.Module]) -> Optional[type[nn.Module]]: + """Optionally return the nn module class that should be used to register nn_cls. + + This method loops through the registry to see if there are any nn_cls matches, i.e., any + subclass with a shared forward method. + """ + for nn_cls_ in self._registry: + if issubclass(nn_cls, nn_cls_) and nn_cls.forward is nn_cls_.forward: + return nn_cls_ + return None + + def __contains__(self, item: Union[nn.Module, type[nn.Module], str]): + if isinstance(item, str): + return item in self._key_registry.values() + nn_cls = type(item) if isinstance(item, nn.Module) else item + return self._get_registered_nn_class(nn_cls) is not None + + def __getitem__(self, nn_cls: Union[type[nn.Module], str]) -> type[DynamicModule]: + """Return the dynamic module class for the given nn_cls type or registered string.""" + dm_cls = self.get(nn_cls) + if dm_cls is None: + raise KeyError(f"{nn_cls} is not registered for a dynamic module!") + return dm_cls + + def get( + self, nn_cls: Union[type[nn.Module], str], default: Any = None + ) -> Optional[type[DynamicModule]]: + """Return the dynamic module class for the given nn_cls type or registered string.""" + # check string case first + if isinstance(nn_cls, str): + for nn_cls_, key in self._key_registry.items(): + if key == nn_cls: + return self.get(nn_cls_, default) + return default + + assert issubclass(nn_cls, nn.Module), f"{nn_cls} is not a subclass of nn.Module!" + + # see if dynamic class is already available and return + if nn_cls in self._dynamic_classes: + return self._dynamic_classes[nn_cls] + + # see if we can generate a new dynamic class and then return + nn_cls_ = self._get_registered_nn_class(nn_cls) + if nn_cls_: + name = self._get_dynamic_class_name(nn_cls) + dm_base_class = self._registry[nn_cls_] + dm_class = type(name, (dm_base_class, nn_cls), {}) + self._dynamic_classes[nn_cls] = dm_class + self._rule_classes[dm_class] = self._generate_rule_class(dm_class) + return self.get(nn_cls, default) + + # default return + return default + + def get_key_from_dm(self, dm_cls: Union[type[DynamicModule], DynamicModule]) -> str: + """Retrieve the key that is registered for a given dynamic module class.""" + dm_cls = type(dm_cls) if isinstance(dm_cls, DynamicModule) else dm_cls + for nn_cls, dm_cls_ in self._dynamic_classes.items(): + if dm_cls == dm_cls_: + return self.get_key(nn_cls) + raise KeyError(f"{dm_cls} is not registered for a dynamic module!") + + def get_key(self, nn_cls: Union[type[nn.Module], str]) -> str: + """Retrieve the key that is registered for a given nn module class.""" + # sanity check + assert nn_cls in self, f"{nn_cls} is not registered for a dynamic module!" + + # handle string case + if isinstance(nn_cls, str): + return nn_cls + + # handle module case + nn_cls_ = self._get_registered_nn_class(nn_cls) + assert nn_cls_ is not None + return self._key_registry[nn_cls_] + + def get_rule_class(self, nn_cls: Union[type[nn.Module], str]) -> type[ModeloptBaseRule]: + """Retrieve the rule config class that is registered for a given nn module class.""" + dm_cls = self.get(nn_cls) + if dm_cls is None: + raise KeyError(f"{nn_cls} is not registered for a dynamic module!") + return self._rule_classes[dm_cls] + + def register(self, cls_to_key: dict[type[nn.Module], str]) -> Callable[[type[T]], type[T]]: + """Use this to register a new dynamic base module. + + Usage: + + .. code-block:: python + + @DMRegistry.register({nn.Linear: "nn.Linear"}) + class DynamicLinear(DynamicModule): + pass + + .. note:: + + The dynamic base module must **NOT** inherit from the nn.Module that is registered. + Instead the registry will automatically generate a new class that inherits from both + the dynamic class (``DynamicLinear`` above) and the nn.Module (``nn.Linear`` above) for + an MRO that corresponds to ``class AutoGenerated(DynamicLinear, nn.Linear): pass``. + """ + + def decorator(dm_class: type[_DMRegistryCls.T]) -> type[_DMRegistryCls.T]: + """Register dnn_class with appropriate nn_class.""" + for nn_cls_, key in cls_to_key.items(): + assert nn_cls_ not in self._registry, f"{nn_cls_} already registered!" + self._registry[nn_cls_] = dm_class + self._key_registry[nn_cls_] = key + return dm_class + + return decorator + + def unregister(self, nn_cls: Union[type[nn.Module], type[T]]) -> None: + """Unregister a previously registered dynamic base module and all its inherited modules. + + It throws a KeyError if the dynamic base module is not registered. + """ + # 0. sanity check + if nn_cls not in self._registry: + raise KeyError(f"{nn_cls} is not registered!") + + # 1. unregister any generated dynamic classes + for nn_cls_ in list(self._dynamic_classes): + if nn_cls == self._get_registered_nn_class(nn_cls_): + dm_cls = self._dynamic_classes.pop(nn_cls_) + self._rule_classes.pop(dm_cls) + + # 2. unregister the base dynamic class + self._registry.pop(nn_cls) + self._key_registry.pop(nn_cls) + + def convert(self, nn_mod: nn.Module) -> DynamicModule: + """Converts the module into a dynamic module if registered or raise KeyError.""" + return self[type(nn_mod)].convert(nn_mod) + + @property + def prefix(self) -> str: + """Return the prefix used for the dynamic classes.""" + return self._prefix + + +class DynamicSpace: + """A class to represent all dynamic model choices over a model with multiple submodules.""" + + def __init__(self, model: nn.Module) -> None: + """Initialize the dynamic space from the model.""" + self.model = model + + def _should_be_converted(self, mod: nn.Module) -> bool: + """Check if the module should be converted.""" + return True + + def convert_to_dynamic( + self, rules: Optional[RulesDict], dm_registry: _DMRegistryCls + ) -> dict[str, nn.Module]: + """Convert the model to dynamic modules according to the rules and provided registry. + + Args: + rules: A dictionary containing rules for the dynamic modules. + dm_registry: A registry containing the dynamic modules to be converted to. + + Returns: + A dictionary containing the converted modules with submodule names as keys and the + converted dynamic modules as values. + """ + # check if the model is DataParallel + unwrap_model(self.model, raise_error=True) + + # check that it's not channels last + assert not is_channels_last(self.model) + + # 1. patch the model with dynamic units + # NOTE: duplicate modules are handled automatically since we convert/patch modules in-place + mods_converted: dict[str, nn.Module] = {} + for name, mod in self.model.named_modules(): + if mod in dm_registry and self._should_be_converted(mod): + dm_registry.convert(mod) + mods_converted[name] = mod + + # 2. modify search space if rules are provided + if rules is None: + return mods_converted + + # change all keys to strings + rules = {dm_registry.get_key(key): rule for key, rule in rules.items()} + + # iterator through all dynamic modules and modify them according to rules + for mod_name, mod in mods_converted.items(): + # get the key for this module that is used in the rules + key = dm_registry.get_key_from_dm(mod) + + # validate + construct custom rules using the rule class stored in the dm_registry + rule_custom = dm_registry.get_rule_class(key).customize_rule(rules.get(key), mod_name) + + if rule_custom is None: + # freeze module if no rule is provided + mod.freeze() + else: + # modify module according to rule + mod.modify(**rule_custom) + + return mods_converted + + def named_dynamic_modules(self) -> Iterator[tuple[str, DynamicModule]]: + """Recursively yield the name and instance of *all* DynamicModules. + + Yields: + (name, DynamicModule): tuple containing the name and module. + """ + for name, module in self.model.named_modules(): + if isinstance(module, DynamicModule): + yield name, module + + def is_dynamic(self) -> bool: + """Check if any module is dynamic. + + Returns: + True if the model contains DynamicModule(s). + """ + return any(True for _ in self.named_dynamic_modules()) + + def named_hparams(self, configurable: Optional[bool] = None) -> Iterator[tuple[str, Hparam]]: + """Recursively yield the name and instance of *all* hparams. + + Args: + configurable: Whether to include configurable hparams. + + Yields: + (name, Hparam): tuple containing the name and hparam. + + Default behavior is to iterate over all hparams. If ``configurable`` is set to ``True``, + only configurable, non-duplicate symbols are iterated over. + """ + _memo = set() + assert configurable in [None, True], "Only all or configurable hparams are supported!" + for mod_name, mod in self.named_dynamic_modules(): + for hp_name, hp in mod.named_hparams(configurable=configurable): + if configurable is None or hp not in _memo: + yield mod_name + ("." if mod_name else "") + hp_name, hp + _memo.add(hp) + + def get_hparam(self, name: str) -> Hparam: + """Get the hparam with the given name.""" + mod_name, _, hp_name = name.rpartition(".") + mod = self.model.get_submodule(mod_name) + assert isinstance(mod, DynamicModule), f"Module {mod} must be a DynamicModule!" + return mod.get_hparam(hp_name) + + def is_configurable(self) -> bool: + """Check if the model has any configurable hyperparameters. + + Args: + model: A model to be checked for DynamicModule(s) with configurable hyperparameters. + + Returns: + True if the model contains DynamicModule(s) with configurable hyperparameters w/ more + than one choice. + """ + return any(True for _ in self.named_hparams(configurable=True)) + + def size(self) -> int: + """Get the search space size of the model. + + Returns: + A int representing the search space size of the model. + """ + space_size = 1 + for _, hp in self.named_hparams(configurable=True): + space_size *= len(hp.choices) + return space_size + + def config(self, configurable: Optional[bool] = None) -> dict[str, Any]: + """Return the config dict of all hyperparameters. + + Args: + model: A model that contains DynamicModule(s). + configurable: None -> all hps, True -> configurable hps, False -> non-configurable hps + + Returns: + A dict of ``(parameter_name, choice)`` that specifies an active subnet. + """ + return { + get_unwrapped_name(name): hp.active for name, hp in self.named_hparams(configurable) + } + + def select(self, config: dict[str, Any], strict: bool = True) -> None: + """Select the subnet provided by config. + + If `strict` is set, then `config` must contain the exact set of keys representing both the + configurable and non-configurable hparams. + """ + configurables = dict(self.named_hparams(configurable=True)) + + # check if we have any overlap with dependent keys + check_non_configurable = any( + name in config and name not in configurables for name, hp in self.named_hparams() + ) + + # go through config, select based on current config, and track key errors + missing_keys = [] + inconsistent_keys = [] + unexpected_keys = {k: True for k in config.keys()} + + # assign free/searchable hparams from config + for name, hparam in configurables.items(): + if name in config: + hparam.active = config[name] + unexpected_keys[name] = False + elif strict: + missing_keys.append(name) + + # do a sanity check on the provided dynamic keys + if check_non_configurable and strict: + for name, hparam in self.named_hparams(): + if name in configurables: + continue + if name not in config: + missing_keys.append(name) + continue + unexpected_keys[name] = False + if hparam.active != config[name]: + inconsistent_keys.append( + f"{name}: active={hparam.active}, config={config[name]}" + ) + + # raise error for missing and unexpected keys if strict + unexpected_keys = [k for k, val in unexpected_keys.items() if val] + error_msg = "" + if strict and len(missing_keys) > 0: + error_msg += "\n\t".join(["Missing keys in config for:"] + missing_keys) + error_msg += "\nMake sure all keys are present in config or set strict=False." + raise RuntimeError(error_msg) + if strict and len(unexpected_keys) > 0: + error_msg += "\n\t".join(["Unexpected keys in config for:"] + unexpected_keys) + error_msg += "\nMake sure only required keys are in config or set strict=False." + if strict and len(inconsistent_keys) > 0: + error_msg += "\n\t".join(["Inconsistent keys in config:"] + inconsistent_keys) + error_msg += "\nMake sure keys in config are consistent or set strict=False." + if error_msg: + raise RuntimeError(f"Error in selecting config:\n{error_msg}") + + def export(self, dm_registry: _DMRegistryCls) -> nn.Module: + """Recursively export the module including self and return the result. + + Args: + dm_registry: A dynamic module registry to check for dynamic modules that should be + exported. + + Returns: + The model after exporting the dynamic modules found in the registry. + """ + + def _recursive_export(mod: nn.Module) -> None: + """Recursively export the module.""" + for n, m in mod.named_children(): + if isinstance(m, DynamicModule) and m.original_cls in dm_registry: + setattr(mod, n, m.export()) # re-assigning needed for DynamicParallelKDModule + _recursive_export(m) + + if isinstance(self.model, DynamicModule): + self.model = self.model.export() + _recursive_export(self.model) + + return self.model + + def __repr__(self): + lines = [f"{type(self).__name__}("] + spaces = " " + attr_line = [f"{x}={getattr(self, x)()}" for x in ["is_dynamic", "is_configurable", "size"]] + lines.append(spaces + ", ".join(attr_line) + ",") + model_lines = str(self.model).split("\n") + lines.extend([spaces + ("" if i else "model=") + x for i, x in enumerate(model_lines)]) + lines.append(")") + return "\n".join(lines) diff --git a/modelopt/torch/opt/hparam.py b/modelopt/torch/opt/hparam.py new file mode 100644 index 0000000000..302b4930a8 --- /dev/null +++ b/modelopt/torch/opt/hparam.py @@ -0,0 +1,261 @@ +# 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. + +"""Standard hyperparameter class for regular symbol.""" + +from abc import ABC, abstractmethod +from contextlib import contextmanager +from typing import Callable, Generator, Optional, Sequence, Union + +import torch + +__all__ = ["Hparam", "HPType"] + + +class CustomHPType(ABC): + """Custom hyperparameter type base class for user-defined hparam types.""" + + @abstractmethod + def __repr__(self) -> str: + """Return string representation with relevant properties of the class.""" + + @abstractmethod + def __lt__(self, other) -> bool: + """Return True if self is less than or equal to other.""" + + @abstractmethod + def __eq__(self, other) -> bool: + """Return True if self is equal to other.""" + + +HPType = Union[tuple[int, ...], int, float, CustomHPType] + + +# NOTE: Hparam class should be picklable for distributed broadcasting of hparams for Megatron pruning. +class Hparam: + """A base hyperparameter of a DynamicModule. + + An example of such a Hparam could be an hparam with identity dependencies. + """ + + Importance = Optional[torch.Tensor] + ImportanceEstimator = Callable[[], Importance] + ActiveSlice = Union[slice, torch.LongTensor] + + def __init__(self, choices: Sequence[HPType], original: Optional[HPType] = None) -> None: + """Initializes Hparam with original value and choices.""" + self._original = max(choices) if original is None else original + self._choices = sorted(set(choices) | {self.original}) + self._active = self.original + self._is_configurable = True # in case we want to manually overwritte configurability. + self._strict_len = True # whether the importance must be of length equal to max choice + + # Callback to compute hparam importance + self._importance_estimators: Optional[list[Hparam.ImportanceEstimator]] = [ + self._default_get_importance + ] + + # an optional order to enforce for active_slice + self._slice_order: Optional[torch.LongTensor] = None + + def __iter__(self) -> Generator[HPType, None, None]: + """Iterate over choices.""" + yield from self.choices + + @property + def is_configurable(self): + """Return whether the hparam is configurable.""" + return self._is_configurable and len(self.choices) > 1 + + @contextmanager + def _force_configurable(self): + """Context manager to temporarily set hparam to be configurable.""" + original_value = self._is_configurable + self._is_configurable = True + yield + self._is_configurable = original_value + + @property + def is_sortable(self): + """Return whether hparam in sortable.""" + return self._importance_estimators is not None + + @property + def active(self) -> HPType: + """Return the currently active value.""" + return self._active + + @active.setter + def active(self, val: Optional[HPType]): + """Set the active value with a sanity check for choices and dynamic hparams.""" + val = self.original if val is None else val + assert val in self._choices, f"val = {val}, choices = {self.choices}" + if self.is_configurable: + self._active = val + else: + assert self._active == val + + @property + def active_slice(self) -> ActiveSlice: + """Return the currently active sorted indices or slice corresponding to the active value.""" + assert isinstance(self.active, int), "active_slice only supported for int hparams" + if self._slice_order is not None: + return self._slice_order[: self.active] + return slice(self.active) + + @property + def choices(self) -> Sequence[HPType]: + """Return available choices.""" + return self._choices + + @choices.setter + def choices(self, val: Sequence[HPType]): + """Set available choices as subset of current choices (original must be preserved!).""" + # sanity checks + assert self.original in val, f"Original choice not in choices: {self.original} not in {val}" + assert self.active in val, f"Active choice not in choices: {self.active} not in {val}" + curr = set(self.choices) + val_set = set(val) + + # update local choices + if self.is_configurable: + assert val_set <= curr, f"New choices must be a subset: current: {curr}, new: {val_set}" + self._choices = sorted(val_set) + else: + assert val_set == curr, f"Cannot update choices: current {curr}, new: {val_set}" + + @property + def min(self) -> HPType: + """Return min value from among choices.""" + return min(self.choices) + + @property + def max(self) -> HPType: + """Return max value from among choices.""" + return max(self.choices) + + @property + def original(self) -> HPType: + """Return original value from among choices.""" + return self._original + + @property + @torch.no_grad() + def importance(self) -> Importance: + """Computes and returns the normalized importance among the features the hparam represents. + + Note that the importance is represented as a 1d-tensor with the length equal to the max + choice (Hparam.max) of the hparam. + + For example, if the hparam represents the number of in_channels to a Conv2d layer, the + importance should be a 1d-tensor of importance score with length equal to the number of + in_channels. + + Note that each module should register appropriate importance callbacks to compute the + actual importance associated with the hparam choices. If there is no notion of importance + for the hparam, this function returns None. + """ + # TODO: will we ever need a cycle detector here? + assert self.is_configurable + return self._get_importance() if self.is_sortable else None + + # TODO: eventually importance should be registered by the search algorithm and _not_ by the + # module implementation. This is a temporary solution to get things working... + def register_importance(self, importance_estimator: ImportanceEstimator): + """Register importance estimator for the hparam. + + This estimator does not take any arguments and should return a single argument ( + optional 1d-tensor) representing the importance among features the hparam represents. + If the return argument is a tensor, the length of the tensor must be equal to the max choice + (Hparam.max) of the hparam. + """ + if self._importance_estimators is None: # use instead of self.is_sortable for mypy + raise RuntimeError("Cannot register importance for unsortable hparams.") + else: + self._importance_estimators.append(importance_estimator) + + def _default_get_importance(self) -> Importance: + """Default importance estimator.""" + return None + + def _get_importance(self) -> Importance: + """Private method to retrieve importance without sanity checks.""" + # iterate though all importance estimators + imps_all = [] + for estimator in self._importance_estimators or []: + imp = estimator() + if imp is not None: + assert ( + not self._strict_len or len(imp) == self.max + ), "Length of importance must be equal to max choice!" + imp = imp / (imp.max() + 1e-9) # normalize importance + imps_all.append(imp) + + return sum(imps_all) if imps_all else None + + @torch.no_grad() + def enforce_order(self, order: Optional[torch.Tensor] = None) -> None: + """Store a reference to this order and enforce the order for active_slice. + + This function enables the user to enforce an order how the active_slice is generated. + + Example: + If the hparam has active value 16 and the max value is 32, the active_slice by default + will be ``slice(16)``, which is equivalent to ``range(16)`` (although faster). + When order is set active_slice will instead return ``self._order[:16]``. + + TODO: will we ever need a cycle detector here? + """ + if order is not None: + # convert to long (torch's index type) + order = order.long() + + # check if the order is valid + assert not self._strict_len or torch.equal( + torch.arange(self.max, device=order.device), torch.sort(order)[0] + ), "order must be a permutation of range(self.max) to be valid!" + + self._enforce_order(order) + + def _enforce_order(self, order: Optional[torch.Tensor] = None) -> None: + """Private method to enforce order without sanity checks.""" + if isinstance(order, torch.Tensor): # Always store the order as a CPU tensor + order = order.cpu() + self._slice_order = order + + def __repr__(self) -> str: + """Return string representation with relevant properties of the class.""" + attrs = ["choices", "active", "original"] + return f"{type(self).__name__}({', '.join(f'{x}={getattr(self, x)}' for x in attrs)})" + + def __iand__(self, hp: "Hparam"): + """Merge another hparam into self.""" + assert isinstance(hp, (Hparam)), f"Cannot merge {type(hp)} into {type(self)}!" + + # merge choices + self.choices = sorted(set(self.choices) & set(hp.choices)) + + # remove slice order + self._slice_order = None + + # merge importance estimators + if isinstance(self._importance_estimators, list) and isinstance( + hp._importance_estimators, list + ): + self._importance_estimators.extend(hp._importance_estimators) + else: + self._importance_estimators = None + + return self diff --git a/modelopt/torch/opt/mode.py b/modelopt/torch/opt/mode.py new file mode 100644 index 0000000000..4bd27a6c64 --- /dev/null +++ b/modelopt/torch/opt/mode.py @@ -0,0 +1,313 @@ +# 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. + +"""Interface and utilities for optimization modes/algorithms. + +A mode is a specific type or algorithms for model optimization, e.g., some type of algorithm for +pruning or quantization. It can also specify a single step within an optimization algorithm instead +of the whole algorithm. For example, a mode can prepare a model for pruning or export (i.e. fix the +optimal model configuration) after pruning. + +Within ``modelopt``, a ``mode`` constitutes the unit for model optimization. We can define arbitrary +modes, each mode gets recorded in the model's modelopt state dict, and we can define workflows as a +sequence of modes. +""" + +from abc import ABC, abstractmethod +from typing import Any, Callable, Optional, TypeVar, Union + +import torch.nn as nn + +from modelopt.torch.utils import val2list + +from .config import ConfigDict, ModeloptBaseConfig +from .searcher import BaseSearcher + +MetadataDict = dict[str, Any] # metadata dict for one mode +ModeConfigList = list[tuple[str, ConfigDict]] # config list for multiple modes +ModeState = dict[str, Union[ConfigDict, MetadataDict]] # state dict for one mode + +ModeEntrypoint = Callable[ + [nn.Module, ModeloptBaseConfig, MetadataDict], tuple[nn.Module, MetadataDict] +] +ConvertReturnType = tuple[nn.Module, MetadataDict] +ConvertEntrypoint = Callable[[nn.Module, ModeloptBaseConfig], ConvertReturnType] +RestoreEntrypoint = Callable[[nn.Module, ModeloptBaseConfig, MetadataDict], nn.Module] +UpdateEntrypoint = Callable[[nn.Module, ModeloptBaseConfig, MetadataDict], None] + +__all__ = [] + + +class _ModeDescriptor(ABC): + """Abstract class to describe a mode.""" + + def __str__(self) -> str: + return str(self.name) + + def __repr__(self) -> str: + return f"<{type(self).__name__}: '{self.name}'>" + + def __hash__(self): + return hash(self.name) + + @property + @abstractmethod + def name(self) -> str: + """Returns the string name of the mode.""" + + @property + @abstractmethod + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + + @property + def next_modes(self) -> Optional[set[str]]: + """Modes that must immediately follow this mode. + + Certain modes only makes sense if they are followed by certain other modes. + + An empty set indicates that _no_ mode can follow this mode. A None value indicates that + there are no restrictions on the following mode. + + Returns: + A set of mode names that must immediately follow this mode. Defaults to None. + """ + return None + + @property + def export_mode(self) -> Optional[str]: + """The mode that corresponds to the export mode of this mode. + + Certain modes require a subsequent export step. For example, after pruning, we might want to + fine-tune the model and then export the model. This property specifies that mode if it + exists. + + None indicates that there exists no such mode. + + Returns: + The (optional) mode name that corresponds to the export mode of this mode. Defaults to + None. + """ + return None + + @property + def is_export_mode(self) -> bool: + """Whether the mode is an export mode. + + Returns: + True if the mode is an export mode, False otherwise. Defaults to False. + """ + return False + + @property + def search_algorithm(self) -> type[BaseSearcher]: + """Specifies the search algorithm to use for this mode (if any).""" + raise RuntimeError("Search is not supported for this mode.") + + @property + @abstractmethod + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model. + + The function signature of the convert entrypoint is described below: + + Args: + model: Model to be restored. + config: Config used for the model that was also used during convert. + + Returns: + A tuple consisting of + 1. the in-place modified model. If the modification failed, the entrypoint can + return None instead + 2. The config dict that can be used to call the restore entrypoint to instantly *restore* + the modified model. + 3. The metatdata that can be used to call the restore entrypoint to instantly + *restore* the modified model from the provided initial state, see below's + description for the restore entrypoint to get more info about ``metadata``. + + Raises: + :meth:`ApplyModeError` to indicate that the + conversion process failed. This error can be caught by user-facing APIs if they want to + enable a fall-back behavior. + """ + + @property + @abstractmethod + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model. + + The function signature of the restore entrypoint is described below: + + Args: + model: Model to be restored. + config: Config used for the model that was also used during convert. + metadata: The metadata is used during restoration of the model architecture to instantly + restore the modified model. The metadata is used on top of the config to ensure that + the model can be instantly restored/modified from the provided state. This is + helpful when the ``convert`` entrypoint contains non-deterministic operations whose + outcome can be stored in the metadata to ensure that the model can be restored + reliably. A few examples of potential non-deterministic operations are provided + below: + * Latency measurements: if the conversion leverages latency measurements during + conversion the conversion process may become non-deterministic. + * Random operations: if the conversion leverages random operations during + conversion, we should store the samples or random seed. + * Module's train flag: the conversion process might be affected by the module's + train flag (e.g. tracing is indirectly affected by train flag since the + forward may be affected by the train flag). If so, we should store the train + flag in the metadata and set the model into the correct mode. + + Returns: + The in-place modified and restored model. If the modification failed, the entrypoint can + return None instead + + Raises: + :meth:`ApplyModeError` to indicate that the + conversion process failed. This error can be caught by user-facing APIs if they want to + enable a fall-back behavior. + """ + + @property + def update_for_save(self) -> UpdateEntrypoint: + """The mode's (optional) entrypoint for updating a model's config and metadata before saving. + + This is useful if metadata or config needs to be updated for saving (and restoring) the mode. + + The function signature of this update entrypoint is described below: + + Args: + model: Model to be restored. + config: The config as described above. It should be modified IN-PLACE. + metadata: The metadata as described above. It should be modified IN-PLACE. + + Returns: + None. + """ + return lambda model, config, metadata: None + + @property + def update_for_new_mode(self) -> UpdateEntrypoint: + """The mode's (optional) entrypoint for updating a model's config and metadata before a new mode. + + This is useful if metadata or config needs to be updated before adding a new mode. For example, a + after adding a new mode, the current mode's restore might only need a subset of the metadata/config. + + The function signature of this update entrypoint is described below: + + Args: + model: Model to be restored. + config: The config as described above. It should be modified IN-PLACE. + metadata: The metadata as described above. It should be modified IN-PLACE. + + Returns: + None. + """ + return lambda model, config, metadata: None + + @property + def require_model_like(self) -> bool: + """Whether the mode requires a ModelLike input? + + Returns: + False + """ + return False + + +ModeType = Union[_ModeDescriptor, str] +ModeLike = Union[ModeType, list[ModeType], ModeConfigList] + + +class _ModeRegistryCls: + """A registry to keep track of available modes.""" + + T = TypeVar("T", bound=_ModeDescriptor) + + # global list to keep track of all registries we initialize + _all_registries: list["_ModeRegistryCls"] = [] + + def __init__(self) -> None: + """Initialize the registry with the lookup dictionaries.""" + self._name2descriptor: dict[str, _ModeDescriptor] = {} + self._all_registries.append(self) + + def register_mode(self, cls_descriptor: type[T]) -> type[T]: + """Register a new mode with the given descriptor.""" + # initialize descriptor and get name + descriptor = cls_descriptor() + name = descriptor.name + + # check if we have a descriptor instance already and use that instance instead. + if self.contained_in_any(name): + descriptor = self.get_from_any(name) + + # check if mode_name/value is already taken + if name in self._name2descriptor: + raise ValueError(f"Mode {name} already registered: {self._name2descriptor}") + + # register mode + self._name2descriptor[name] = descriptor + return cls_descriptor + + def remove_mode(self, mode: ModeType) -> None: + """Remove a mode from the registry.""" + # remove mode + del self._name2descriptor[str(mode)] + + def get(self, mode: ModeType) -> Optional[_ModeDescriptor]: + """Get the mode by value or throw an error.""" + return self._name2descriptor.get(str(mode)) + + def __getitem__(self, mode: ModeType) -> _ModeDescriptor: + """Get the mode by value or throw an error.""" + return self._name2descriptor[str(mode)] + + def __contains__(self, mode: ModeType) -> bool: + """Check if mode is registered in this registry.""" + return str(mode) in self._name2descriptor + + def __del__(self) -> None: + """Remove the registry from the global list.""" + self._all_registries.remove(self) + + @classmethod + def contained_in_any(cls, mode: ModeType) -> bool: + """Check if mode is registered in any registry.""" + for registry in cls._all_registries: + if str(mode) in registry._name2descriptor: + return True + return False + + @classmethod + def get_from_any(cls, mode: ModeType) -> _ModeDescriptor: + """Get the mode by value from any registry or throw a KeyError. + + Adds a sanity check to ensure that the mode is not ambiguous, i.e., there is only one + instance. + """ + mode_ds = [registry[mode] for registry in cls._all_registries if mode in registry] + if not mode_ds: + raise KeyError(f"Mode {mode} not found in any registry.") + assert all(mode_ds[0] == m_d for m_d in mode_ds), f"Mode {mode} is ambiguous." + return mode_ds[0] + + +def get_mode_config(mode_like: ModeLike) -> ModeConfigList: + """Standardize mode to ModeConfigDict and return.""" + mode_and_config = [ + ((m, {}) if isinstance(m, str) else (m[0], m[1] or {})) for m in val2list(mode_like) + ] + + return mode_and_config diff --git a/modelopt/torch/opt/plugins/__init__.py b/modelopt/torch/opt/plugins/__init__.py new file mode 100644 index 0000000000..ac36784308 --- /dev/null +++ b/modelopt/torch/opt/plugins/__init__.py @@ -0,0 +1,32 @@ +# 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. + +"""Handles plugins for third-party modules.""" + +from modelopt.torch.utils import import_plugin + +from .huggingface import * + +with import_plugin("megatron core dist checkpointing"): + from .mcore_dist_checkpointing import * + +with import_plugin("transformers"): + from .transformers import * + +with import_plugin("diffusers"): + from .diffusers import * + +with import_plugin("peft"): + from .peft import * diff --git a/modelopt/torch/opt/plugins/diffusers.py b/modelopt/torch/opt/plugins/diffusers.py new file mode 100644 index 0000000000..9040d05658 --- /dev/null +++ b/modelopt/torch/opt/plugins/diffusers.py @@ -0,0 +1,24 @@ +# 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. + +"""ModelOpt plugin for enabling automatic save/restore of ModelOpt state for `diffusers` library.""" + +from diffusers import ModelMixin + +from .huggingface import patch_pretrained_methods + +__all__ = [] + +patch_pretrained_methods(ModelMixin, "diffusers") diff --git a/modelopt/torch/opt/plugins/huggingface.py b/modelopt/torch/opt/plugins/huggingface.py new file mode 100644 index 0000000000..a632c59e01 --- /dev/null +++ b/modelopt/torch/opt/plugins/huggingface.py @@ -0,0 +1,192 @@ +# 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. + +"""ModelOpt plugin for enabling automatic save/restore of ModelOpt state for HuggingFace models.""" + +import functools +import os +import types +from contextlib import contextmanager +from typing import Any + +import torch + +from modelopt.torch.utils import print_rank_0 + +from ..conversion import ( + ModeloptStateManager, + has_nas_modelopt_state, + modelopt_state, + restore_from_modelopt_state, +) + +__all__ = ["enable_huggingface_checkpointing"] + +_PATCHED_LIBRARIES = set() +_MODELOPT_STATE_SAVE_NAMES = { + "default": "modelopt_state.pth", +} + + +def _get_modelopt_state_path(obj: Any, model_name_or_path: str) -> str: + modelopt_state_save_name = _MODELOPT_STATE_SAVE_NAMES["default"] + for type, value in _MODELOPT_STATE_SAVE_NAMES.items(): + if type == "default": + continue + if isinstance(obj, type): # type: ignore [arg-type] + modelopt_state_save_name = value + return os.path.join(model_name_or_path, modelopt_state_save_name) + + +@contextmanager +def _patch_model_init_for_modelopt(cls, model_path): + """Patch for `cls.init` method to restore ModelOpt state after `init`.""" + cls._original__init__ = cls.__init__ + + @functools.wraps(cls._original__init__) + def new_init_fn(self, *args, **kwargs): + modelopt_state_path = _get_modelopt_state_path(self, model_path) + cls._original__init__(self, *args, **kwargs) + if os.path.isfile(modelopt_state_path): + modelopt_state = torch.load(modelopt_state_path, map_location="cpu") + if not has_nas_modelopt_state(modelopt_state): + restore_from_modelopt_state(self, modelopt_state) + print_rank_0(f"Restored ModelOpt state from {modelopt_state_path}") + + cls.__init__ = new_init_fn + try: + yield + finally: + cls.__init__ = cls._original__init__ + delattr(cls, "_original__init__") + + +def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs): + """Patch for `cls.from_pretrained` method to restore ModelOpt state. + + NOTE: This does not work with NAS/Pruned model in Bert example probably due to using accelerate. + (Error during tracing when restoring since __init__ is called under a context manager that disables weight init). + Hence, we defer the restoration to the `_load_pretrained_model` method. + """ + with _patch_model_init_for_modelopt(cls, pretrained_model_name_or_path): + model = types.MethodType(cls._modelopt_cache["from_pretrained"].__func__, cls)( + pretrained_model_name_or_path, *args, **kwargs + ) + return model + + +def _new_from_config(cls, /, config, **kwargs): + """Patch for `cls.from_config` method to restore ModelOpt state.""" + with _patch_model_init_for_modelopt(cls, config._name_or_path): + model = types.MethodType(cls._modelopt_cache["_from_config"].__func__, cls)( + config, **kwargs + ) + return model + + +def _new__load_pretrained_model( + cls, + /, + model, + state_dict, + loaded_keys, + resolved_archive_file, + pretrained_model_name_or_path, + *args, + **kwargs, +): + """Patch for `cls._load_pretrained_model` method to restore ModelOpt state for NAS/Pruned models. + + NOTE: Since ths is a lower level method, make sure to check if the order of arguments is ever changed! + """ + modelopt_state_path = _get_modelopt_state_path(model, pretrained_model_name_or_path) + if os.path.isfile(modelopt_state_path) and not ModeloptStateManager.is_converted(model): + modelopt_state_dict = torch.load(modelopt_state_path, map_location="cpu") + restore_from_modelopt_state(model, modelopt_state_dict) + print_rank_0(f"Restored ModelOpt state after init from {modelopt_state_path}") + + return types.MethodType(cls._modelopt_cache["_load_pretrained_model"].__func__, cls)( + model, + state_dict, + loaded_keys, + resolved_archive_file, + pretrained_model_name_or_path, + *args, + **kwargs, + ) + + +def _new_save_pretrained(self, save_directory, *args, **kwargs): + self._modelopt_cache["save_pretrained"](self, save_directory, *args, **kwargs) + if ModeloptStateManager.is_converted(self): + path = _get_modelopt_state_path(self, save_directory) + torch.save(modelopt_state(self), path) + print_rank_0(f"Saved ModelOpt state to {path}") + + +_DEFAULT_PATCH_METHODS_MAP = { + "from_pretrained": classmethod(_new_from_pretrained), + "_load_pretrained_model": classmethod(_new__load_pretrained_model), + "save_pretrained": _new_save_pretrained, + "_from_config": classmethod(_new_from_config), +} + + +def patch_pretrained_methods(cls, library_name: str, patch_methods_map: dict = None): + if hasattr(cls, "_modelopt_cache"): + return + cls._modelopt_cache = {} + patch_methods_map = patch_methods_map or _DEFAULT_PATCH_METHODS_MAP + for method_name in patch_methods_map: + if not hasattr(cls, method_name): + continue + cls._modelopt_cache[method_name] = getattr(cls, method_name) + setattr(cls, method_name, patch_methods_map[method_name]) + + _PATCHED_LIBRARIES.add(library_name) + + +def enable_huggingface_checkpointing(): + """Enables automatic save/restore of ModelOpt state with HuggingFace checkpointing APIs. + + ModelOpt automatically saves `modelopt_state` to `save_directory/modelopt_state.pth` when + a Huggingface model is saved using + `model.save_pretrained(save_directory) `_. + + Conversely, ModelOpt restores the saved state from `pretrained_model_name_or_path/modelopt_state.pth` if it exists + when a Huggingface model is loaded using + `cls.from_pretrained(pretrained_model_name_or_path) `_. + + + This function should be called once in the program before loading/saving any HuggingFace models. + + Here is an example usage: + + .. code-block:: python + + from transformers import AutoModelForCausalLM + import modelopt.torch.opt as mto + + # Enable ModelOpt save/restore for HuggingFace models + # This only needs to be called once in the program. + mto.enable_huggingface_checkpointing() + + # Instantiate a HuggingFace model, modelopt_state will be automatically loaded if it exists. + model = AutoModelForCausalLM.from_pretrained(model_path).cuda() + + """ + # This method simply prints if ModelOpt save/restore is enabled for the HuggingFace libraries. + for library_name in _PATCHED_LIBRARIES: + print_rank_0(f"ModelOpt save/restore enabled for `{library_name}` library.") diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py new file mode 100644 index 0000000000..cb609380c0 --- /dev/null +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -0,0 +1,412 @@ +# 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. + +"""Megatron core distributed checkpointing plugin for sharded ``modelopt_state``.""" + +# TODO: Add unit tests for this plugin +import copy +import os +from pathlib import Path +from typing import Any, Optional, Union + +import torch +from megatron.core import dist_checkpointing, mpu +from megatron.core.models.gpt import GPTModel +from megatron.core.models.mamba import MambaModel +from megatron.core.parallel_state import ( + get_pipeline_model_parallel_rank, + get_pipeline_model_parallel_world_size, +) +from megatron.core.transformer.module import Float16Module +from megatron.core.transformer.utils import make_sharded_object_for_checkpoint + +import modelopt.torch.opt as mto +import modelopt.torch.utils.distributed as dist +from modelopt.torch.utils.network import SUPPORTED_WRAPPERS + +SUPPORTED_MODELS = [GPTModel, MambaModel] +SUPPORTED_WRAPPERS[Float16Module] = "module" + + +def _remap_quantizer_state_with_prefix( + local_state: dict[str, Any], prefix: str, quantizer_state: dict[str, Any] +) -> None: + """Remap quantizer_state using the given prefix. + + The prefix is usually the local state_dict prefix. The remapped local_state + will be stored as a ShardedObject with its global state_dict prefix attached. + For quantizer_state, we further remove _amax and _pre_quant_scale which are + stored as a part of the model sharded_state_dict. + + Args: + local_state: the local quantizer_state to be sharded that all has the same local prefix + prefix: the quantizer_state prefix to extract + quantizer_state: all local quantizer_state + """ + prefix_with_dot = prefix + "." + + for key, val in quantizer_state.items(): + if not key.startswith(prefix_with_dot): + continue + new_key = key[len(prefix) :] + if "_amax" in val: + val.pop("_amax") + if "_pre_quant_scale" in val: + val.pop("_pre_quant_scale") + if new_key in local_state: + local_state[new_key]["quantizer_state"] = val + else: + local_state[new_key] = {"quantizer_state": val} + + +def _remap_sparsity_state_with_prefix( + local_state: dict[str, Any], + prefix: str, + subnet_config: dict[str, Any], +) -> None: + """Remap sparsity subnet_config using the given prefix. + + The prefix is usually the local state_dict prefix. The remapped local_state + will be stored as a ShardedObject with its global state_dict prefix attached. + + Args: + local_state: the local subnet_config to be sharded that all has the same local prefix + prefix: the subnet_config prefix to extract + subnet_config: all local subnet_config for sparsity + """ + prefix_with_dot = prefix + "." + + for key, val in subnet_config.items(): + if not key.startswith(prefix_with_dot): + continue + new_key = key[len(prefix) :] + if new_key in local_state: + local_state[new_key]["subnet_config"] = val + else: + local_state[new_key] = {"subnet_config": val} + + +def remove_modelopt_state_metadata( + modelopt_state: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + """Remove metadata from the modelopt_state. + + The metadata of the modelopt_state contains keys which may change with different pipeline + parallelism. As a result, the metadata must be stored as several ShardedObject with + global and local layer offset mapping. + + Args: + modelopt_state: the state_dict that contains all algorithms have have been applied + to the given model. + """ + quantizer_state = {} + subnet_config = {} + # TODO (chenhany): the assumption here is that there is at most one quantize and sparsity mode + if "modelopt_state_dict" in modelopt_state: + for mode, config in modelopt_state["modelopt_state_dict"]: + metadata = config.pop("metadata") if "metadata" in config else {} + if "quantizer_state" in metadata: + quantizer_state = metadata.get("quantizer_state") + if "subnet_config" in metadata: + subnet_config = metadata.get("subnet_config") + + return quantizer_state, subnet_config + + +def restore_modelopt_state_metadata(sharded_modelopt_state: dict[str, Any]) -> dict[str, Any]: + """Restore the per rank modelopt_state after the metadata is loaded. + + Args: + sharded_modelopt_state: the state_dict that contains all algorithms have have been applied + to the given model. + """ + modelopt_state_dict = sharded_modelopt_state.pop("modelopt_state_dict") + modelopt_version = sharded_modelopt_state.pop("modelopt_version") + quantizer_state = {} + subnet_config = {} + + for key, val in sharded_modelopt_state.items(): + for subkey, subval in val.items(): + if "quantizer_state" in subval: + quantizer_state[key + subkey] = subval["quantizer_state"] + if "subnet_config" in subval: + subnet_config[key + subkey] = subval["subnet_config"] + + for mode, config in modelopt_state_dict: + if mode == "quantize": + config["metadata"] = {"quantizer_state": quantizer_state} + elif mode == "sparsegpt": + config["metadata"] = {"subnet_config": subnet_config} + else: + config["metadata"] = {} + + return {"modelopt_state_dict": modelopt_state_dict, "modelopt_version": modelopt_version} + + +def _get_gpt_sharded_modelopt_state( + num_layers: int = -1, + num_medusa_heads: int = 0, + num_eagle_layers: int = 0, + model: Optional[torch.nn.Module] = None, + prefix: str = "", +) -> dict[str, Any]: + """Return the sharded modelopt_state for a GPTModel. + + If a GPTModel is not provided, then the sharded modelopt_state will still return a + dictionary of ShardedObject. This is used to load the sharded modelopt_state before + the initialization of a GPTModel. + + Args: + num_layers: number of decoder layers in the GPTModel + num_medusa_heads: number of Medusa heads in the GPTModel + num_eagle_layers: number of Eagle layers in the GPTModel + model: optionally provide a GPTModel instance + prefix: the prefix to add to the modelopt_state keys + """ + if model is not None: + num_layers = model.config.num_layers + if hasattr(model, "medusa_heads") and model.medusa_heads is not None: + num_medusa_heads = len(model.medusa_heads) + if hasattr(model, "eagle_module") and model.eagle_module is not None: + num_eagle_layers = len(model.eagle_module.decoder.layers) + elif num_layers < 0: + raise ValueError("Either num_layers or a model instance must be provided!") + + modelopt_state = {} if model is None else copy.deepcopy(mto.modelopt_state(model)) + quantizer_state, subnet_config = remove_modelopt_state_metadata(modelopt_state) + + # The sharded modelopt_state remains the part that is shared across all DP, TP, PP ranks. + sharded_modelopt_state = modelopt_state + sharded_offsets = [] + + # Compute per pp rank num_layers and global_layer_offset + local_num_layers = num_layers // get_pipeline_model_parallel_world_size() + global_layer_offset = local_num_layers * get_pipeline_model_parallel_rank() + + # First pp stage + if get_pipeline_model_parallel_rank() == 0: + local_key = "embedding" + global_key = f"{prefix}embedding" + local_state = {} + _remap_quantizer_state_with_prefix(local_state, local_key, quantizer_state) + _remap_sparsity_state_with_prefix(local_state, local_key, subnet_config) + sharded_modelopt_state[local_key] = make_sharded_object_for_checkpoint( + local_state, + global_key, + sharded_offsets, + ) + + # Last pp stage + if get_pipeline_model_parallel_rank() == (get_pipeline_model_parallel_world_size() - 1): + local_key = "output_layer" + global_key = f"{prefix}output_layer" + local_state = {} + _remap_quantizer_state_with_prefix(local_state, local_key, quantizer_state) + _remap_sparsity_state_with_prefix(local_state, local_key, subnet_config) + sharded_modelopt_state[local_key] = make_sharded_object_for_checkpoint( + local_state, + global_key, + sharded_offsets, + ) + + local_key = "medusa_heads" + global_key = f"{prefix}medusa_heads" + local_state = {} + _remap_quantizer_state_with_prefix(local_state, local_key, quantizer_state) + _remap_sparsity_state_with_prefix(local_state, local_key, subnet_config) + if num_medusa_heads > 0: + sharded_modelopt_state[local_key] = make_sharded_object_for_checkpoint( + local_state, + global_key, + sharded_offsets, + ) + + local_key = "eagle_module" + global_key = f"{prefix}eagle_module" + local_state = {} + _remap_quantizer_state_with_prefix(local_state, local_key, quantizer_state) + _remap_sparsity_state_with_prefix(local_state, local_key, subnet_config) + if num_eagle_layers > 0: + sharded_modelopt_state[local_key] = make_sharded_object_for_checkpoint( + local_state, + global_key, + sharded_offsets, + ) + + # Each pp rank owns some stages + for local_layer_id in range(local_num_layers): + global_layer_id = global_layer_offset + local_layer_id + local_key = f"decoder.layers.{local_layer_id}" + global_key = f"{prefix}decoder.layers.{global_layer_id}" + local_state = {} + _remap_quantizer_state_with_prefix(local_state, local_key, quantizer_state) + _remap_sparsity_state_with_prefix(local_state, local_key, subnet_config) + sharded_modelopt_state[local_key] = make_sharded_object_for_checkpoint( + local_state, + global_key, + sharded_offsets, + ) + + return sharded_modelopt_state + + +def get_sharded_modelopt_state( + num_layers: int, + model: torch.nn.Module, + prefix: str = "", + num_medusa_heads: int = 0, + num_eagle_layers: int = 0, +) -> dict[str, Any]: + """Return the sharded modelopt_state. + + Most of the modelopt_state is shared across all DP, TP, PP ranks. The metadata of each + mode, however, is pipeline-parallelism (PP) dependent. The major reason is because the + state is using the module name which contains decoder layer id that may change + when pipeline-parallelism is changing. As a result, we need to store metadata + with ShardedObject which provides the functionality to map the local state to + the global state. For example, ``embedding.`` only exists in the first PP stage, and + the same ``decoder.layers.{id}.`` module name can be found on all PP stages, but they + are mapped to different global layer id. + + For example, for a quantizer_state metadata, ``{"decoder.layers.3.mlp.linear_fc1.input_quantizer": state}``, + we need to store it as the following which maps local key ``decoder.layers.3`` to its global key + ``decoder.layers.11`` because of PP: + + .. code-block:: python + + sharded_modelopt_state["decoder.layers.3"] = ShardedObject( + {"mlp.linear_fc1.input_quantizer": {"quantizer_state": state}}, + global_key="decoder.layers.11", + ) + + To restore the metadata, we simply revert the process. + + Args: + num_layers: number of decoder layers in the MCore model + model: a MCore model instance + prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) + num_medusa_heads: number of Medusa heads + """ + support_list = SUPPORTED_MODELS + + if any([isinstance(model, arch) for arch in support_list]): + return _get_gpt_sharded_modelopt_state( + num_layers=num_layers, + model=model, + prefix=prefix, + num_medusa_heads=num_medusa_heads, + ) + else: + raise ValueError( + f"get_sharded_modelopt_state() only supports {SUPPORTED_MODELS}, received {type(model)}" + ) + + +def save_modelopt_state(model: list[torch.nn.Module], state_dict: dict[str, Any]) -> None: + """Save modelopt_state as a part of the per rank state_dict. + + NOTE: Only used for Megatron-LM. + + Args: + model: the modelopt optimized model + state_dict: the current modelopt optimized model state_dict to store + """ + if not mto.ModeloptStateManager.is_converted(model[0]): + return + if len(model) == 1: + state_dict["modelopt_state"] = mto.modelopt_state(model[0]) + else: + for i in range(len(model)): + mpu.set_virtual_pipeline_model_parallel_rank(i) + state_dict[f"modelopt_state_{i}"] = mto.modelopt_state(model[i]) + + +def restore_modelopt_state(model: list[torch.nn.Module], state_dict: dict[str, Any]) -> None: + """Restore modelopt_state from the per rank state_dict. + + NOTE: Only used for Megatron-LM. + + Args: + model: the model to restore the modelopt optimization + state_dict: the loaded state_dict to extract + """ + if ( + len(model) == 1 + and "modelopt_state" in state_dict + and not mto.ModeloptStateManager.is_converted(model[0]) + ): + model[0] = mto.restore_from_modelopt_state(model[0], state_dict["modelopt_state"]) + else: + for i in range(len(model)): + mpu.set_virtual_pipeline_model_parallel_rank(i) + if f"modelopt_state_{i}" in state_dict and not mto.ModeloptStateManager.is_converted( + model[i] + ): + model[i] = mto.restore_from_modelopt_state( + model[i], state_dict[f"modelopt_state_{i}"] + ) + + +def save_sharded_modelopt_state( + model: list[torch.nn.Module], + checkpoint_name: Union[str, Path], + sharded_strategy: Optional[tuple[str, int]] = None, + prefix: str = "", +) -> None: + """Save modelopt_state in the sharded state_dict format. + + Args: + model: the model to restore the modelopt optimization + checkpoint_name: the checkpoint folder path + sharded_strategy: configures sharded tensors saving behavior and backend + prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) + """ + if not mto.ModeloptStateManager.is_converted(model[0]): + return + if len(model) > 1: + raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") + modelopt_checkpoint_name = f"{checkpoint_name}/modelopt_state" + if dist.is_master(): + os.makedirs(modelopt_checkpoint_name, exist_ok=True) + + dist_checkpointing.save( + get_sharded_modelopt_state(-1, model[0], prefix), + modelopt_checkpoint_name, + sharded_strategy, + ) + + +def restore_sharded_modelopt_state( + model: list[torch.nn.Module], checkpoint_name: Union[str, Path], prefix: str = "" +) -> None: + """Restore modelopt_state from the sharded state_dict format. + + Args: + model: the model to restore the modelopt optimization + checkpoint_name: the checkpoint folder path + prefix: the prefix to add to the modelopt_state keys ("model." for NeMo) + """ + modelopt_checkpoint_name = f"{checkpoint_name}/modelopt_state" + if os.path.exists(modelopt_checkpoint_name): + if len(model) > 1: + raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!") + if not mto.ModeloptStateManager.is_converted(model[0]): + modelopt_state = restore_modelopt_state_metadata( + dist_checkpointing.load( + get_sharded_modelopt_state(-1, model[0], prefix), + modelopt_checkpoint_name, + ) + ) + model[0] = mto.restore_from_modelopt_state(model[0], modelopt_state) diff --git a/modelopt/torch/opt/plugins/peft.py b/modelopt/torch/opt/plugins/peft.py new file mode 100644 index 0000000000..bf65eb77bc --- /dev/null +++ b/modelopt/torch/opt/plugins/peft.py @@ -0,0 +1,100 @@ +# 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. + +"""ModelOpt plugin for enabling automatic save/restore of ModelOpt state for `peft` library.""" + +import os + +import torch +from peft import PeftModel + +from modelopt.torch.utils import get_unwrapped_name, silence_matched_warnings + +from ..conversion import ModeloptStateManager, restore_from_modelopt_state +from .huggingface import ( + _MODELOPT_STATE_SAVE_NAMES, + _get_modelopt_state_path, + _new_save_pretrained, + patch_pretrained_methods, +) + +__all__ = [] + +_MODELOPT_STATE_SAVE_NAMES[PeftModel] = "peft_modelopt_state.pth" + + +def _get_quantizer_state_save_path(dir): + return os.path.join(dir, "quantizer_state_dict.pth") + + +def _new_save_pretrained_peft(self, save_directory, *args, **kwargs): + _new_save_pretrained(self, save_directory, *args, **kwargs) + if not ModeloptStateManager.is_converted(self): + return + + # Lets save the quantizer state_dict separately + # PEFT save_pretrained only saves the state_dict corresponding to the adapters + # However our quantizers are part of the LoraLinear layers and not the adapters + # Also there might non-LoraLinear layers which have quantizers in the model which also wont be saved by PEFT + # So we need to save the quantizer state_dict separately + + # TODO: Move this to modelopt.torch.quantization.plugins.peft + from modelopt.torch.quantization.nn import TensorQuantizer + + # We should not call self/model.state_dict() here. HF Trainer calls model.save_pretrained() only from process 0 + # With FSDP, model.state_dict() will hang if it is not called from all processes + quantizer_state_dict = {} + for name, module in self.named_modules(): + if isinstance(module, TensorQuantizer): + quantizer_state_dict[get_unwrapped_name(name)] = module.state_dict() + if len(quantizer_state_dict) > 0: + torch.save(quantizer_state_dict, _get_quantizer_state_save_path(save_directory)) + + +def _new_load_adapter(self, model_id, adapter_name, *args, **kwargs): + modelopt_state_path = _get_modelopt_state_path(self, model_id) + + if os.path.isfile(modelopt_state_path): + assert ( + adapter_name in self.peft_config + ), f"ModelOpt modified model should have adapter_name={adapter_name} in peft_config" + restore_from_modelopt_state(self, torch.load(modelopt_state_path, map_location="cpu")) + + # The adapter state_dictionary does not contain the quantizer weights for the layers which are not LoraLinear + # However this is okay since the quantizer weights have been loaded from the quantizer_state_dict.pth in + # the previous step + with silence_matched_warnings(": No amax in state_dict."): + self._modelopt_cache["load_adapter"](self, model_id, adapter_name, *args, **kwargs) + + # TODO: Move this to modelopt.torch.quantization.plugins.peft + if os.path.isfile(_get_quantizer_state_save_path(model_id)): + from modelopt.torch.quantization.nn import TensorQuantizer + + quantizer_state_dict = torch.load( + _get_quantizer_state_save_path(model_id), map_location="cpu" + ) + for name, module in self.named_modules(): + if isinstance(module, TensorQuantizer): + module.load_state_dict(quantizer_state_dict[get_unwrapped_name(name)]) + + +patch_pretrained_methods( + PeftModel, + "peft", + patch_methods_map={ + "save_pretrained": _new_save_pretrained_peft, + "load_adapter": _new_load_adapter, + }, +) diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py new file mode 100644 index 0000000000..59293975f6 --- /dev/null +++ b/modelopt/torch/opt/plugins/transformers.py @@ -0,0 +1,24 @@ +# 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. + +"""ModelOpt plugin for enabling automatic save/restore of ModelOpt state for HuggingFace models.""" + +from transformers import PreTrainedModel + +from .huggingface import patch_pretrained_methods + +__all__ = [] + +patch_pretrained_methods(PreTrainedModel, "transformers") diff --git a/modelopt/torch/opt/searcher.py b/modelopt/torch/opt/searcher.py new file mode 100644 index 0000000000..010dc90dc6 --- /dev/null +++ b/modelopt/torch/opt/searcher.py @@ -0,0 +1,364 @@ +# 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. + +"""Standard interface to implement a searcher algorithm. + +A searcher is useful whenever we want to search/optimize over a set of hyperparameters in the model. +Searchers are usually used in conjunction with a mode, which can define a search space via its +entrypoints, i.e., convert the model into a search space. The searcher then optimizes over this +search space. +""" + +import copy +import os +from abc import ABC, abstractmethod +from contextlib import nullcontext +from typing import Any, Callable, Optional, Union, final + +import numpy as np +import pulp +import torch +import torch.nn as nn + +from modelopt.torch.utils import distributed as dist +from modelopt.torch.utils import no_stdout, run_forward_loop + +LimitsTuple = tuple[float, float] +ConstraintsDict = dict[str, Union[str, float, dict, None]] +Deployment = dict[str, str] +ForwardLoop = Callable[[nn.Module], None] +ScoreFunc = Callable[[nn.Module], float] + +SearchConfig = dict[str, Any] # config dict for searcher +SearchStateDict = dict[str, Any] # state dict for searcher + +__all__ = ["BaseSearcher"] + + +class BaseSearcher(ABC): + """A basic search interface that can be used to search/optimize a model. + + The base interface supports basic features like setting up a search, checkpointing, and + loading logic and defines a minimal workflow to follow. + """ + + model: nn.Module + config: SearchConfig + constraints: ConstraintsDict + deployment: Optional[Deployment] + dummy_input: Union[Any, tuple] + forward_loop: Optional[ForwardLoop] + + @final + def __init__(self) -> None: + """We don't allow to override __init__ method.""" + super().__init__() + + # TODO: see if we really want to keep all the config here. + @property + def default_search_config(self) -> SearchConfig: + """Get the default config for the searcher.""" + return { + "checkpoint": None, + "verbose": dist.is_master(), + "forward_loop": None, + "data_loader": None, + "collect_func": None, + "max_iter_data_loader": None, + "score_func": None, + "loss_func": None, + "deployment": None, + } + + @property + @abstractmethod + def default_state_dict(self) -> SearchStateDict: + """Return default state dict.""" + + def sanitize_search_config(self, config: Optional[SearchConfig]) -> SearchConfig: + """Sanitize the search config dict.""" + # supply with defaults (for verbose we wanna make sure it's on master only) + config = {**self.default_search_config, **(config or {})} + config["verbose"] = config["verbose"] and self.default_search_config["verbose"] + + # sanity checks + assert ( + config.keys() == self.default_search_config.keys() + ), f"Unexpected config keys: {config.keys() - self.default_search_config.keys()}" + + return config + + # TODO: double-check if we want all these args here. + @final + def search( + self, + model: nn.Module, + constraints: ConstraintsDict, + dummy_input: Optional[Union[Any, tuple]] = None, + config: Optional[SearchConfig] = None, + ) -> SearchStateDict: + """Search a given prunable model for the best sub-net and return the search model. + + The best sub-net maximizes the score given by ``score_func`` while satisfying the + ``constraints``. + + Args: + model: The converted model to be searched. + constraints: The dictionary from constraint name to upper bound the searched model has + to satisfy. + dummy_input: Arguments of ``model.forward()``. This is used for exporting and + calculating inference-based metrics, such as latency/FLOPs. The format of + ``dummy_inputs`` follows the convention of the ``args`` argument in + `torch.onnx.export `_. + config: Additional optional arguments to configure the search. + + Returns: A tuple (subnet, state_dict) where + subnet is the searched subnet (nn.Module), which can be used for subsequent tasks like + fine-tuning, state_dict contains the history and detailed stats of the search procedure. + """ + # check model train state + is_training = model.training + + # reset the search + self.reset_search() + + # update and initialize searcher + self.model = model + self.config = self.sanitize_search_config(config) + self.deployment = self.config["deployment"] + + self.constraints = constraints + self.dummy_input = dummy_input + self.forward_loop = self.construct_forward_loop(silent=not self.config["verbose"]) + + # load checkpoint if it exists + self.load_search_checkpoint() + + # run initial step and sanity checks before the search + self.before_search() + + # run actual search + self.run_search() + + # run clean-up steps after search + self.after_search() + + # make sure model is in original state + model.train(is_training) + + # return the config for the best result + return self.best + + def reset_search(self) -> None: + """Reset search at the beginning.""" + # reset self.best where we store results + self.best: SearchStateDict = {} + + # reset state dict (do it afterwards in case best is in state_dict) + for key, val in self.default_state_dict.items(): + setattr(self, key, copy.deepcopy(val)) + + def before_search(self) -> None: + """Optional pre-processing steps before the search.""" + pass + + @abstractmethod + def run_search(self) -> None: + """Run actual search.""" + + def after_search(self) -> None: + """Optional post-processing steps after the search.""" + pass + + @property + def has_score(self) -> bool: + """Check if the model has a score function.""" + return self.config["score_func"] is not None + + def eval_score(self, silent=True) -> float: + """Optionally silent evaluation of the score function.""" + assert self.has_score, "Please provide `score_func`!" + score_func: ScoreFunc = self.config["score_func"] + with no_stdout() if silent else nullcontext(): + return float(score_func(self.model)) + + def construct_forward_loop( + self, + silent=True, + progress_bar_msg=None, + max_iter_data_loader=None, + post_process_fn=False, + ) -> Optional[ForwardLoop]: + """Get runnable forward loop on the model using the provided configs.""" + # check config + data_loader = self.config["data_loader"] + forward_loop = self.config.get("forward_loop", None) + assert None in [data_loader, forward_loop], "Only provide `data_loader` or `forward_loop`!" + + # check trivial case + if not (data_loader or forward_loop): + return None + + def forward_loop_with_silence_check(m: nn.Module) -> None: + with no_stdout() if silent else nullcontext(): + if data_loader is not None: + run_forward_loop( + m, + data_loader=data_loader, + max_iters=( + max_iter_data_loader + if max_iter_data_loader is not None + else self.config.get("max_iter_data_loader", None) + ), + collect_func=self.config["collect_func"], + progress_bar=progress_bar_msg, + post_process=post_process_fn, + ) + elif forward_loop is not None: + forward_loop(m) + + return forward_loop_with_silence_check + + @final + def state_dict(self) -> SearchStateDict: + """The state dictionary that can be stored/loaded.""" + return {key: getattr(self, key) for key in self.default_state_dict} + + def load_search_checkpoint(self) -> bool: + """Load function for search checkpoint returning indicator whether checkpoint was loaded.""" + # check if checkpoint exists + checkpoint: Optional[str] = self.config["checkpoint"] + if checkpoint is None or not os.path.exists(checkpoint): + return False + + # iterate through state dict and load keys + print(f"Loading searcher state from {checkpoint}...") + state_dict = torch.load(checkpoint) + assert state_dict.keys() == self.state_dict().keys(), "Keys in checkpoint don't match!" + for key, state in state_dict.items(): + setattr(self, key, state) + return True + + def save_search_checkpoint(self) -> None: + """Save function for search checkpoint.""" + # check if save requirements are satisfied + checkpoint: Optional[str] = self.config["checkpoint"] + if checkpoint is None or not dist.is_master(): + return + + # save state dict + save_dirname, _ = os.path.split(checkpoint) + if save_dirname: + os.makedirs(save_dirname, exist_ok=True) + torch.save(self.state_dict(), checkpoint) + + +class LPS: + """A wrapper on top of PuLP Linear Programming Solver. + + This solver maximizes/minimizes the candidate scores while meeting the cost constraints. + """ + + def __init__( + self, + name: str, + constraints: dict[str, Union[float, tuple[float]]], # upper bound or (lower, upper) bound + constraints_to_candidate_costs: dict[str, list[list[float]]], + candidate_scores: list[list[float]], + objective_type: str = "minimize", + verbose: bool = False, + ): + """Initialize the LPS solver.""" + assert set(constraints.keys()) == set(constraints_to_candidate_costs.keys()) + assert objective_type in ("minimize", "maximize") + num_candidates_per_layer = list(map(len, candidate_scores)) + for candidate_costs in constraints_to_candidate_costs.values(): + assert list(map(len, candidate_costs)) == num_candidates_per_layer + + self.name = name + self.constraints = constraints + self.constraints_to_candidate_costs = constraints_to_candidate_costs + self.candidate_scores = candidate_scores + self.objective_type = pulp.LpMinimize if objective_type == "minimize" else pulp.LpMaximize + self.solver = pulp.PULP_CBC_CMD(msg=verbose) + + self.num_layers = len(self.candidate_scores) + self.num_candidates_per_layer = list(map(len, self.candidate_scores)) + + def _build_selection_vars(self) -> list[list[pulp.LpVariable]]: + vars = [] + for li in range(self.num_layers): + num_candidates = self.num_candidates_per_layer[li] + layer_vars = [ + pulp.LpVariable(f"z{li}_{ci}", lowBound=0, upBound=1, cat=pulp.LpBinary) + for ci in range(num_candidates) + ] + vars.append(layer_vars) + return vars + + def _build_objective_problem( + self, selection_vars: list[list[pulp.LpVariable]] + ) -> pulp.LpProblem: + problem = pulp.LpProblem(name=self.name, sense=self.objective_type) + objective_value = 0 + for layer_id, layer_vars in enumerate(selection_vars): + objective_value += sum( + [z * a for z, a in zip(layer_vars, self.candidate_scores[layer_id])] + ) + problem += (objective_value, "L") + return problem + + def _build_one_hot_constraints(self, selection_vars: list[list[pulp.LpVariable]]) -> list: + return [sum(layer_vars) == 1 for layer_vars in selection_vars] + + def _build_budget_constraints(self, selection_vars: list[list[pulp.LpVariable]]) -> list: + budget_constraints = [] + for ( + constraint_name, + candidate_costs_list, + ) in self.constraints_to_candidate_costs.items(): + cost = 0 + for layer_vars, candidate_costs in zip(selection_vars, candidate_costs_list): + cost += sum([z * b for z, b in zip(layer_vars, candidate_costs)]) + if isinstance(self.constraints[constraint_name], tuple): + lower_bound, upper_bound = self.constraints[constraint_name] # type: ignore[misc] + else: + lower_bound, upper_bound = None, self.constraints[constraint_name] + if upper_bound is not None: + budget_constraints.append(cost <= upper_bound) + if lower_bound is not None: + budget_constraints.append(cost >= lower_bound) + + return budget_constraints + + def __call__(self) -> tuple[list[int], str]: + """Run the solver. + + Returns: + selections: A list of selected candidate indices per layer. + status: Status of the solver. + """ + selection_vars = self._build_selection_vars() + + problem = self._build_objective_problem(selection_vars) + for one_hot_constraint in self._build_one_hot_constraints(selection_vars): + problem += one_hot_constraint + for budget_constraint in self._build_budget_constraints(selection_vars): + problem += budget_constraint + + problem.solve(self.solver) + + selections = [np.argmax([z.varValue for z in layer_vars]) for layer_vars in selection_vars] + return selections, pulp.LpStatus[problem.status] diff --git a/modelopt/torch/opt/utils.py b/modelopt/torch/opt/utils.py new file mode 100644 index 0000000000..9323692bee --- /dev/null +++ b/modelopt/torch/opt/utils.py @@ -0,0 +1,77 @@ +# 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. + +"""Utilities for optimization.""" + +from typing import Generator, Optional + +import torch.nn as nn + +from modelopt.torch.utils import unwrap_model + +from .dynamic import DynamicModule, DynamicSpace +from .hparam import Hparam + +__all__ = [ + "get_hparam", + "is_configurable", + "is_dynamic", + "named_dynamic_modules", + "named_hparams", + "search_space_size", +] + + +class _DynamicSpaceUnwrapped(DynamicSpace): + """A wrapper for the DynamicSpace class to handle unwrapping of model wrappers like DDP. + + This is useful to ensure that configurations are valid for both vanilla models and wrapped + models, see :meth:`unwrap_models` for supported wrappers. + """ + + def __init__(self, model: nn.Module) -> None: + super().__init__(unwrap_model(model)) + + +def is_configurable(model: nn.Module) -> bool: + """Check if the model is configurable.""" + return _DynamicSpaceUnwrapped(model).is_configurable() + + +def is_dynamic(model: nn.Module) -> bool: + """Check if the model is dynamic.""" + return _DynamicSpaceUnwrapped(model).is_dynamic() + + +def named_hparams( + model: nn.Module, configurable: Optional[bool] = None +) -> Generator[tuple[str, Hparam], None, None]: + """Recursively yield the name and instance of *all* hparams.""" + yield from _DynamicSpaceUnwrapped(model).named_hparams(configurable) + + +def named_dynamic_modules(model: nn.Module) -> Generator[tuple[str, DynamicModule], None, None]: + """Recursively yield the name and instance of *all* dynamic modules.""" + yield from _DynamicSpaceUnwrapped(model).named_dynamic_modules() + + +def get_hparam(model: nn.Module, name: str) -> Hparam: + """Get the hparam with the given name.""" + return _DynamicSpaceUnwrapped(model).get_hparam(name) + + +def search_space_size(model: nn.Module) -> int: + """Return the size of the search space.""" + return _DynamicSpaceUnwrapped(model).size() diff --git a/modelopt/torch/prune/__init__.py b/modelopt/torch/prune/__init__.py new file mode 100644 index 0000000000..918ebcb098 --- /dev/null +++ b/modelopt/torch/prune/__init__.py @@ -0,0 +1,26 @@ +# 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. + +"""A simplified API for :meth:`modelopt.torch.nas` for pruning algorithms. + +This module provides a simplified API for pruning that is based on the NAS infrastructure but +simplifies the overall workflow to accommodate for the simpler nature of pruning algorithms. +""" + +# nas is a required - so let's check if it's available +import modelopt.torch.nas + +from . import config, mode, plugins +from .pruning import * diff --git a/modelopt/torch/prune/config.py b/modelopt/torch/prune/config.py new file mode 100644 index 0000000000..5f234f5c07 --- /dev/null +++ b/modelopt/torch/prune/config.py @@ -0,0 +1,86 @@ +# 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. + +"""Default configurations for prune modes.""" + +from pydantic import create_model + +from modelopt.torch.nas.registry import DMRegistry +from modelopt.torch.opt.config import ModeloptBaseConfig, get_kwargs_for_create_model_with_rules + + +def _conv_config(): + return { + "channels_ratio": tuple(0.05 * i for i in range(1, 21)), + "kernel_size": (), + "channel_divisor": 32, + } + + +def _norm_lin_config(): + return { + "features_ratio": tuple(0.05 * i for i in range(1, 21)), + "feature_divisor": 32, + } + + +def _get_fastnas_default_rules(): + return { + "nn.Conv1d": _conv_config(), + "nn.Conv2d": _conv_config(), + "nn.Conv3d": _conv_config(), + "nn.ConvTranspose1d": _conv_config(), + "nn.ConvTranspose2d": _conv_config(), + "nn.ConvTranspose3d": _conv_config(), + "nn.Linear": _norm_lin_config(), + "nn.BatchNorm1d": _norm_lin_config(), + "nn.BatchNorm2d": _norm_lin_config(), + "nn.BatchNorm3d": _norm_lin_config(), + "nn.SyncBatchNorm": _norm_lin_config(), + "nn.InstanceNorm1d": _norm_lin_config(), + "nn.InstanceNorm2d": _norm_lin_config(), + "nn.InstanceNorm3d": _norm_lin_config(), + "nn.LayerNorm": _norm_lin_config(), + "nn.GroupNorm": {k: v for k, v in _conv_config().items() if k != "kernel_size"}, + } + + +FastNASConfig: type[ModeloptBaseConfig] = create_model( + "FastNASConfig", + **get_kwargs_for_create_model_with_rules( + registry=DMRegistry, + default_rules=_get_fastnas_default_rules(), + doc='Configuration for the ``"fastnas"`` mode.', + ), +) + + +GradNASConfig: type[ModeloptBaseConfig] = create_model( + "GradNASConfig", + **get_kwargs_for_create_model_with_rules( + registry=DMRegistry, + default_rules=_get_fastnas_default_rules(), + doc='Configuration for the ``"gradnas"`` mode.', + ), +) + +MCoreGPTMinitronConfig: type[ModeloptBaseConfig] = create_model( + "MCoreGPTMinitronConfig", + **get_kwargs_for_create_model_with_rules( + registry=DMRegistry, + default_rules={}, # Dynamically generated rules if Megatron-core is available + doc='Configuration for the ``"mcore_gpt_minitron"`` mode.', + ), +) diff --git a/modelopt/torch/prune/fastnas.py b/modelopt/torch/prune/fastnas.py new file mode 100644 index 0000000000..963aafdf91 --- /dev/null +++ b/modelopt/torch/prune/fastnas.py @@ -0,0 +1,263 @@ +# 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. + +"""Module implementing ``fasnas`` pruning algorithm for search.""" + +import warnings +from typing import Any, Callable, Optional + +import torch.nn as nn +import tqdm + +from modelopt.torch.nas.autonas import ( + AutoNASPatchManager, + IterativeSearcher, + convert_searchspace, + restore_searchspace, + update_autonas_metadata, +) +from modelopt.torch.nas.hparams import ConcatTracedHp +from modelopt.torch.nas.utils import get_subnet_config, sample, select, sort_parameters +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict +from modelopt.torch.opt.searcher import SearchStateDict +from modelopt.torch.opt.utils import named_hparams +from modelopt.torch.utils import random + +ConstraintsRes = dict[str, float] +ConstraintEvalFunc = Callable[[Optional[ConstraintsRes]], float] + + +class FastNASPatchManager(AutoNASPatchManager): + """A patch manager for FastNAS (same as AutoNAS except no sampling during training).""" + + @property + def sample_during_training(self): + """Indicates whether we should sample a new subnet during training.""" + return False + + +update_fastnas_metadata = update_autonas_metadata + + +def convert_fastnas_searchspace(model: nn.Module, config: ModeloptBaseConfig) -> ConvertReturnType: + """Convert search space for FastNAS mode with correct patch manager.""" + return convert_searchspace(model, config, FastNASPatchManager) + + +def restore_fastnas_searchspace( + model: nn.Module, config: ModeloptBaseConfig, metadata: MetadataDict +) -> nn.Module: + """Restore search space for FastNAS mode with correct patch manager.""" + return restore_searchspace(model, config, metadata, FastNASPatchManager) + + +class BinarySearcher(IterativeSearcher): + """An iterative searcher that uses binary search to find the best configuration.""" + + sensitivity_map: dict[str, dict[int, float]] + min_degrade: float + max_degrade: float + middle_value: float + original_score: float + + @property + def hparam_names_for_search(self) -> set[str]: + """We can only optimize over certain types of hparams in binary search.""" + # TODO: eventually improve this! + return {"out_channels", "out_features", "num_attention_heads"} + + @property + def hparam_types_for_search(self) -> tuple[type]: + """We can only optimize over certain types of hparams in binary search.""" + # TODO: eventually improve this! + return (ConcatTracedHp,) + + def before_search(self) -> None: + """Build sensitivity map before search that we use to approximate the cost function.""" + # run super method + super().before_search() + + # sort parameters in model + sort_parameters(self.model) + + # for binary search, we currently only support 1 constraint + assert len(self.constraints_func.effective_limits) == 1, ( + f"Currently, we only support one constraint for `{type(self).__name__}`, but got" + f" {len(self.constraints_func.effective_limits)}" + ) + + # compute and register the construction of sensitivity map + self._build_sensitivity_map(self.config["verbose"]) + self.max_degrade = max([max(v.values()) for v in self.sensitivity_map.values()]) + self.min_degrade = min([min(v.values()) for v in self.sensitivity_map.values()]) + + # overwrite the score function to be a fake function, returning the -max degrade + def max_degrade(_model): + return -max( + self.sensitivity_map[k][active] + for k, active in get_subnet_config(_model, True).items() + if k in self.sensitivity_map + ) + + self.config["score_func"] = max_degrade + + def before_step(self) -> None: + """Check what the middle value is to determine where we recurse.""" + self.middle_value = 0.5 * (self.max_degrade + self.min_degrade) + + def _apply_fastnas_according_to_threshold(self, threshold): + cfg = { + name: min([k for k, v in sensitivity.items() if v <= threshold]) + for name, sensitivity in self.sensitivity_map.items() + } + select(self.model, cfg, strict=False) + + def sample(self) -> dict[str, Any]: + """Check in which interval we should recurse and sets the corresponding subnet.""" + self._apply_fastnas_according_to_threshold(self.middle_value) + return {"config": get_subnet_config(self.model)} + + def after_step(self) -> None: + """Update boundaries of the interval after recursing.""" + if self.candidate["is_satisfied"]: + self.min_degrade, self.max_degrade = self.min_degrade, self.middle_value + else: + self.min_degrade, self.max_degrade = self.middle_value, self.max_degrade + + def early_stop(self) -> bool: + """Early stop if the interval is small enough.""" + return abs(0.5 * (self.min_degrade + self.max_degrade) - self.middle_value) < 1e-5 + + @property + def default_state_dict(self) -> SearchStateDict: + """We also store the sensitivity map and related arguments.""" + return { + **super().default_state_dict, + "sensitivity_map": {}, + "middle_value": None, + "min_degrade": None, + "max_degrade": None, + "original_score": None, + } + + def load_search_checkpoint(self) -> bool: + """We only want to load sensitivity map and original_score here and keep the rest.""" + state_dict = self.state_dict() + keys_to_load = ("sensitivity_map", "original_score") + is_loaded = super().load_search_checkpoint() + if is_loaded: + for k, v in state_dict.items(): + if k not in keys_to_load: + setattr(self, k, v) + return is_loaded + + def _get_binary_search_hps(self): + hps_configurable = {k: hp for k, hp in named_hparams(self.model, configurable=True)} + + # TODO: eventually improve this! + binary_search_hps = { + k: hp + for k, hp in hps_configurable.items() + if k.split(".")[-1] in self.hparam_names_for_search + or isinstance(hp, self.hparam_types_for_search) + } + + # Sanity check if there are other searchable hparams not supported by the searcher. + if len(binary_search_hps) < len(hps_configurable): + warnings.warn( + f"`{type(self).__name__}` does not support search for" + f" {hps_configurable.keys() - binary_search_hps.keys()}" + ) + + return binary_search_hps + + def _build_sensitivity_map(self, verbose=False) -> None: + if self.original_score is None: + sample(self.model, random.original) + self.original_score = self.eval_score() + + binary_search_hps = self._get_binary_search_hps() + + # Measure sensitivity + + # Getting the number of choices needed to validate + total_choices_to_validate = sum( + [len(hparam.choices) for hparam in binary_search_hps.values()] + ) + + assert total_choices_to_validate != 0, f"{type(self).__name__}: no searchable hparams found" + + binary_search_hps = { + k: hp + for k, hp in binary_search_hps.items() + if ( + k not in self.sensitivity_map + or sorted(hp.choices) != sorted(self.sensitivity_map[k].keys()) + ) + } + + remaining_choices_to_validate = sum( + [len(hparam.choices) for hparam in binary_search_hps.values()] + ) + + if remaining_choices_to_validate == 0: + return + + print( + "\nBeginning pre-search estimation. If the runtime of score function is longer than", + "a few minutes, consider subsampling the dataset used in score function.", + "\nA PyTorch dataset can be subsampled using torch.utils.data.Subset", + "(https://pytorch.org/docs/stable/data.html#torch.utils.data.Subset) as following:\n", + "subset_dataset = torch.utils.data.Subset(dataset, indices)", + ) + with tqdm.tqdm( + initial=total_choices_to_validate - remaining_choices_to_validate, + total=total_choices_to_validate, + desc="Collecting pre-search statistics", + smoothing=0.1, + ) as pbar: + sample(self.model, sample_func=max) # reset for sanity + for name, hparam in binary_search_hps.items(): + reach_zero = False # if already reach 0 or <0, we can skip later computations + cur_sensitivity_scores = {} + for choice in sorted(hparam.choices): + select(self.model, {name: choice}, strict=False) + if hparam.active == hparam.max or reach_zero: + cur_sensitivity_scores[hparam.active] = 0.0 + else: + score = self.eval_score() + cur_sensitivity_scores[hparam.active] = self.original_score - score + if self.original_score - score < 1e-5: + reach_zero = True # will skip the computation for later choices + # better visualization for debug + if verbose: + pbar.set_postfix( + { + "cur": ( + f"{name}({hparam.active}/{hparam.max}):" + f" {cur_sensitivity_scores[hparam.active]:.2f}" + ) + } + ) + pbar.update(1) + + # make sure the sensitivity is monotonic + prev_max = 0.0 + for k in list(cur_sensitivity_scores.keys())[::-1]: + cur_sensitivity_scores[k] = prev_max = max(cur_sensitivity_scores[k], prev_max) + + self.sensitivity_map[name] = cur_sensitivity_scores + self.save_search_checkpoint() diff --git a/modelopt/torch/prune/gradnas.py b/modelopt/torch/prune/gradnas.py new file mode 100644 index 0000000000..4c3c9048d4 --- /dev/null +++ b/modelopt/torch/prune/gradnas.py @@ -0,0 +1,314 @@ +# 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. + +"""Module implementing ``gradnas`` pruning algorithm for search. + +Summary: +-------- +`gradnas` algorithm gives a better score to sort various pruning choices over L1 norm (`fastnas`) for language models. + +Details: +-------- +Further, we can get scores for hparams which are implemented even abstractly. +For example, we can use this algorithm to sort the heads in a multi-head attention layer. The attention heads +do not have a unique tensor parameter associated to it. + +We are ranking the prunable choices for a particular hparam based on Sum((gradient of loss wrt pruning mask)^2). +The pruning mask of an hparam is a binary mask indicating which choices of the hparam are pruned +(0 means pruned and 1 means not pruned). + +While calculating the backward gradient of loss, the masks are set to 1 at all tensors. +See more about masks being used to measure sensitivity in this paper: https://arxiv.org/pdf/1905.10650.pdf +""" + +import types +import warnings +from contextlib import contextmanager +from functools import partial +from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional + +import torch +import torch.nn as nn +from torch.utils.hooks import RemovableHandle +from tqdm import tqdm + +import modelopt.torch.nas.modules as dnn +from modelopt.torch.nas.registry import DMRegistry +from modelopt.torch.nas.traced_hp import TracedHp +from modelopt.torch.nas.utils import get_subnet_config, sample, select +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.opt.searcher import SearchConfig +from modelopt.torch.opt.utils import named_hparams +from modelopt.torch.utils import get_module_device, standardize_model_args, torch_to, unwrap_model + +from .fastnas import BinarySearcher + +if TYPE_CHECKING: + from modelopt.torch.nas.plugins.transformers import ( + _DynamicAttention, + _DynamicBertAttention, + _DynamicGPTJAttention, + ) + +try: + from transformers.models.bert.modeling_bert import BertAttention + from transformers.models.gptj.modeling_gptj import GPTJAttention + + HAS_HF = True +except ImportError: + HAS_HF = False + + +class GradientDataManager: + """Class for managing gradient data for an hparam.""" + + def __init__(self, shape, model, reduce_func=lambda x: x): + """Initialize GradientDataManager.""" + self.mask = torch.ones(shape, requires_grad=True, device=get_module_device(model)) + self._score = torch.zeros_like(self.mask, requires_grad=False) + self._reduce_func = reduce_func + + # TODO: Implement score averaging and early stopping based on score convergence + def process_gradient(self): + """Process gradient of the mask.""" + self._score += self.mask.grad.detach().pow(2) + self.mask.grad = None + + @property + def score(self): + """The score of the hparam based on the stored gradients.""" + return self._reduce_func(self._score) + + +def _setup_grad_manager_linear( + module: dnn._DynamicLinear, +) -> tuple[GradientDataManager, RemovableHandle]: + def forward_hook(_modelopt_mask, module, input, output): + return output * _modelopt_mask + + grad_data = GradientDataManager(module.get_hparam("out_features").max, module) + hook_handle = module.register_forward_hook(partial(forward_hook, grad_data.mask)) + return grad_data, hook_handle + + +def _setup_grad_manager_hf_attention( + module: "_DynamicAttention", head_mask_idx: int +) -> tuple[GradientDataManager, RemovableHandle]: + def forward_pre_hook(_modelopt_mask, module, args, kwargs): + head_mark_in_args = False + head_mask_in_kwargs = False + if "head_mask" in kwargs: + head_mask = kwargs["head_mask"] + head_mask_in_kwargs = True + elif len(args) > head_mask_idx: + head_mask = args[head_mask_idx] + head_mark_in_args = True + else: + head_mask = None + + if head_mask is None: + head_mask = _modelopt_mask + else: + # head_mask shape: 1 x num_attention_heads x 1 x 1 + # https://github.com/huggingface/transformers/blob/8f796960f6bb6e68c673a2f682d5ac911f671ed8/ + # src/transformers/modeling_utils.py#L841 + head_mask = head_mask * _modelopt_mask + + if head_mask_in_kwargs or not head_mark_in_args: + kwargs["head_mask"] = head_mask + elif head_mark_in_args: + args = list(args) + args[head_mask_idx] = head_mask + args = tuple(args) + + return args, kwargs + + grad_data = GradientDataManager( + (module.get_hparam("num_attention_heads").max, 1, 1), + module, + reduce_func=lambda x: x.squeeze(), + ) + hook_handle = module.register_forward_pre_hook( + partial(forward_pre_hook, grad_data.mask), with_kwargs=True + ) + return grad_data, hook_handle + + +def _setup_grad_manager_bert_attention( + module: "_DynamicBertAttention", +) -> tuple[GradientDataManager, RemovableHandle]: + # See forward signature here: https://github.com/huggingface/transformers/blob/ + # b8648290d2d97e7c7dbccd2d4a6a4f44e70d3b63/src/transformers/models/bert/modeling_bert.py#L415-L424 + + return _setup_grad_manager_hf_attention(module, head_mask_idx=2) + + +def _setup_grad_manager_gptj_attention( + module: "_DynamicGPTJAttention", +) -> tuple[GradientDataManager, RemovableHandle]: + # See forward signature here: https://github.com/huggingface/transformers/blob/ + # 0ea42ef0f9f71deba7775ead33afa0e493823d60/src/transformers/models/gptj/modeling_gptj.py#L194-L202 + + return _setup_grad_manager_hf_attention(module, head_mask_idx=4) + + +class GradientBinarySearcher(BinarySearcher): + """Binary searcher for gradient algorithm.""" + + SETUP_GRADIENT_FUNC: dict[ + type[DynamicModule], Callable[[DynamicModule], tuple[GradientDataManager, RemovableHandle]] + ] + + @property + def default_search_config(self) -> SearchConfig: + """Get the default config for the searcher.""" + config = super().default_search_config + config["max_iter_data_loader"] = 128 # Default 50 is not optimal for gradient estimation + return config + + def before_search(self) -> None: + """Setup search with gradient-based score.""" + # initialize here to keep DMRegistry empty during imports + GradientBinarySearcher.SETUP_GRADIENT_FUNC = { + DMRegistry[nn.Linear]: _setup_grad_manager_linear, + } + if HAS_HF: + GradientBinarySearcher.SETUP_GRADIENT_FUNC.update( + { + DMRegistry[BertAttention]: _setup_grad_manager_bert_attention, + DMRegistry[GPTJAttention]: _setup_grad_manager_gptj_attention, + } + ) + # TODO: Support distributed models for gradient search + self.model = unwrap_model(self.model, raise_error=True) + + hps_for_grad_calc = self._estimate_gradient_scores( + self.config["data_loader"], + self.config["loss_func"], + self.config["collect_func"], + self.config["max_iter_data_loader"], + ) + + # This is a hack so that `sort_parameters` will use `score_tensor` + # to sort channels/features of the respective hparams in `super().before_search` + with self._overwrite_hp_importance(hps_for_grad_calc): + super().before_search() + + # features/heads corresponding to the hparams have already been sorted, now sort the + # score_tensors + for _, hp in hps_for_grad_calc.items(): + assert hasattr(hp, "score_tensor") + hp.score_tensor = hp.score_tensor.sort(descending=True).values + + def sanitize_search_config(self, config: Optional[SearchConfig]) -> SearchConfig: + """Sanitize the search config dict.""" + config = config or {} + if config.get("score_func") is not None: + warnings.warn("`score_func` should not be provided for `gradnas`.") + config["score_func"] = GradientBinarySearcher.gradnas_score_func + config = super().sanitize_search_config(config) + assert config["data_loader"] is not None, "data_loader must be provided for `gradnas`." + return config + + @property + def hparam_names_for_search(self) -> set[str]: + """We can only optimize over certain types of hparams in gradient binary search.""" + return {"out_features", "num_attention_heads"} + + @staticmethod + def gradnas_score_func(model: nn.Module) -> float: + """Score function for `gradnas` algorithm. + + If we prune N neurons from layer L, the total degradation is the sum of degradation values of the + N pruned neurons. In `fast` algorithm, the degradation due to pruning is estimated directly from + `validation_score(model after pruning)`. Rest of the algorithm is exactly the same as `fast` algorithm. + """ + score = 0 + for _, hp in named_hparams(model, configurable=True): + if hasattr(hp, "score_tensor"): + # Get the score for the pruned away parameters + hp_score = torch.sum(hp.score_tensor) - torch.sum(hp.score_tensor[hp.active_slice]) + score = max(score, hp_score) + return -score + + def _estimate_gradient_scores( + self, + data_loader: Iterable, + loss_func: Callable[[Any, Any], torch.Tensor], + collect_func: Optional[Callable] = None, + max_iter_data_loader: int = 128, + ) -> dict[str, TracedHp]: + """Estimate gradient scores for the searchable hparams in the model.""" + if collect_func is None: + + def collect_func(data): + return data[0] + + config = get_subnet_config(self.model) + sample(self.model, max) + + hp_grad_data: dict[str, GradientDataManager] = {} + mod_to_hook = {} + hps_for_grad_calc = self._get_binary_search_hps() + for hp_name in hps_for_grad_calc: + module_name = hp_name.rpartition(".")[0] + module = self.model.get_submodule(module_name) + hp_grad_data[hp_name], mod_to_hook[hp_name] = ( + GradientBinarySearcher.SETUP_GRADIENT_FUNC[type(module)](module) + ) + + device = get_module_device(self.model) + for idx, batch in tqdm( + enumerate(data_loader), desc="Estimating gradient scores", total=max_iter_data_loader + ): + args = standardize_model_args(self.model, collect_func(batch), use_kwargs=True) + args = torch_to(args, device) + output = self.model(*args[:-1], **args[-1]) + + loss = loss_func(output, batch) + loss.backward() + + for grad_data in hp_grad_data.values(): + grad_data.process_gradient() + + if idx >= max_iter_data_loader: + break + + # TODO: Normalize the score_tensors so max(sum(hp.score_tensor)) = 1 + for hp_name, grad_data in hp_grad_data.items(): + hps_for_grad_calc[hp_name].score_tensor = grad_data.score + + # remove all forward hooks + for hook in mod_to_hook.values(): + hook.remove() + + select(self.model, config) + + return hps_for_grad_calc + + @contextmanager + def _overwrite_hp_importance(self, hps_for_grad_calc: dict[str, TracedHp]): + """Context manager to overwrite `_get_importance` for hparams in the model.""" + for _, hp in hps_for_grad_calc.items(): + assert getattr(hp, "score_tensor") is not None + setattr(hp, "_get_importance_bkp", hp._get_importance) + hp._get_importance = types.MethodType(lambda x: x.score_tensor, hp) + + yield + + for _, hp in hps_for_grad_calc.items(): + assert hasattr(hp, "_get_importance_bkp") + hp._get_importance = hp._get_importance_bkp + delattr(hp, "_get_importance_bkp") diff --git a/modelopt/torch/prune/mcore_gpt_minitron.py b/modelopt/torch/prune/mcore_gpt_minitron.py new file mode 100644 index 0000000000..540328751d --- /dev/null +++ b/modelopt/torch/prune/mcore_gpt_minitron.py @@ -0,0 +1,155 @@ +# 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. + +"""Module implementing ``mcore_gpt_minitron`` pruning algorithm for NVIDIA Megatron-Core / NeMo models. + +Minitron pruning algorithm uses activation magnitudes to estimate importance of neurons / attention heads in the model. +More details on Minitron pruning algorithm can be found here: https://arxiv.org/pdf/2407.14679 +""" + +from typing import Optional + +import torch + +from modelopt.torch.nas.utils import sort_parameters +from modelopt.torch.opt.searcher import BaseSearcher, SearchConfig, SearchStateDict +from modelopt.torch.opt.utils import named_hparams + +SUPPORTED_MODEL_CONFIG_MAP = {} # map from model type to config attribute name +SUPPORTED_HPARAMS = { + # Width pruning + "ffn_hidden_size", + "num_attention_heads", + "num_query_groups", + "hidden_size", + # Depth pruning + "num_layers", +} + +try: + from megatron.core.models.gpt import GPTModel + + SUPPORTED_MODEL_CONFIG_MAP[GPTModel] = "config" +except Exception: + pass + +try: + from nemo.collections import llm + from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel + + SUPPORTED_MODEL_CONFIG_MAP[MegatronGPTModel] = "cfg" + SUPPORTED_MODEL_CONFIG_MAP[llm.GPTModel] = "config" +except Exception: + pass + + +class MCoreGPTMinitronSearcher(BaseSearcher): + """Searcher for Minitron pruning algorithm.""" + + @property + def default_search_config(self) -> SearchConfig: + """Get the default config for the searcher.""" + return {**super().default_search_config, "max_iter_data_loader": 1024} + + @property + def default_state_dict(self) -> SearchStateDict: + """Return default state dict.""" + return {} # Not used + + def sanitize_search_config(self, config: Optional[SearchConfig]) -> SearchConfig: + """Sanitize the search config dict.""" + config = super().sanitize_search_config(config) + assert ( + config["data_loader"] or config["forward_loop"] + ), "Data loader or forward loop must be provided for importance estimation!" + return config + + def before_search(self) -> None: + """Optional pre-processing steps before the search.""" + super().before_search() + + # Check that the constraint is valid + assert self.constraints.keys() == { + "export_config" + }, "Only `export_config` constraint is supported for pruning!" + + export_config = self.constraints["export_config"] + assert isinstance(export_config, dict) # to keep mypy happy + assert ( + export_config.keys() <= SUPPORTED_HPARAMS + ), f"Only {SUPPORTED_HPARAMS} are supported for pruning!" + + assert ("num_attention_heads" in export_config and "num_query_groups" in export_config) or ( + "num_attention_heads" not in export_config and "num_query_groups" not in export_config + ), "Both `num_attention_heads` and `num_query_groups` should be provided together!" + + # Only sort the parameters that are to be pruned + # If a user only prunes depth, we should not sort width parameters + self.hps_to_sort = SUPPORTED_HPARAMS & export_config.keys() + + # Convert `num_attention_heads` to `num_heads_per_group` + # Still keep `num_attention_heads` for updating model_cfg below + if "num_attention_heads" in export_config and "num_query_groups" in export_config: + export_config["num_heads_per_group"] = ( + export_config["num_attention_heads"] // export_config["num_query_groups"] + ) + self.hps_to_sort.add("num_heads_per_group") + + for n, hp in named_hparams(self.model, configurable=True): + hp_name = n.split(".")[-1] + if hp_name in export_config: + assert ( + export_config[hp_name] in hp.choices + ), f"Invalid choice for {hp_name}! Available choices: {hp.choices}" + + def run_search(self) -> None: + """Run actual search.""" + # Run forward loop to collect activations and sort parameters + assert self.forward_loop is not None + is_training = self.model.training + self.model.eval() + with torch.no_grad(): + self.forward_loop(self.model) + sort_parameters(self.model, self.hps_to_sort) + self.model.train(is_training) + + # Prune homogeneously + export_config = self.constraints["export_config"] + assert isinstance(export_config, dict) # to keep mypy happy + for n, hp in named_hparams(self.model, configurable=True): + hp_name = n.split(".")[-1] + if hp_name in export_config: + hp.active = export_config[hp_name] + + model_cfg = None + for m_type, cfg_name in SUPPORTED_MODEL_CONFIG_MAP.items(): + if isinstance(self.model, m_type): + model_cfg = getattr(self.model, cfg_name) + break + if model_cfg is None: + raise NotImplementedError( + f"Only {SUPPORTED_MODEL_CONFIG_MAP.keys()} models are supported! Got: {type(self.model)}" + ) + + # kv_channels can be None so we need to save original from original hidden_size and num_attention_heads + orig_kv_channels = getattr(model_cfg, "kv_channels") + if orig_kv_channels is None: + orig_kv_channels = getattr(model_cfg, "hidden_size") // getattr( + model_cfg, "num_attention_heads" + ) + setattr(model_cfg, "kv_channels", orig_kv_channels) + for n in SUPPORTED_HPARAMS: + if n in export_config: + setattr(model_cfg, n, export_config[n]) diff --git a/modelopt/torch/prune/mode.py b/modelopt/torch/prune/mode.py new file mode 100644 index 0000000000..e004adfbac --- /dev/null +++ b/modelopt/torch/prune/mode.py @@ -0,0 +1,146 @@ +# 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. + +"""Module implementing and describing modes that can be used during the NAS convert process. + +Check out :meth:`mtp.prune ` to learn more about modes. +""" + +from typing import Optional + +from modelopt.torch.nas.mode import NASModeRegistry +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.mode import ( + ConvertEntrypoint, + RestoreEntrypoint, + UpdateEntrypoint, + _ModeDescriptor, + _ModeRegistryCls, +) +from modelopt.torch.opt.searcher import BaseSearcher + +from .config import FastNASConfig, GradNASConfig, MCoreGPTMinitronConfig +from .fastnas import ( + BinarySearcher, + convert_fastnas_searchspace, + restore_fastnas_searchspace, + update_fastnas_metadata, +) +from .gradnas import GradientBinarySearcher +from .mcore_gpt_minitron import MCoreGPTMinitronSearcher + +PruneModeRegistry = _ModeRegistryCls() + + +@NASModeRegistry.register_mode +@PruneModeRegistry.register_mode +class FastNASModeDescriptor(_ModeDescriptor): + """Class to describe the ``"fastnas"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "fastnas" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return FastNASConfig + + @property + def next_modes(self) -> Optional[set[str]]: + """Modes that must immediately follow this mode.""" + return {"export", "kd_loss", "quantize", "sparse_magnitude", "sparse_gpt"} + + @property + def export_mode(self) -> Optional[str]: + """The mode that corresponds to the export mode of this mode.""" + return "export" + + @property + def search_algorithm(self) -> type[BaseSearcher]: + """Specifies the search algorithm to use for this mode (if any).""" + return BinarySearcher + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return convert_fastnas_searchspace + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_fastnas_searchspace + + @property + def update_for_save(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models state before saving.""" + return update_fastnas_metadata + + @property + def update_for_new_mode(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models state before new mode.""" + return update_fastnas_metadata + + +@NASModeRegistry.register_mode +@PruneModeRegistry.register_mode +class GradNASModeDescriptor(FastNASModeDescriptor): + """Class to describe the ``"gradnas"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "gradnas" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return GradNASConfig + + @property + def search_algorithm(self) -> type[BaseSearcher]: + """Specifies the search algorithm to use for this mode (if any).""" + return GradientBinarySearcher + + +@NASModeRegistry.register_mode +@PruneModeRegistry.register_mode +class MCoreGPTMinitronModeDescriptor(FastNASModeDescriptor): + """Class to describe the ``"mcore_gpt_minitron"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "mcore_gpt_minitron" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return MCoreGPTMinitronConfig + + @property + def search_algorithm(self) -> type[BaseSearcher]: + """Specifies the search algorithm to use for this mode (if any).""" + return MCoreGPTMinitronSearcher diff --git a/modelopt/torch/prune/plugins/__init__.py b/modelopt/torch/prune/plugins/__init__.py new file mode 100644 index 0000000000..927c20944d --- /dev/null +++ b/modelopt/torch/prune/plugins/__init__.py @@ -0,0 +1,24 @@ +# 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. + +"""Handles prune plugins for third-party modules.""" + +from modelopt.torch.utils import import_plugin + +with import_plugin("megatron"): + from .megatron import * + +with import_plugin("transformers"): + from .transformers import * diff --git a/modelopt/torch/prune/plugins/megatron.py b/modelopt/torch/prune/plugins/megatron.py new file mode 100644 index 0000000000..213db81492 --- /dev/null +++ b/modelopt/torch/prune/plugins/megatron.py @@ -0,0 +1,32 @@ +# 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. + +"""Plugins for pruning for Megatron-Core / NeMo modules using Minitron algorithm.""" + +# import nas plugin to check if it is enabled else raises an Exception +from modelopt.core.torch.nas.plugins.megatron import * # noqa: F403 + +from ..config import MCoreGPTMinitronConfig + +MCoreGPTMinitronConfig.register_default( + { + "megatron.core.models.gpt.GPTModel": { + "hidden_size_divisor": 64, + "num_heads_per_group_divisor": 1, + "num_query_groups_divisor": 1, + "ffn_hidden_size_divisor": 64, + }, + } +) diff --git a/modelopt/torch/prune/plugins/transformers.py b/modelopt/torch/prune/plugins/transformers.py new file mode 100644 index 0000000000..be01d8c5eb --- /dev/null +++ b/modelopt/torch/prune/plugins/transformers.py @@ -0,0 +1,40 @@ +# 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. + +"""Plugins for pruning for Transformers Attention.""" + +# import nas plugin to check if it is enabled else raises an Exception +from modelopt.torch.nas.plugins.transformers import * # noqa: F403 + +from ..config import FastNASConfig, GradNASConfig + + +def _n_heads_config(): + return {"n_heads_ratio": None, "n_heads_divisor": 1} + + +FastNASConfig.register_default( + { + "hf.BertAttention": _n_heads_config(), + "hf.GPTJAttention": _n_heads_config(), + } +) + +GradNASConfig.register_default( + { + "hf.BertAttention": _n_heads_config(), + "hf.GPTJAttention": _n_heads_config(), + } +) diff --git a/modelopt/torch/prune/pruning.py b/modelopt/torch/prune/pruning.py new file mode 100644 index 0000000000..58b1c3df02 --- /dev/null +++ b/modelopt/torch/prune/pruning.py @@ -0,0 +1,211 @@ +# 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. + +"""High-level API to automatically prune and optimize your model with various algorithms.""" + +from typing import Any, Optional, Union + +from torch import nn + +import modelopt.torch.nas as mtn +from modelopt.torch.opt.conversion import apply_mode +from modelopt.torch.opt.mode import ModeLike +from modelopt.torch.opt.searcher import ConstraintsDict, SearchConfig + +from .mode import PruneModeRegistry + +__all__ = ["prune"] + + +def prune( + model: nn.Module, + mode: ModeLike, + constraints: ConstraintsDict, + dummy_input: Union[Any, tuple], + config: Optional[SearchConfig] = None, +) -> tuple[nn.Module, dict[str, Any]]: + """Prune a given model by searching for the best architecture within the design space. + + Args: + model: A standard model that contains standard building blocks to be pruned in-place. + mode: A (list of) string(s) or Mode(s) or a list of tuples containing the mode and its + config indicating the desired mode(s) (and configurations) for the convert + process. Modes set up the model for different algorithms for model optimization. The + following modes are available: + + * :class:`"fastnas"`: The ``model`` will + be converted into a search space and set up to automatically perform operations + required for FastNAS pruning & search. The mode's config + is described in :class:`FastNASConfig`. + This mode is recommended to prune Computer Vision models. + * :class:`"gradnas"`: The ``model`` will + be converted into a search space and set up to automatically perform operations + required for gradient-based pruning & search. The mode's config + is described in :class:`GradNASConfig`. + This mode is recommended to prune Hugging Face language models like BERT and GPT-J. + * :class:`"mcore_gpt_minitron"`: The ``model`` + will be converted into a search space and set up to automatically perform operations + required for Minitron-style pruning & search. The mode's config + is described in :class:`MCoreGPTMinitronConfig`. + This mode is required to prune NVIDIA Megatron-Core / NeMo GPT-type models. + + If the mode argument is specified as a dictionary, the keys should indicate the mode and + the values specify the per-mode configuration. If not provided, then default + configuration would be used. + constraints: A dictionary mapping constraint names to their respective values that the pruned model must + satisfy. Currently, the supported constraints are ``flops``, ``params``, and ``export_config``. If the key + is ``flops`` or ``params``, the value should be an upper bound number or percentage of original. For + ``export_config``, the value is a dictionary mapping hyperparameter names to their pruned values. For e.g.,: + + .. code-block:: python + + # Specify a flops upper bound as 4.5 GFLOPs + constraints = {"flops": 4.5e6} + + # Specify a percentage-based constraint + # (e.g., search for a model with <= 60% of the original model params) + constraints = {"params": "60%"} + + # Specify export_config with pruned hyperparameters + # This is supported and required if the model is converted via ``mcore_gpt_minitron`` mode. + constraints = { + "export_config": { + "ffn_hidden_size": 128, + "num_attention_heads": 16, + "num_query_groups": 4, + } + } + + dummy_input: Arguments of ``model.forward()``. This is used for exporting and calculating + inference-based metrics, such as FLOPs. The format of ``dummy_inputs`` follows the + convention of the ``args`` argument in + `torch.onnx.export `_. + Specifically, ``dummy_input`` can be: + + #. a single argument (``type(dummy_input) != tuple``) corresponding to + + .. code-block:: python + + model.forward(dummy_input) + + #. a tuple of arguments corresponding to + + .. code-block:: python + + model.forward(*dummy_input) + + #. a tuple of arguments such that ``type(dummy_input[-1]) == dict`` corresponding to + + .. code-block:: python + + model.forward(*dummy_input[:-1], **dummy_input[-1]) + + .. warning:: + + In this case the model's ``forward()`` method **cannot** contain keyword-only + arguments (e.g. ``forward(..., *, kw_only_args)``) or variable keyword arguments + (e.g. ``forward(..., **kwargs)``) since these cannot be sorted into positional + arguments. + + .. note:: + + In order to pass a dict as last non-keyword argument, you need to use a tuple as + ``dummy_input`` and add an *empty* dict as the last element, e.g., + + .. code-block:: python + + dummy_input = (x, {"y": y, "z": z}, {}) + + The empty dict at the end will then be interpreted as the keyword args. + + See `torch.onnx.export `_ + for more info. + + Note that if you provide a ``{arg_name}`` with batch size ``b``, the results will be + computed based on batch size ``b``. + config: Additional optional arguments to configure the search. Currently, we support: + + * ``checkpoint``: Path to save/restore checkpoint with dictionary containing intermediate + search state. If provided, the intermediate search state will be automatically + restored before search (if exists) and stored/saved during search. + * ``verbose``: Whether to print detailed search space profiling and search stats during search. + * ``forward_loop``: A ``Callable`` that takes a model as input and runs a forward loop + on it. It is recommended to choose the data loader used inside the forward loop + carefully to reduce the runtime. Cannot be provided at the same time as + ``data_loader`` and ``collect_func``. + * ``data_loader``: An iterator yielding batches of data for calibrating the + normalization layers in the model or compute gradient scores. It is recommended to use + the same data loader as for training but with significantly fewer iterations. Cannot + be provided at the same time as ``forward_loop``. + * ``collect_func``: A ``Callable`` that takes a batch of data from the data loader as + input and returns the input to ``model.forward()`` as described in + :meth:`run_forward_loop `. Cannot + be provided at the same time as ``forward_loop``. + * ``max_iter_data_loader``: Maximum number of iterations to run the data loader. + * ``score_func``: A callable taking the model as input and returning a single accuracy/score + metric (float). This metric will be maximized during search. + + .. note:: + + The ``score_func`` is required only for ``fastnas`` mode. It will be + evaluated on models in eval mode (``model.eval()``). + * ``loss_func``: A ``Callable`` which takes the model output (i.e output of ``model.forward()``) + and the batch of data as its inputs and returns a scalar loss. + This is a required argument if the model is converted via ``gradnas`` mode. + + It should be possible to run a backward pass on the loss value returned by this method. + + ``collect_func`` will be used to gather the inputs to ``model.forward()`` + from a batch of data yielded by``data_loader``. + + ``loss_func`` should support the following usage: + + .. code-block:: python + + for i, batch in enumerate(data_loader): + if i >= max_iter_data_loader: + break + + # Assuming collect_func returns a tuple of arguments + output = model(*collect_func(batch)) + + loss = loss_func(output, batch) + loss.backward() + + .. note:: + + Additional configuration options may be added by individual algorithms. Please + refer to the documentation of the individual algorithms for more information. + + Returns: A tuple (subnet, state_dict) where + subnet is the searched subnet (nn.Module), which can be used for subsequent tasks like + fine-tuning, state_dict contains the history and detailed stats of the search procedure. + + .. note:: + + The given model is modified (exported) in-place to match the best subnet found by the + search algorithm. The returned subnet is thus a reference to the same model instance as the + input model. + """ + # apply prune mode(s) to model + model = apply_mode(model, mode, registry=PruneModeRegistry) + + # now run the search and return the result + return mtn.search( + model, + constraints, + dummy_input, + config=config, + ) diff --git a/modelopt/torch/quantization/__init__.py b/modelopt/torch/quantization/__init__.py new file mode 100644 index 0000000000..8f04f660fc --- /dev/null +++ b/modelopt/torch/quantization/__init__.py @@ -0,0 +1,25 @@ +# 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. + +"""Quantization package.""" + +# Initialize mode and plugins +from . import mode, plugins, utils + +# Add methods to mtq namespace +from .config import * +from .conversion import * +from .model_quant import * +from .nn.modules.quant_module import QuantModuleRegistry diff --git a/modelopt/torch/quantization/calib/__init__.py b/modelopt/torch/quantization/calib/__init__.py new file mode 100644 index 0000000000..c9f100a3d9 --- /dev/null +++ b/modelopt/torch/quantization/calib/__init__.py @@ -0,0 +1,24 @@ +# 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. + +"""Calibrator classes. + +``modelopt.torch.quantization.calib`` provides Calibrator classes that +collect data statistics and determine modelopt.torch.quantization parameters. +""" + +from .calibrator import * +from .histogram import * +from .max import * diff --git a/modelopt/torch/quantization/calib/calibrator.py b/modelopt/torch/quantization/calib/calibrator.py new file mode 100644 index 0000000000..cff7a75e73 --- /dev/null +++ b/modelopt/torch/quantization/calib/calibrator.py @@ -0,0 +1,63 @@ +# 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. + +"""Abstract base class for calibrators.""" + +__all__ = ["_Calibrator"] + + +class _Calibrator: + """Abstract base class of calibrators. + + Args: + num_bits: An integer. Number of bits of quantization. + axis: A tuple. see :class:`QuantizerAttributeConfig <..config.QuantizerAttributeConfig>`. + unsigned: A boolean. using unsigned quantization. + + Readonly Properties: + axis: + """ + + def __init__(self, num_bits=8, axis=None, unsigned=False): + """Initialize.""" + self._num_bits = num_bits + self._axis = axis + self._unsigned = unsigned + + def collect(self, x): + """Abstract method: collect tensor statistics used to compute amax. + + Args: + x: A tensor + """ + raise NotImplementedError + + def reset(self): + """Abstract method: reset calibrator to initial state.""" + raise NotImplementedError + + def compute_amax(self, *args, **kwargs): + """Abstract method: compute the amax from the collected data. + + Returns: + amax: a tensor + """ + raise NotImplementedError + + def __repr__(self): + s = "num_bits={_num_bits}" + s += " axis={_axis}" + s += " unsigned={_unsigned}" + return s.format(**self.__dict__) diff --git a/modelopt/torch/quantization/calib/histogram.py b/modelopt/torch/quantization/calib/histogram.py new file mode 100644 index 0000000000..721921d6fb --- /dev/null +++ b/modelopt/torch/quantization/calib/histogram.py @@ -0,0 +1,436 @@ +# 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. + +"""Histogram based calibrators.""" + +import warnings +from collections import Counter + +import numpy as np +import torch +from scipy.stats import entropy + +from modelopt.torch.utils import distributed as dist + +from .. import nn as qnn +from .. import utils as quant_utils +from ..tensor_quant import fake_tensor_quant, scaled_e4m3 +from .calibrator import _Calibrator + +__all__ = ["HistogramCalibrator", "calibrate_weights"] + + +class HistogramCalibrator(_Calibrator): + """Unified histogram calibrator. + + Histogram will be only collected once. compute_amax() performs entropy, percentile, or mse + calibration based on arguments + + Args: + num_bits: An integer. Number of bits of quantization. + axis: A tuple. see :class:`QuantizerAttributeConfig <..config.QuantizerAttributeConfig>`. + unsigned: A boolean. using unsigned quantization. + num_bins: An integer. Number of histograms bins. Default 2048. + grow_method: A string. DEPRECATED. default None. + skip_zeros: A boolean. If True, skips zeros when collecting data for histogram. Default False. + torch_hist: A boolean. If True, collect histogram by torch.histc instead of np.histogram. If input tensor + is on GPU, histc will also be running on GPU. Default True. + """ + + def __init__( + self, + num_bits=8, + axis=None, + unsigned=False, + num_bins=2048, + grow_method=None, + skip_zeros=False, + torch_hist=True, + ): + """Initialize.""" + super(HistogramCalibrator, self).__init__(num_bits, axis, unsigned) + self._num_bins = num_bins + self._skip_zeros = skip_zeros + + self._calib_bin_edges = None + self._calib_hist = None + + self._torch_hist = torch_hist + + if axis is not None: + raise NotImplementedError( + "Calibrator histogram collection only supports per tensor scaling" + ) + + def collect(self, x): + """Collect histogram.""" + if torch.min(x) < 0.0: + x = x.abs() + + x = x.float() + + if not self._torch_hist: + x_np = x.cpu().detach().numpy() + + if self._skip_zeros: + x_np = x_np[np.where(x_np != 0)] + + if self._calib_bin_edges is None and self._calib_hist is None: + # first time it uses num_bins to compute histogram. + self._calib_hist, self._calib_bin_edges = np.histogram(x_np, bins=self._num_bins) + else: + temp_amax = np.max(x_np) + if temp_amax > self._calib_bin_edges[-1]: # type: ignore + # increase the number of bins + width = self._calib_bin_edges[1] - self._calib_bin_edges[0] # type: ignore + # NOTE: np.arange may create an extra bin after the one containing temp_amax + new_bin_edges = np.arange( + self._calib_bin_edges[-1] + width, # type: ignore[index] + temp_amax + width, + width, + ) + self._calib_bin_edges = np.hstack((self._calib_bin_edges, new_bin_edges)) + hist, self._calib_bin_edges = np.histogram(x_np, bins=self._calib_bin_edges) + hist[: len(self._calib_hist)] += self._calib_hist # type: ignore + self._calib_hist = hist + else: + # This branch of code is designed to match numpy version as close as possible + with torch.no_grad(): + if self._skip_zeros: + x = x[torch.where(x != 0)] + + # Because we collect histogram on absolute value, setting min=0 simplifying the rare case where + # minimum value is not exactly 0 and first batch collected has larger min value than later batches + x_max = x.max() + if self._calib_bin_edges is None and self._calib_hist is None: + self._calib_hist = torch.histc(x, bins=self._num_bins, min=0, max=x_max) + self._calib_bin_edges = torch.linspace(0, x_max, self._num_bins + 1) + else: + if x_max > self._calib_bin_edges[-1]: # type: ignore + width = self._calib_bin_edges[1] - self._calib_bin_edges[0] # type: ignore + self._num_bins = int((x_max / width).ceil().item()) + self._calib_bin_edges = torch.arange( + 0, x_max + width, width, device=x.device + ) + + hist = torch.histc(x, bins=self._num_bins, min=0, max=self._calib_bin_edges[-1]) # type: ignore + hist[: self._calib_hist.numel()] += self._calib_hist # type: ignore + self._calib_hist = hist + + def reset(self): + """Reset the collected histogram.""" + self._calib_bin_edges = None + self._calib_hist = None + + def compute_amax( + self, + method: str, + *, + stride: int = 1, + start_bin: int = 128, + percentile: float = 99.99, + ): + """Compute the amax from the collected histogram. + + Args: + method: A string. One of ['entropy', 'mse', 'percentile'] + + Keyword Arguments: + stride: An integer. Default 1 + start_bin: An integer. Default 128 + percentils: A float number between [0, 100]. Default 99.99. + + Returns: + amax: a tensor + """ + if dist.is_initialized(): + warnings.warn( + "This method does not perform any synchronization across DistributedDataParallel" + " (DDP) https://pytorch.org/docs/stable/notes/ddp.html modules. The recommended" + " method is to use the same calibration dataset across all distributed data" + " parallel groups so that `amax` is the same for all DDP modules." + ) + + if isinstance(self._calib_hist, torch.Tensor): + calib_hist = self._calib_hist.int().cpu().numpy() + calib_bin_edges = self._calib_bin_edges.cpu().numpy() # type: ignore + else: + calib_hist = self._calib_hist + calib_bin_edges = self._calib_bin_edges + + if method == "entropy": + calib_amax = _compute_amax_entropy( + calib_hist, calib_bin_edges, self._num_bits, self._unsigned, stride, start_bin + ) + elif method == "mse": + calib_amax = _compute_amax_mse( + calib_hist, calib_bin_edges, self._num_bits, self._unsigned, stride, start_bin + ) + elif method == "percentile": + calib_amax = _compute_amax_percentile(calib_hist, calib_bin_edges, percentile) + else: + raise TypeError("Unknown calibration method {}".format(method)) + + return calib_amax + + def __str__(self): + s = "HistogramCalibrator(" + if self._calib_bin_edges is None: + bin_edge_str = "None" + else: + bin_edge_str = "[{:.3f}, ..., {:.3f}]({})".format( + self._calib_bin_edges[0], self._calib_bin_edges[-1], len(self._calib_bin_edges) + ) + s += "calib_bin_edges={})".format(bin_edge_str) + return s + + def __repr__(self): + s = "HistogramCalibrator(" + s += super(HistogramCalibrator, self).__repr__() + s += " calib_bin_edges={_calib_bin_edges}" + s += " calib_hist={_calib_hist})" + return s.format(**self.__dict__) + + +# Ideally, we want to decouple collector (collect histogram) and calibrator (compute amax) as opposed to +# the current calibrator design. The following compute amax functions are broken out from the calibrator +# as first step towards there. +def _compute_amax_entropy(calib_hist, calib_bin_edges, num_bits, unsigned, stride=1, start_bin=128): + """Returns amax that minimizes KL-Divergence of the collected histogram.""" + # If calibrator hasn't collected any data, return none + if calib_bin_edges is None and calib_hist is None: + return None + + def _normalize_distr(distr): + summ = np.sum(distr) + if summ != 0: + distr = distr / summ + + bins = calib_hist[:] + bins[0] = bins[1] + + total_data = np.sum(bins) + + divergences = [] + arguments = [] + + # we are quantizing to 128 values + sign if num_bits=8 + nbins = 1 << (num_bits - 1 + int(unsigned)) + + starting = start_bin + stop = len(bins) + + new_density_counts = np.zeros(nbins, dtype=np.float64) + + for i in range(starting, stop + 1, stride): + new_density_counts.fill(0) + space = np.linspace(0, i, num=nbins + 1) + digitized_space = np.digitize(range(i), space) - 1 + + digitized_space[bins[:i] == 0] = -1 + + for idx, digitized in enumerate(digitized_space): + if digitized != -1: + new_density_counts[digitized] += bins[idx] + + counter = Counter(digitized_space) + for key, val in counter.items(): + if key != -1: + new_density_counts[key] = new_density_counts[key] / val + + new_density = np.zeros(i, dtype=np.float64) + for idx, digitized in enumerate(digitized_space): + if digitized != -1: + new_density[idx] = new_density_counts[digitized] + + total_counts_new = np.sum(new_density) + np.sum(bins[i:]) + _normalize_distr(new_density) + + reference_density = np.array(bins[: len(digitized_space)]) + reference_density[-1] += np.sum(bins[i:]) + + total_counts_old = np.sum(reference_density) + if round(total_counts_new) != total_data or round(total_counts_old) != total_data: + raise RuntimeError( + "Count mismatch! total_counts_new={}, total_counts_old={}, total_data={}".format( + total_counts_new, total_counts_old, total_data + ) + ) + + _normalize_distr(reference_density) + + ent = entropy(reference_density, new_density) + divergences.append(ent) + arguments.append(i) + + divergences = np.array(divergences) + last_argmin = len(divergences) - 1 - np.argmin(divergences[::-1]) + calib_amax = calib_bin_edges[last_argmin * stride + starting] + calib_amax = torch.tensor(calib_amax.item()) + + return calib_amax + + +def _compute_amax_mse(calib_hist, calib_bin_edges, num_bits, unsigned, stride=1, start_bin=128): + """Returns amax that minimizes MSE of the collected histogram.""" + # If calibrator hasn't collected any data, return none + if calib_bin_edges is None and calib_hist is None: + return None + + counts = torch.from_numpy(calib_hist[:]).float() + edges = torch.from_numpy(calib_bin_edges[:]).float() + + device = None + if torch.cuda.is_available(): + device = counts.device + counts = counts.cuda() + edges = edges.cuda() + + centers = (edges[1:] + edges[:-1]) / 2 + + mses = [] + arguments = [] + + for i in range(start_bin, len(centers), stride): + amax = centers[i] + if isinstance(num_bits, int) and num_bits >= 0: + quant_centers = fake_tensor_quant(centers, amax, num_bits, unsigned) + elif num_bits == (4, 3): + quant_centers = scaled_e4m3(centers, amax, num_bits[0], num_bits[1]) + else: + raise TypeError("Invalid num_bits. num_bits must be a positive integer or tuple (4,3).") + + mse = ((quant_centers - centers) ** 2 * counts).mean() + + mses.append(mse.cpu()) + arguments.append(i) + + argmin = np.argmin(mses) + calib_amax = centers[arguments[argmin]] + + if device is not None: + calib_amax = calib_amax.to(device) + + return calib_amax + + +def _compute_amax_percentile(calib_hist, calib_bin_edges, percentile): + """Returns amax that clips the percentile fraction of collected data.""" + if percentile < 0 or percentile > 100: + raise ValueError("Invalid percentile. Must be in range 0 <= percentile <= 100.") + + # If calibrator hasn't collected any data, return none + if calib_bin_edges is None and calib_hist is None: + return None + + total = calib_hist.sum() + cdf = np.cumsum(calib_hist / total) + idx = np.searchsorted(cdf, percentile / 100) + calib_amax = calib_bin_edges[idx] + calib_amax = torch.tensor(calib_amax.item()) + + return calib_amax + + +def calibrate_weights(model, method="percentile", perchannel=True, percentile=99.99, num_bins=2048): + """Calibrate weights of all child quantized modules. + + Ideally, we would split calibration functionality to histogram collector and calibrator which + takes histogram and compute amax. But since we haven't decoupled collector and calibrator, it + is easier to create a separate function to calibrate weight. + + .. note:: + This function uses `method` specified by the argument to decide which method to use, NOT the one + specified by the calibrator embedded in weight_quantizer. + We haven't moved calibration to GPU, so everything is transfered to CPU + + Args: + model: A torch.nn.Module. + method: A string of calibration method. Supports "mse" and "percentile". Default "percentile" + perchannel: A bool. Set channel/neuron axis if True. Default True. + percentile: A float. Default 99.99 + num_bins: A integer. Number of bins of histogram. Default 2048. + + """ + for name, module in model.named_modules(): + if hasattr(module, "weight") and hasattr(module, "weight_quantizer"): + num_bits = module.weight_quantizer.num_bits + unsigned = module.weight_quantizer.unsigned + channel_second_modules = ( + qnn.QuantConvTranspose1d, + qnn.QuantConvTranspose2d, + qnn.QuantConvTranspose3d, + ) + if perchannel: + axis = 1 if isinstance(module, channel_second_modules) else 0 + else: + axis = None + axis_size = module.weight.shape[axis] if axis is not None else 1 + + # Histogram is always collected even if method is "max". Although "max" is supported here + # but it is not the primary usage of this function + if axis is None: + input_weights = module.weight.abs().cpu().detach().numpy() + calib_hist, calib_bin_edges = np.histogram( + input_weights, bins=2048, range=(0, input_weights.max()) + ) + calib_hist = [calib_hist] + calib_bin_edges = [calib_bin_edges] + else: + calib_hist = [] + calib_bin_edges = [] + for i in range(axis_size): + input_weights = ( + module.weight.index_select( + axis, torch.tensor(i, device=module.weight.device) + ) + .abs() + .cpu() + .detach() + .numpy() + ) + hist, bin_edges = np.histogram( + input_weights, bins=num_bins, range=(0, input_weights.max()) + ) + calib_hist.append(hist) + calib_bin_edges.append(bin_edges) + + calib_amax = [] + if method == "max": + reduce_axis = list(range(module.weight.dim())) + reduce_axis.remove(axis) # type: ignore + calib_amax.append(quant_utils.reduce_amax(module.weight, axis=reduce_axis)) + elif method == "mse": + for i in range(axis_size): + calib_amax.append( + _compute_amax_mse(calib_hist[i], calib_bin_edges[i], num_bits, unsigned) + ) + elif method == "percentile": + for i in range(axis_size): + calib_amax.append( + _compute_amax_percentile(calib_hist[i], calib_bin_edges[i], percentile) + ) + else: + raise TypeError("Unsupported calibration method {}".format(method)) + + if axis is None: + calib_amax_t = calib_amax[0] + else: + calib_amax_shape = [1] * module.weight.dim() + calib_amax_shape[axis] = module.weight.shape[axis] + calib_amax_t = torch.stack(calib_amax).reshape(calib_amax_shape) + if calib_amax_t.numel() == 1: + calib_amax_t.squeeze_() + module.weight_quantizer.amax = calib_amax_t.detach().cpu().numpy() diff --git a/modelopt/torch/quantization/calib/max.py b/modelopt/torch/quantization/calib/max.py new file mode 100644 index 0000000000..221c88c77a --- /dev/null +++ b/modelopt/torch/quantization/calib/max.py @@ -0,0 +1,103 @@ +# 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. + +"""Calibrator that returns the absolute max of all collected tensors.""" + +import torch + +from .. import utils as quant_utils +from .calibrator import _Calibrator + +__all__ = ["MaxCalibrator"] + + +class MaxCalibrator(_Calibrator): + """Max calibrator, tracks the maximum value globally. + + Args: + calib_desc: A MaxCalibDescriptor. + num_bits: An integer. Number of bits of quantization. + axis: A tuple. see :class:`QuantizerAttributeConfig <..config.QuantizerAttributeConfig>`. + unsigned: A boolean. using unsigned quantization. + + Readonly Properties: + amaxs: A list of amax. Numpy array is saved as it is likely to be used for some plot. + """ + + def __init__(self, num_bits=8, axis=None, unsigned=False, track_amax=False): + """Initialize.""" + super(MaxCalibrator, self).__init__(num_bits, axis, unsigned) + self._track_amax = track_amax + if self._track_amax: + self._amaxs = [] # shall we have a better name? + self._calib_amax = None + + @property + def amaxs(self): + """Returns the list of amax`s collected so far.""" + return self._amaxs + + @torch.no_grad() + def collect(self, x): + """Tracks the absolute max of all tensors. + + Args: + x: A tensor + + Raises: + RuntimeError: If amax shape changes + """ + # Swap axis to reduce. + axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis] + # Handle negative axis. + axis = [x.dim() + i if isinstance(i, int) and i < 0 else i for i in axis] + reduce_axis = [] + for i in range(x.dim()): + if i not in axis: + reduce_axis.append(i) + local_amax = quant_utils.reduce_amax(x, axis=reduce_axis).detach() + if self._calib_amax is None: + self._calib_amax = local_amax + else: + if local_amax.shape != self._calib_amax.shape: + raise RuntimeError("amax shape changed!") + self._calib_amax = torch.max(self._calib_amax, local_amax) + + if self._track_amax: + self._amaxs.append(local_amax.cpu().numpy()) + + def reset(self): + """Reset the collected absolute max.""" + self._calib_amax = None + + def compute_amax(self): + """Return the absolute max of all tensors collected.""" + return self._calib_amax + + def __str__(self): + s = "MaxCalibrator(" + s += "track_amax={_track_amax}" + s += ")" + return s.format(**self.__dict__) + + def __repr__(self): + s = "MaxCalibrator(" + s += super(MaxCalibrator, self).__repr__() + s += " calib_amax={_calib_amax}" + s += " track_amax={_track_amax}" + if self._track_amax: + s += " amaxs={_amaxs}" + s += ")" + return s.format(**self.__dict__) diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py new file mode 100644 index 0000000000..a995e63d9d --- /dev/null +++ b/modelopt/torch/quantization/config.py @@ -0,0 +1,812 @@ +# 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. + +"""This document lists the quantization formats supported by Model Optimizer and example quantization configs. + +.. _quantization-formats: + +Quantization Formats +========================================== + +The following table lists the quantization formats supported by Model Optimizer and the corresponding quantization +config. See :ref:`Quantization Configs ` for the +specific quantization config definitions. + +Please see :doc:`choosing the right quantization formats <../../guides/_choosing_quant_methods>` to +learn more about the formats and their use-cases. + +.. note:: + + The recommended configs given below are for LLM models. For CNN models, only INT8 quantization + is supported. Please use quantization config ``INT8_DEFAULT_CFG`` for CNN models. + +================================= ======================================================= +Quantization Format Model Optimizer config +================================= ======================================================= +INT8 ``INT8_SMOOTHQUANT_CFG`` + +FP8 ``FP8_DEFAULT_CFG`` + +INT4 Weights only AWQ (W4A16) ``INT4_AWQ_CFG`` + +INT4-FP8 AWQ (W4A8) ``W4A8_AWQ_BETA_CFG`` + +================================= ======================================================= + +.. _quantization-configs: + +Quantization Configs +================================ + +Quantization config is dictionary specifying the values for keys ``"quant_cfg"`` and +``"algorithm"``. The ``"quant_cfg"`` key specifies the quantization configurations. The +``"algorithm"`` key specifies the ``algorithm`` argument to +:meth:`calibrate `. Please see :class:`QuantizeConfig` +for the quantization config definition. + +'Quantization configurations' is a dictionary mapping wildcards or filter functions +to its 'quantizer attributes'. The wildcards or filter functions are matched +against the quantizer module names. The quantizer modules have names ending with +``weight_quantizer`` and ``input_quantizer`` and they perform weight quantization and +input quantization (or activation quantization) respectively. The quantizer modules are generally +instances of +:class:`TensorQuantizer `. +The quantizer attributes are defined by :class:`QuantizerAttributeConfig`. See :class:`QuantizerAttributeConfig` +for details on the quantizer attributes and their values. + +The key `"default"` from the quantization configuration dictionary is applied if no other wildcard or filter functions +match the quantizer module name. + +The quantizer attributes are applied in the order they are specified. For the missing attributes, the default attributes +as defined by :class:`QuantizerAttributeConfig` are used. + +Quantizer attributes can also be a list of dictionaries. In this case, the matched quantizer module +is replaced with a +:class:`SequentialQuantizer ` +module which is used to quantize a tensor in multiple formats sequentially. Each quantizer attribute +dictionary in the list specifies the quantization formats for each quantization step of the +sequential quantizer. For example, `SequentialQuantizer` is used in 'INT4 Weights, FP8 Activations' +quantization in which the weights are quantized in INT4 followed by FP8. + +In addition, the dictionary entries could also be pytorch module class names mapping the class specific +quantization configurations. The pytorch modules should have a quantized equivalent. + +To get the string representation of a module class, do: + +.. code-block:: + + from modelopt.torch.quantization import QuantModuleRegistry + + # Get the class name for nn.Conv2d + class_name = QuantModuleRegistry.get_key(nn.Conv2d) + +Here is an example of a quantization config: + +.. code-block:: + + MY_QUANT_CFG = { + "quant_cfg": { + # Quantizer wildcard strings mapping to quantizer attributes + "*weight_quantizer": {"num_bits": 8, "axis": 0}, + "*input_quantizer": {"num_bits": 8, "axis": None}, + + # Module class names mapping to quantizer configurations + "nn.LeakyReLU": {"*input_quantizer": {"enable": False}}, + + } + } + +.. _example-quantization-configs: + +Example Quantization Configurations +========================================== + +Here are the recommended quantization configs from Model Optimizer for +quantization formats such as FP8, INT8, INT4, etc.: + +.. code-block:: + + INT8_DEFAULT_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": 8, "axis": 0}, + "*input_quantizer": {"num_bits": 8, "axis": None}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "default": {"enable": False}, + }, + "algorithm": "max", + } + + INT8_SMOOTHQUANT_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": 8, "axis": 0}, + "*input_quantizer": {"num_bits": 8, "axis": -1}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "nn.Conv2d": { + "*weight_quantizer": {"num_bits": 8, "axis": 0}, + "*input_quantizer": {"num_bits": 8, "axis": None}, + }, + "default": {"enable": False}, + }, + "algorithm": "smoothquant", + } + + FP8_DEFAULT_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": (4, 3), "axis": None}, + "*input_quantizer": {"num_bits": (4, 3), "axis": None}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "default": {"enable": False}, + }, + "algorithm": "max", + } + + INT4_BLOCKWISE_WEIGHT_ONLY_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": 4, "block_sizes": {-1: 128}, "enable": True}, + "*input_quantizer": {"enable": False}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "default": {"enable": False}, + }, + "algorithm": "max", + } + + INT4_AWQ_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": 4, "block_sizes": {-1: 128}, "enable": True}, + "*input_quantizer": {"enable": False}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "default": {"enable": False}, + }, + "algorithm": {"method": "awq_lite", "alpha_step": 0.1}, + # "algorithm": {"method": "awq_full", "alpha_step": 0.1, "max_co_batch_size": 1024}, + # "algorithm": {"method": "awq_clip", "max_co_batch_size": 2048}, + } + + W4A8_AWQ_BETA_CFG = { + "quant_cfg": { + "*weight_quantizer": [ + {"num_bits": 4, "block_sizes": {-1: 128}, "enable": True}, + {"num_bits": (4, 3), "axis": None, "enable": True}, + ], + "*input_quantizer": {"num_bits": (4, 3), "axis": None, "enable": True}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "default": {"enable": False}, + }, + "algorithm": "awq_lite", + } + +These config can be accessed as attributes of ``modelopt.torch.quantization`` and can be given as +input to :meth:`mtq.quantize() `. For example: + +.. code-block:: + + import modelopt.torch.quantization as mtq + model = mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop) + +You can also create your own config by following these examples. +For instance, if you want to quantize a model with int4 AWQ algorithm, but need to skip quantizing +the layer named ``lm_head``, you can create a custom config and quantize your model as following: + +.. code-block:: + + # Create custom config + CUSTOM_INT4_AWQ_CFG = copy.deepcopy(mtq.INT4_AWQ_CFG) + CUSTOM_INT4_AWQ_CFG["quant_cfg"]["*lm_head*"] = {"enable": False} + + # quantize model + model = mtq.quantize(model, CUSTOM_INT4_AWQ_CFG, forward_loop) + +""" + +from typing import Callable, Optional, Union + +from pydantic import ValidationInfo, field_validator, model_validator + +from modelopt.core.torch.quantization.config import ( # noqa: F401 + NVFP4_AWQ_CLIP_CFG, + NVFP4_AWQ_FULL_CFG, + NVFP4_AWQ_LITE_CFG, + NVFP4_DEFAULT_CFG, + NVFP4_REAL_QUANT_CFG, +) +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField +from modelopt.torch.utils.network import ConstructorLike + +_default_disabled_quantizer_cfg = { + "nn.BatchNorm1d": {"*": {"enable": False}}, + "nn.BatchNorm2d": {"*": {"enable": False}}, + "nn.BatchNorm3d": {"*": {"enable": False}}, + "nn.LeakyReLU": {"*": {"enable": False}}, +} + +INT8_DEFAULT_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": 8, "axis": 0}, + "*input_quantizer": {"num_bits": 8, "axis": None}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": "max", +} + +INT8_SMOOTHQUANT_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": 8, "axis": 0}, + "*input_quantizer": {"num_bits": 8, "axis": None}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": "smoothquant", +} + +FP8_DEFAULT_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": (4, 3), "axis": None}, + "*input_quantizer": {"num_bits": (4, 3), "axis": None}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": "max", +} + +FP8_PER_CHANNEL_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": (4, 3), "axis": 0}, + "*input_quantizer": {"num_bits": (4, 3), "axis": -1, "type": "dynamic"}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": "max", +} + +FP8_WA_FP8_KV_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": (4, 3), "axis": None}, + "*input_quantizer": {"num_bits": (4, 3), "axis": None}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*bmm_quantizer": {"num_bits": (4, 3), "axis": None}, + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": "max", +} + +INT4_BLOCKWISE_WEIGHT_ONLY_CFG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": 4, "block_sizes": {-1: 128}, "enable": True}, + "*input_quantizer": {"enable": False}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": "max", +} + +NF4_REAL_QUANT_CFG = { + "quant_cfg": { + "*weight_quantizer": { + "num_bits": 4, + "block_sizes": {-1: 64, "scale_bits": 8, "scale_block_sizes": {-1: 256}}, + "enable": True, + "fake_quant": False, + }, + "*input_quantizer": {"enable": False}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + "default": {"enable": False}, + }, + "algorithm": {"method": "real_quantize"}, +} + +INT4_AWQ_REAL_QUANT_CFG = { + "quant_cfg": { + "*weight_quantizer": { + "num_bits": 4, + "block_sizes": {-1: 128, "type": "static"}, + "enable": True, + "fake_quant": False, + }, + "*input_quantizer": {"enable": False}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": { + "method": "real_quantize", + "additional_algorithm": {"method": "awq_lite", "alpha_step": 0.1}, + }, +} + +INT4_AWQ_CFG = { + "quant_cfg": { + "*weight_quantizer": { + "num_bits": 4, + "block_sizes": {-1: 128, "type": "static"}, + "enable": True, + }, + "*input_quantizer": {"enable": False}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": {"method": "awq_lite", "alpha_step": 0.1}, + # "algorithm": {"method": "awq_full", "alpha_step": 0.1, "max_co_batch_size": 1024}, + # "algorithm": {"method": "awq_clip", "max_co_batch_size": 2048}, +} + +# W4A8 currently uses INT4 blockwise quantization (block size = 128) followed by FP8 quantization +# for weights. This could change in the future +W4A8_AWQ_BETA_CFG = { + "quant_cfg": { + "*weight_quantizer": [ + {"num_bits": 4, "block_sizes": {-1: 128, "type": "static"}, "enable": True}, + {"num_bits": (4, 3), "axis": None, "enable": True}, + ], + "*input_quantizer": {"num_bits": (4, 3), "axis": None, "enable": True}, + "*lm_head*": {"enable": False}, + "*block_sparse_moe.gate*": {"enable": False}, # Skip the MOE router + "*router*": {"enable": False}, # Skip the MOE router + "*output_layer*": {"enable": False}, + "output.*": {"enable": False}, + **_default_disabled_quantizer_cfg, + "default": {"enable": False}, + }, + "algorithm": "awq_lite", +} + + +choices: set[str] = { + "INT8_DEFAULT_CFG", + "INT8_SMOOTHQUANT_CFG", + "FP8_DEFAULT_CFG", + "INT4_BLOCKWISE_WEIGHT_ONLY_CFG", + "INT4_AWQ_CFG", + "W4A8_AWQ_BETA_CFG", + "NVFP4_DEFAULT_CFG", + "NVFP4_AWQ_LITE_CFG", + "NVFP4_AWQ_CLIP_CFG", + "NVFP4_AWQ_FULL_CFG", + "NF4_REAL_QUANT_CFG", + "NVFP4_REAL_QUANT_CFG", + "INT4_AWQ_REAL_QUANT_CFG", +} + + +class QuantizerAttributeConfig(ModeloptBaseConfig): + """Quantizer attribute type.""" + + enable: bool = ModeloptField( + default=True, + title="Enable quantizer.", + description="""If True, enables the quantizer. If False, by-pass the quantizer and returns the input tensor.""", + ) + + num_bits: Union[int, tuple[int, int]] = ModeloptField( + default=8, + title="An integer or a tuple of two integers specifying the number of quantization bits.", + description="""`num_bits` can be: + + #. A positive integer argument for integer quantization. `num_bits` specify + the number of bits used for integer quantization. + + #. Constant integer tuple (E,M) for floating point quantization emulating + Nvidia's FPx quantization. E is the number of exponent bits and M is the number + of mantissa bits. Supported FPx quantization formats: FP8 (E4M3, E5M2), FP6(E3M2, E2M3), FP4(E2M1).""", + ) + + @model_validator(mode="after") + def validate_num_bits(self): + """Validate `num_bits`.""" + num_bits = self.num_bits + if not isinstance(num_bits, tuple): + return self + + block_sizes = self.block_sizes + if num_bits not in [ + (4, 3), + (5, 2), + (2, 1), + (1, 2), + (0, 3), + (3, 0), + (3, 2), + (2, 3), + ]: + raise ValueError( + "Supported FPx quantization formats: FP8 (E4M3, E5M2), FP6(E3M2, E2M3), FP4(E2M1)." + ) + elif num_bits != (4, 3) and ( + block_sizes is None or block_sizes.get("type", None) != "dynamic" + ): + raise ValueError( + ( + "Only blockwise dynamic quantization is supported with quantization " + "formats E{num_bis[0]}M{num_bits[1]}." + ) + ) + return self + + axis: Optional[Union[int, tuple[int, ...]]] = ModeloptField( + default=None, + title="None, integer or a tuple of integers specifying the axis to quantize.", + description="""The specified axis/axes will have its own amax for + computing scaling factor. If None (the default), use per tensor scale. Must be in the + range [-rank(input_tensor), rank(input_tensor)). E.g. For a KCRS weight tensor, + ``quant_axis=(0)`` will yield per channel scaling.""", + ) + + fake_quant: bool = ModeloptField( + default=True, + title="Enable fake quantization.", + description="""If True, enable fake quantization.""", + ) + + unsigned: bool = ModeloptField( + default=False, + title="Enable unsigned quantization.", + description="""If True, enable unsigned quantization. Used only for integer quantization.""", + ) + + narrow_range: bool = ModeloptField( + default=False, + title="Enable narrow range quantization.", + description="""If True, enable narrow range quantization. Used only for integer quantization.""", + ) + + learn_amax: bool = ModeloptField( + default=False, + title="Enable learning amax.", + description="""If True, enable learning amax.""", + ) + + type: str = ModeloptField( + default="static", + title="""Specify whether the quantization is static or dynamic.""", + description="""The value is a string from ``["static", "dynamic"]``. + If ``"dynamic"``, dynamic quantization will be enabled which does not collect any statistics during + calibration.""", + pattern=r"^static$|^dynamic$", + ) + + @field_validator("type") + @classmethod + def validate_type(cls, v, info: ValidationInfo): + """Validate type.""" + if v == "dynamic": + assert info.data["learn_amax"] is False + return v + + block_sizes: Optional[ + dict[Union[int, str], Union[int, tuple[int, int], str, dict[int, int]]] + ] = ModeloptField( + default=None, + title="Optional dictionary specifying block quantization parameters.", + description="""The keys are the axes for block quantization and the + values are block sizes for quantization along the respective axes. Keys must be in the + range ``[-tensor.dim(), tensor.dim())``. Values, which are the block sizes + for quantization must be positive integers. + + In addition, there can be special string keys ``"type"``, ``"scale_bits"`` and ``"scale_block_sizes"``. + + Key ``"type"`` should map to ``"dynamic"`` or ``"static"`` where ``"dynamic"`` + indicates dynamic block quantization and "static" + indicates static calibrated block quantization. By default, the type is ``"static"``. + + Key ``"scale_bits"`` specify the quantization bits for the per-block quantization scale factor + (i.e a double quantization scheme). + + Key ``"scale_block_sizes"`` specify the block size for double quantization. + By default per-block quantization scale is not quantized. + + For example, ``block_sizes = {-1: 32}`` will quantize the last axis of the input tensor in + blocks of size 32 with static calibration and ``block_sizes = {-1: 32, "type": "dynamic"}`` + will perform dynamic block quantization. If None, block + quantization is not performed. ``axis`` must be None when ``block_sizes`` is not None. + """, + ) + + @staticmethod + def _get_block_quant_axes_and_sizes(block_sizes): + if block_sizes is None: + return None + return { + k: v + for k, v in block_sizes.items() + if k not in ["type", "scale_bits", "scale_block_sizes"] + } + + @field_validator("block_sizes") + @classmethod + def validate_block_sizes(cls, v, info: ValidationInfo): + """Validate block sizes.""" + if v is None: + return v + assert info.data["axis"] is None, "axis must be None when block_sizes is not None." + if v.get("type", None) == "dynamic": + assert ( + len(cls._get_block_quant_axes_and_sizes(v)) == 1 + ), "Dynamic block quantization only supports quantization last axis." + for _k, _v in v.items(): + if isinstance(_k, str): + assert _k in ["type", "scale_bits", "scale_block_sizes"] + else: + assert isinstance(_k, int) and isinstance(_v, int) + return v + + trt_high_precision_dtype: str = ModeloptField( + default="Float", + title="TRT StronglyType requires all weights and amax to be in the same dtype.", + description="""The value is a string from ``["Float", "Half", "BFloat16"]``. + The QDQs will be assigned the appropriate data type, and this variable will only be + used when the user is exporting the quantized ONNX model.""", + pattern=r"^Float$|^Half$|^BFloat16$", + ) + + calibrator: Union[str, ConstructorLike] = ModeloptField( + default="max", + title="""Specify the calibrator to use.""", + description="""The calibrator can be a string from ``["max", "histogram"]`` or a constructor + to create a calibrator which subclasses :class:`_Calibrator `. + See :meth:`standardize_constructor_args ` + for more information on how to specify the constructor.""", + ) + + @field_validator("calibrator") + @classmethod + def validate_calibrator(cls, v, info: ValidationInfo): + """Validate calibrator.""" + if isinstance(v, str): + assert v in ["max", "histogram"] + return v + + +class QuantizeAlgorithmConfig(ModeloptBaseConfig): + """Calibration algorithm config base.""" + + method: str = ModeloptField( + default="max", + title="The calibration algorithm.", + description="""The algorithm used for calibration. Supported algorithms include + ``"max", "smoothquant", "awq_lite", "awq_full", and "awq_clip"``.""", + ) + + +class MaxCalibConfig(QuantizeAlgorithmConfig): + """The config for max calibration algorithm. + + Max calibration estimates max values of activations or weights and use this max values + to set the quantization scaling factor. + See `Integer Quantization `_ for the concepts. + """ + + +class SmoothQuantCalibConfig(QuantizeAlgorithmConfig): + """The config for ``smoothquant`` algorithm (SmoothQuant). + + SmoothQuant applies a smoothing factor which balances the scale of outliers in weights and activations. + See `SmoothQuant paper `_ for more details. + """ + + alpha: Optional[float] = ModeloptField( + default=1.0, + ge=0.0, + le=1.0, + title="SmoothQuant hyper-parameter alpha.", + description=( + "This hyper-parameter controls the migration strength." + "The migration strength is within [0, 1], " + "a larger value migrates more quantization difficulty to weights." + ), + ) + + +class AWQLiteCalibConfig(QuantizeAlgorithmConfig): + """The config for ``awq_lite`` (AWQ lite) algorithm. + + AWQ lite applies a channel-wise scaling factor which minimizes the output difference after quantization. + See `AWQ paper `_ for more details. + """ + + alpha_step: Optional[float] = ModeloptField( + default=0.1, + gt=0.0, + le=1.0, + title="Step size for the searching alpha.", + description="The alpha will be searched from 0 to 1 with the step size specified.", + ) + + debug: Optional[bool] = ModeloptField( + default=False, + title="Debug mode.", + description="If True, module's search metadata will be kept as a module attribute named `awq_lite`.", + ) + + +class AWQClipCalibConfig(QuantizeAlgorithmConfig): + """The config for ``awq_clip`` (AWQ clip) algorithm. + + AWQ clip searches clipped amax for per-group quantization, This search requires much more compute + compared to AWQ lite. To avoid any OOM, the linear layer weights are batched along the ``out_features`` + dimension of batch size ``max_co_batch_size``. AWQ clip calibration also takes longer than AWQ lite. + """ + + max_co_batch_size: Optional[int] = ModeloptField( + default=1024, + title="Maximum output channel batch size while searching clip values.", + description="Reduce this number if CUDA Out of Memory error occurs.", + ) + + max_tokens_per_batch: Optional[int] = ModeloptField( + default=64, + title="Maximum tokens per batch while searching clip values.", + description="""The total tokens used for clip search would be ``max_tokens_per_batch * number of batches``. + Original AWQ uses a total of 512 tokens to search for clip values.""", + ) + + min_clip_ratio: Optional[float] = ModeloptField( + default=0.5, + gt=0.0, + lt=1.0, + title="Minimum clip ratio to search for.", + description="""It should be in (0, 1.0). Clip will search for the optimal clipping value in the range + ``[original block amax * min_clip_ratio, original block amax]``.""", + ) + + shrink_step: Optional[float] = ModeloptField( + default=0.05, + gt=0.0, + le=1.0, + title="Step size to search for clip values.", + description="""It should be in range (0, 1.0]. The clip ratio will be searched from ``min_clip_ratio`` to 1 + with the step size specified.""", + ) + + debug: Optional[bool] = ModeloptField( + default=False, + title="Debug mode.", + description="If True, module's search metadata will be kept as a module attribute named ``awq_clip``.", + ) + + +class AWQFullCalibConfig(AWQLiteCalibConfig, AWQClipCalibConfig): + """The config for ``awq`` or ``awq_full`` algorithm (AWQ full). + + AWQ full performs ``awq_lite`` followed by ``awq_clip``. + """ + + debug: Optional[bool] = ModeloptField( + default=False, + title="Debug mode.", + description=( + "If True, module's search metadata will be kept as " + "module attributes named ``awq_lite`` and ``awq_clip``." + ), + ) + + +class RealQuantizeConfig(QuantizeAlgorithmConfig): + """The config for real quantization config. + + The ``additional_algorithm`` will be used for calibration before quantizing weights into low precision. + """ + + additional_algorithm: Optional[ + Union[AWQLiteCalibConfig, AWQClipCalibConfig, AWQFullCalibConfig] + ] = ModeloptField( + default="", + title="Additional algorithm for calibration before applying real quantization.", + description="""The algorithm used for calibration. Supported algorithms include + ``"awq_lite", "awq_full", and "awq_clip"``.""", + ) + + +QuantizeQuantCfgType = dict[ + Union[str, Callable], + Union[ + QuantizerAttributeConfig, + list[QuantizerAttributeConfig], + dict[ + Union[str, Callable], + Union[QuantizerAttributeConfig, list[QuantizerAttributeConfig]], + ], + ], +] + +CalibAlgorithmCfgType = Union[ + MaxCalibConfig, + SmoothQuantCalibConfig, + AWQLiteCalibConfig, + AWQClipCalibConfig, + AWQFullCalibConfig, + RealQuantizeConfig, +] + + +class QuantizeConfig(ModeloptBaseConfig): + """Default configuration for ``quantize`` mode.""" + + quant_cfg: QuantizeQuantCfgType = ModeloptField( + default={"default": {"num_bits": 8, "axis": None}}, + title="Quantization configuration", + validate_default=True, + ) + + algorithm: Union[None, str, CalibAlgorithmCfgType] = ModeloptField( + default="max", + title="Calibration algorithm", + validate_default=True, + ) + + @field_validator("algorithm") + @classmethod + def validate_calibrator(cls, v, info: ValidationInfo): + """Validate algorithm.""" + if isinstance(v, str): + assert v in ["max", "smoothquant", "awq_lite", "awq_clip", "awq_full", "real_quantize"] + return v + + +class _QuantizeExportConfig(ModeloptBaseConfig): + """An empty config.""" diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py new file mode 100644 index 0000000000..6a580f49f3 --- /dev/null +++ b/modelopt/torch/quantization/conversion.py @@ -0,0 +1,337 @@ +# 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. + +"""Quantization conversion/restore utilities.""" + +import fnmatch +from contextlib import contextmanager +from typing import Any, Callable, Union + +import torch.nn as nn + +from modelopt.torch.opt.conversion import ApplyModeError, ModelLikeModule, ModeloptStateManager +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict +from modelopt.torch.utils import get_unwrapped_name + +from .config import ( + QuantizeConfig, + QuantizeQuantCfgType, + QuantizerAttributeConfig, + _QuantizeExportConfig, +) +from .nn import QuantLinearConvBase, QuantModuleRegistry, SequentialQuantizer, TensorQuantizer +from .utils import is_quantized, is_quantized_layer_with_weight + +__all__ = [ + "replace_quant_module", + "set_quantizer_by_cfg", + "set_quantizer_attribute", + "register", + "unregister", + "set_quantizer_by_cfg_context", +] + + +def convert_to_quantized_model(model: nn.Module, config: QuantizeConfig) -> ConvertReturnType: + """Convert the model to a quantized one as per `config`.""" + # initialize the true module if necessary + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + + replace_quant_module(model, version=ModeloptStateManager(model).state_version) + set_quantizer_by_cfg(model, config["quant_cfg"]) + + metadata = {"quantizer_state": quantizer_state(model)} + + return model, metadata + + +def restore_quantized_model( + model: nn.Module, config: QuantizeConfig, metadata: MetadataDict +) -> nn.Module: + """Restores the quantizer states from the given state dict.""" + # initialize the true module if necessary + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + + assert not is_quantized(model), "Model must not be quantized!" + + quantizer_state_dict = metadata["quantizer_state"] + + replace_quant_module(model, version=ModeloptStateManager(model).state_version) + set_quantizer_by_cfg(model, config["quant_cfg"]) + + unmatched_keys = quantizer_state_dict.keys() - quantizer_state(model).keys() + extra_keys = quantizer_state(model).keys() - quantizer_state_dict.keys() + + if unmatched_keys: + raise ApplyModeError(f"Unmatched keys in quantizer state_dict: {unmatched_keys}") + if extra_keys: + raise ApplyModeError(f"Extra keys in quantizer state_dict: {extra_keys}") + + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer): + module.set_from_modelopt_state(quantizer_state_dict[name], name) + + for name, module in model.named_modules(): + if is_quantized_layer_with_weight(module): + with SequentialQuantizer.replace_sequential_quantizer_with_single_quantizer(module): + QuantLinearConvBase.sanitize_dummy_weight(module) + # Disable fake_quant for real quantization because we need to calibrate the model + with set_quantizer_by_cfg_context(model, {"*weight_quantizer": {"fake_quant": True}}): + QuantLinearConvBase.initialize_quantizer_with_dummy_states(module) + with SequentialQuantizer.replace_sequential_quantizer_with_single_quantizer(module): + QuantLinearConvBase.initialize_real_qtensor_with_dummy_weight(module) + if isinstance(module, TensorQuantizer): + module.clean_up_after_set_from_modelopt_state(name) + + return model + + +def update_quantize_metadata( + model: nn.Module, config: QuantizeConfig, metadata: MetadataDict +) -> None: + """Update the quantizer state in the metadata dict.""" + metadata["quantizer_state"] = quantizer_state(model) + + +def quantizer_state(model: nn.Module) -> dict[str, Any]: + """Returns the quantizer state dict describing the quantizer states in the model.""" + return { + get_unwrapped_name(n): m.get_modelopt_state() + for n, m in model.named_modules() + if isinstance(m, (TensorQuantizer, SequentialQuantizer)) + } + + +def replace_quant_module(model: nn.Module, version=None): + """Recursively replace the module with quantized module.""" + from .plugins.custom import register_custom_model_plugins_on_the_fly + + assert not is_quantized(model), "Model must not be quantized!" + + register_custom_model_plugins_on_the_fly(model) + + # If the model has a corresponding quantization module, replace it with it's quantized module + if type(model) in QuantModuleRegistry: + model = QuantModuleRegistry.convert(model) + + def _replace_quant_module(model): + for name, module in model.named_children(): + if type(module) in QuantModuleRegistry: + setattr(model, name, QuantModuleRegistry.convert(module)) + # Setting version for MCore checkpointing + if hasattr(module, "_modelopt_load_version"): + module._modelopt_load_version = version + + # Continue replacing in case of nested quantization as well + _replace_quant_module(getattr(model, name)) + + _replace_quant_module(model) + + replaced_modules = sum(isinstance(m, TensorQuantizer) for _, m in model.named_modules()) + print(f"Inserted {replaced_modules} quantizers") + + +def set_quantizer_by_cfg(quant_model: nn.Module, quant_cfg: Union[QuantizeQuantCfgType, dict]): + """Update the quantizer attributes based on the specified `quant_cfg`. + + `quant_cfg` is a dictionary mapping wildcards or filter functions + to its quantizer attributes which are defined in + :class:`QuantizerAttributeConfig <.config.QuantizerAttributeConfig>`. + The wildcards or filter functions are matched against the quantizer module names. + The specified quantizer attributes of the matched quantizer modules are set accordingly. + The key ``"default"`` is a special key that sets the quantizer attributes of all the quantizers for which + no other wildcard or filter functions match the quantizer module name. + + In addition, the dictionary entries could also be pytorch module class names mapping the class specific + quantization configuration. The pytorch modules should have a quantized equivalent. + + See :meth:`set_quantizer_attribute ` + for more details. + """ + quant_cfg = quant_cfg.copy() + if "default" in quant_cfg: + set_quantizer_attribute(quant_model, "*", quant_cfg["default"]) + quant_cfg.pop("default") + + for pattern, cfg in quant_cfg.items(): + if str(pattern) in QuantModuleRegistry: + parent_class = QuantModuleRegistry[str(pattern)] + assert isinstance( + cfg, dict + ), f"Expected a dictionary for quantizer configuration for child tensor quantizers of {parent_class}." + for sub_pattern, sub_cfg in cfg.items(): + set_quantizer_attribute(quant_model, sub_pattern, sub_cfg, parent_class) + continue + set_quantizer_attribute(quant_model, pattern, cfg) + + +def set_quantizer_attribute( + quant_model: nn.Module, + wildcard_or_filter_func: Union[str, Callable], + attribute: Union[ + QuantizerAttributeConfig, + list[QuantizerAttributeConfig], + dict[ + Union[str, Callable], + Union[QuantizerAttributeConfig, list[QuantizerAttributeConfig]], + ], + dict, + list[dict], + ], + parent_class: Union[None, type] = None, +): + """Finegrained adjustment of quantizer attribute by wildcard or filter function. + + Args: + quant_model: A pytorch model + wildcard_or_filter_func: a wildcard string or a filter function. The wildcard string is matched + against the quantizer module names. The quantizer modules are + instances of + :class:`TensorQuantizer `. + The filter function takes a quantized module name as input and returns ``True`` if the + quantizer should be adjusted and ``False`` otherwise. + attribute: An instance of :class:`QuantizerAttributeConfig <.config.QuantizerAttributeConfig>` or an equivalent + dictionary or a list of these two types. + If ``attribute`` is a list, the matched + :class:`TensorQuantizer ` + modules will be replaced with :class:`SequentialQuantizer ` + modules having one quantizer for each attribute instance from the list. + See + :meth:`set_from_attribute_config() ` + for more details on the supported attributes and their types. + parent_class: (Optional) The parent class of the quantizer modules matching ``wildcard_or_filter_func`` which + should be adjusted. If ``None``, all the matching quantizer modules will be adjusted. + """ + for name, module in quant_model.named_modules(): + if isinstance(module, (TensorQuantizer, SequentialQuantizer)): + if isinstance(wildcard_or_filter_func, str): + if not fnmatch.fnmatch(name, wildcard_or_filter_func): + continue + elif callable(wildcard_or_filter_func): + if not wildcard_or_filter_func(name): + continue + else: + raise NotImplementedError(f"Unsupported type {type(wildcard_or_filter_func)}") + + if parent_class is not None and not isinstance( + quant_model.get_submodule(".".join(name.split(".")[:-1])), parent_class + ): + continue + + if isinstance(attribute, list): + parent_module = quant_model.get_submodule(name.rpartition(".")[0]) + module = SequentialQuantizer(*(TensorQuantizer() for _ in range(len(attribute)))) + setattr(parent_module, name.split(".")[-1], module) + + module.set_from_attribute_config(attribute) + + +@contextmanager +def set_quantizer_by_cfg_context( + quant_model: nn.Module, quant_cfg: Union[QuantizeQuantCfgType, dict] +): + """Context manager for setting quantizer attributes using `quant_cfg`. + + The set attributes will be reset to the original attributes after exiting the context manager. + See :meth:`set_quantizer_by_cfg` for more details. + + Use this context manager with caution. Changing certain attributes of the quantizer such as + `calibrator` can lead to unexpected behavior. + """ + assert not any( + cfg for cfg in quant_cfg.values() if isinstance(cfg, (list, tuple)) + ), "list of config not support." + + original_attributes = {} + for name, module in quant_model.named_modules(): + if isinstance(module, TensorQuantizer): + original_attributes[name] = module.get_modelopt_state(properties_only=True) + + set_quantizer_by_cfg(quant_model, quant_cfg) + yield + for name, module in quant_model.named_modules(): + if isinstance(module, TensorQuantizer): + module.set_from_modelopt_state(original_attributes[name], name) + + +def register(original_cls: nn.Module, quantized_cls: nn.Module): + """Register a quantized class for the given un-quantized original class. + + Args: + original_cls: The original un-quantized class. + quantized_cls: The quantized class. This class should have a `_setup` method which initializes + various quantizers called in the forward. The forward function of the quantized class should call the + quantizers at the correct location. + + Here is an example of defining a quantized class and registering it: + + .. code-block:: python + + import modelopt.torch.quantization as mtq + from modelopt.torch.quantization.nn import TensorQuantizer + + + class QuantLayerNorm(nn.LayerNorm): + def __init__(self, normalized_shape): + super().__init__(normalized_shape) + self._setup() + + def _setup(self): + # Method to setup the quantizers + self.input_quantizer = TensorQuantizer() + self.weight_quantizer = TensorQuantizer() + + def forward(self, input): + input = self.input_quantizer(input) + weight = self.weight_quantizer(self.weight) + return F.layer_norm(input, self.normalized_shape, weight, self.bias, self.eps) + + + # Register the custom quantized module + mtq.register(original_cls=nn.LayerNorm, quantized_cls=QuantLayerNorm) + + """ + assert hasattr( + quantized_cls, "_setup" + ), "Quantized class must have a _setup method which initializes various quantizers." + + quantized_dm_cls = type( + quantized_cls.__name__, (quantized_cls, DynamicModule, original_cls), {} + ) + + QuantModuleRegistry.register({original_cls: original_cls.__name__})(quantized_dm_cls) + + +def unregister(original_cls: nn.Module): + """Unregister the quantized class for the given un-quantized original class. + + Args: + original_cls: The original un-quantized class. + + """ + QuantModuleRegistry.unregister(original_cls) + + +def export_quantized_model(model: nn.Module, config: _QuantizeExportConfig) -> ConvertReturnType: + """Export the quantized model to a quantized model.""" + raise NotImplementedError("Exporting a quantized model is not supported yet.") + + +def restore_export_quantized_model( + model: nn.Module, config: _QuantizeExportConfig, metadata: MetadataDict +) -> nn.Module: + """Restores the quantized model from the given state dict.""" + raise NotImplementedError("Restoring a quantized & exported model is not supported yet.") diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py new file mode 100644 index 0000000000..1dca9d9582 --- /dev/null +++ b/modelopt/torch/quantization/export_onnx.py @@ -0,0 +1,544 @@ +# Adapted from https://github.com/pytorch/pytorch/blob/e48ee2cf50d86d87ef7c7d0839267dbed4903ebf/torch/onnx/symbolic_opset14.py#L137-L207 + +# From PyTorch: + +# Copyright (c) 2016- Facebook, Inc (Adam Paszke) +# Copyright (c) 2014- Facebook, Inc (Soumith Chintala) +# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) +# Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) +# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) +# Copyright (c) 2011-2013 NYU (Clement Farabet) +# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) +# Copyright (c) 2006 Idiap Research Institute (Samy Bengio) +# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) + +# From Caffe2: + +# Copyright (c) 2016-present, Facebook Inc. All rights reserved. + +# All contributions by Facebook: +# Copyright (c) 2016 Facebook Inc. + +# All contributions by Google: +# Copyright (c) 2015 Google Inc. +# All rights reserved. + +# All contributions by Yangqing Jia: +# Copyright (c) 2015 Yangqing Jia +# All rights reserved. + +# All contributions by Kakao Brain: +# Copyright 2019-2020 Kakao Brain + +# All contributions by Cruise LLC: +# Copyright (c) 2022 Cruise LLC. +# All rights reserved. + +# All contributions by Tri Dao: +# Copyright (c) 2024 Tri Dao. +# All rights reserved. + +# All contributions by Arm: +# Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates + +# All contributions from Caffe: +# Copyright(c) 2013, 2014, 2015, the respective contributors +# All rights reserved. + +# All other contributions: +# Copyright(c) 2015, 2016 the respective contributors +# All rights reserved. + +# Caffe2 uses a copyright model similar to Caffe: each contributor holds +# copyright over their contributions to Caffe2. The project versioning records +# all such contribution and copyright details. If a contributor wants to further +# mark their specific copyright on a particular contribution, they should +# indicate their copyright solely in the commit message of the change when it is +# committed. + +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. + +# 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America +# and IDIAP Research Institute nor the names of its contributors may be +# used to endorse or promote products derived from this software without +# specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. +# +# Not a contribution +# Changes made by NVIDIA CORPORATION & AFFILIATES or otherwise documented as +# NVIDIA-proprietary are not a contribution and subject to the following terms and conditions: +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual +# property and proprietary rights in and to this material, related +# documentation and any modifications thereto. Any use, reproduction, +# disclosure or distribution of this material and related documentation +# without an express license agreement from NVIDIA CORPORATION or +# its affiliates is strictly prohibited. + +"""Utility to export a quantized torch model to quantized ONNX.""" + +from typing import Optional, Union + +import onnx +import torch +from torch.onnx import _type_utils, symbolic_helper +from torch.onnx import symbolic_helper as sym_help +from torch.onnx._internal import jit_utils +from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask + +onnx_dtype_map = { + "BFloat16": onnx.TensorProto.BFLOAT16, + "Float": onnx.TensorProto.FLOAT, + "Float8": onnx.TensorProto.FLOAT8E4M3FN, + "Half": onnx.TensorProto.FLOAT16, + "INT8": onnx.TensorProto.INT8, + "UINT8": onnx.TensorProto.UINT8, +} +mha_valid_precisions = {"Half", "BFloat16"} + +torch_dtype_map = {"Float": torch.float32, "Half": torch.float16, "BFloat16": torch.bfloat16} + + +def export_int8( + g: torch.onnx._internal.jit_utils.GraphContext, + inputs: torch.Value, + amax: torch.Tensor, + num_bits: int, + unsigned: bool, + narrow_range: bool, + trt_high_precision_dtype: str, +): + """Export quantized model to INT8 ONNX.""" + assert num_bits == 8, "Only INT8 ONNX export is supported for now." + output_shape = sym_help._get_tensor_sizes(inputs) + maxbound = (1 << (num_bits - 1 + int(unsigned))) - 1 + + if amax.numel() == 1: + zero_point, axis = torch.tensor(0.0, device=amax.device), None + else: + amax_init_shape = amax.shape + amax = amax.squeeze().data + assert len(amax.shape) == 1, "ONNX does not support multi-axis quantization." + zero_point = torch.zeros_like(amax, dtype=torch.int32).data + axis = list(amax_init_shape).index(list(amax.shape)[0]) + + zero_point = g.op("Constant", value_t=zero_point.to(torch_dtype_map[trt_high_precision_dtype])) + + if not unsigned: + assert not narrow_range, "ONNX does not support unsigned narrow range INT8." + zero_point = g.op("Cast", zero_point, to_i=onnx_dtype_map["INT8"]) + else: + zero_point = g.op("Cast", zero_point, to_i=onnx_dtype_map["UINT8"]) + + amax = amax.to(torch_dtype_map[trt_high_precision_dtype]) + scale = amax / maxbound + scale.masked_fill_(scale == 0, 1.0) + scale = g.op("Constant", value_t=scale) + + input_type = inputs.type().scalarType() + + assert ( + trt_high_precision_dtype == input_type or trt_high_precision_dtype == "Float" + ), "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." + + # custom ops, so cast the input if needed. + if trt_high_precision_dtype != input_type: + inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[trt_high_precision_dtype]) + quantized = g.op("QuantizeLinear", inputs, scale, zero_point, axis_i=axis) + out = g.op("DequantizeLinear", quantized, scale, zero_point, axis_i=axis).setType( + inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) + ) + + # custom ops, so cast the output if needed. + if trt_high_precision_dtype != input_type: + inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[input_type]) + + return out + + +def _fp8_quantize( + g: torch.onnx._internal.jit_utils.GraphContext, + inputs: torch.Value, + scale_inv: float, + trt_high_precision_dtype: str, +): + """Helper Function for Quantization.""" + output_shape = sym_help._get_tensor_sizes(inputs) + + # TRT StronglyType only supports FP16 QDQs + # custom ops, so cast the input if needed. + input_type = inputs.type().scalarType() + assert ( + trt_high_precision_dtype == input_type or trt_high_precision_dtype == "Float" + ), "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." + if trt_high_precision_dtype != input_type: + inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[trt_high_precision_dtype]) + + scale = g.op( + "Constant", + value_t=torch.tensor(scale_inv).to(torch_dtype_map[trt_high_precision_dtype]), + ) + q_op = g.op("trt::TRT_FP8QuantizeLinear", inputs, scale).setType( + inputs.type().with_dtype(torch.uint8).with_sizes(output_shape) + ) + return q_op + + +def _fp8_dequantize( + g: torch.onnx._internal.jit_utils.GraphContext, + inputs: torch.Value, + scale_inv: float, + trt_high_precision_dtype: str = "Float", + otype: str = None, +): + """Helper Function for Dequantization.""" + output_shape = sym_help._get_tensor_sizes(inputs) + assert ( + trt_high_precision_dtype == otype or trt_high_precision_dtype == "Float" + ), "TRT StronglyType requires both weights and amax to be in the BF16/FP16, or the QDQ in Float." + scale = g.op( + "Constant", + value_t=torch.tensor(scale_inv, dtype=torch_dtype_map[otype]), + ) + out = g.op("trt::TRT_FP8DequantizeLinear", inputs, scale).setType( + inputs.type().with_dtype(torch_dtype_map[trt_high_precision_dtype]).with_sizes(output_shape) + ) + + # DQ outputs are currently constrained to FP32 due to a similar limitation in ORT + # custom ops, so cast the output if needed. + if trt_high_precision_dtype != otype: + out = g.op("Cast", out, to_i=onnx_dtype_map[otype]) + return out + + +def export_fp8( + g: torch.onnx._internal.jit_utils.GraphContext, + inputs: torch.Value, + amax: float, + trt_high_precision_dtype: str, +): + """Export quantized model to FP8 ONNX.""" + if amax is None: + scale = 1.0 + else: + scale = 448.0 / float(amax) + otype = inputs.type().scalarType() + q_tensor = _fp8_quantize(g, inputs, 1.0 / scale, trt_high_precision_dtype) + return _fp8_dequantize(g, q_tensor, 1.0 / scale, trt_high_precision_dtype, otype) + + +def scaled_dot_product_attention( + g: jit_utils.GraphContext, + query: torch._C.Value, + key: torch._C.Value, + value: torch._C.Value, + attn_mask: Optional[torch._C.Value] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[torch._C.Value] = None, + enable_gqa: bool = False, +): + """Perform scaled dot product attention.""" + assert (not is_causal) or ( + is_causal and symbolic_helper._is_none(attn_mask) + ), "is_causal and attn_mask cannot be set at the same time" + assert ( + not enable_gqa + ), "conversion of scaled_dot_product_attention not implemented if enable_gqa is True" + + if symbolic_helper._is_none(scale): + scale = _attention_scale(g, query) + + if is_causal: + attn_mask = _causal_attention_mask(g, query, key) + + # Swap the last two axes of key + # NOTE: onnx-script has different logic here, because the attribute perms in + # transpose needs list of ints + key_shape_builtin = symbolic_helper._get_tensor_rank(key) + key_transposed_axes = list(range(key_shape_builtin)) + key_transposed_axes[-1], key_transposed_axes[-2] = ( + key_transposed_axes[-2], + key_transposed_axes[-1], + ) + key_transposed = g.op("Transpose", key, perm_i=key_transposed_axes) + + # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 + # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math + query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) + key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) + mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) + + if symbolic_helper._is_none(attn_mask): + mul_qk_add = mul_qk + elif _type_utils.JitScalarType.from_value(attn_mask) == _type_utils.JitScalarType.BOOL: + # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) + const_zero = g.op("Constant", value_t=torch.tensor([0.0])) + const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")])) + attn_mask = g.op("Where", attn_mask, const_zero, const_neg_inf) + mul_qk_add = g.op("Add", mul_qk, attn_mask) + elif _type_utils.JitScalarType.from_value(attn_mask) in ( + _type_utils.JitScalarType.FLOAT, + _type_utils.JitScalarType.HALF, + _type_utils.JitScalarType.BFLOAT16, + ): + mul_qk_add = g.op("Add", mul_qk, attn_mask) + else: + raise ValueError( + f"Unsupported type for attn_mask: {_type_utils.JitScalarType.from_value(attn_mask)}" + ) + + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + + if dropout_p != 0: + attn_weight = g.op( + "Dropout", + attn_weight, + g.op("Constant", value_t=torch.tensor(dropout_p, dtype=torch.float)), + ) + + return g.op("MatMul", attn_weight, value) + + +def export_fp8_mha( + g: torch.onnx._internal.jit_utils.GraphContext, + query: torch._C.Value, + key: torch._C.Value, + value: torch._C.Value, + attn_mask: Optional[torch._C.Value] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[torch._C.Value] = None, + q_quantized_scale: float = 1.0, + k_quantized_scale: float = 1.0, + v_quantized_scale: float = 1.0, + high_precision_flag: str = "Half", + disable_fp8_mha: bool = True, +): + r"""Export quantized fMHA to FP8 ONNX. + + FP8 ONNX graph:: + + Q K V + | | | + \ / | + QDQ QDQ | + \ / | + Cast Cast | + \ / | + BMM1 | + \ | + Cast QDQ + \ | + SoftMax | + | | + QDQ | + \ | + Cast Cast + \ / + BMM2 + | + Cast + """ + from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask + + # Pass all arguments, including x, to the custom ONNX operator + assert (not is_causal) or ( + is_causal and sym_help._is_none(attn_mask) + ), "is_causal and attn_mask cannot be set at the same time" + + scale = sym_help._maybe_get_const(scale, "f") + if sym_help._is_none(scale): + scale = _attention_scale(g, query) + + if is_causal: + attn_mask = _causal_attention_mask(g, query, key) + + # Swap the last two axes of key + # NOTE: onnx-script has different logic here, because the attribute perms in + # transpose needs list of ints + key_shape_builtin = sym_help._get_tensor_rank(key) + key_transposed_axes = list(range(key_shape_builtin)) + key_transposed_axes[-1], key_transposed_axes[-2] = ( + key_transposed_axes[-2], + key_transposed_axes[-1], + ) + key_transposed = g.op("Transpose", key, perm_i=key_transposed_axes) + + # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 + # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math + query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) + key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) + if not disable_fp8_mha: + if high_precision_flag not in mha_valid_precisions: + raise ValueError( + "The Quantized config setting doesn't match TRT's fusion pattern; the qdqs must be in 16 bits." + ) + q_input_dtype = query.type().scalarType() + k_input_dtype = key.type().scalarType() + v_input_dtype = value.type().scalarType() + if {q_input_dtype, k_input_dtype, v_input_dtype} != {high_precision_flag}: + raise ValueError("The quantized MHA must have 16-bit inputs.") + query_scaled = export_fp8(g, query_scaled, q_quantized_scale, high_precision_flag) + query_scaled = g.op("Cast", query_scaled, to_i=onnx_dtype_map["Float"]) + key_transposed_scaled = export_fp8( + g, key_transposed_scaled, k_quantized_scale, high_precision_flag + ) + key_transposed_scaled = g.op("Cast", key_transposed_scaled, to_i=onnx_dtype_map["Float"]) + mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) + if not disable_fp8_mha: + mul_qk = g.op("Cast", mul_qk, to_i=onnx_dtype_map[high_precision_flag]) + + if sym_help._is_none(attn_mask): + mul_qk_add = mul_qk + elif _type_utils.JitScalarType.from_value(attn_mask) == _type_utils.JitScalarType.BOOL: + # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) + const_zero = g.op("Constant", value_t=torch.tensor([0.0])) + const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")])) + attn_mask = g.op("Where", attn_mask, const_zero, const_neg_inf) + mul_qk_add = g.op("Add", mul_qk, attn_mask) + elif _type_utils.JitScalarType.from_value(attn_mask) in ( + _type_utils.JitScalarType.FLOAT, + _type_utils.JitScalarType.HALF, + _type_utils.JitScalarType.BFLOAT16, + ): + mul_qk_add = g.op("Add", mul_qk, attn_mask) + else: + raise ValueError( + f"Unsupported type for attn_mask: {_type_utils.JitScalarType.from_value(attn_mask)}" + ) + + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + + if not disable_fp8_mha: + # Softmax's output scale is hard coded to 1.0 + attn_weight = export_fp8(g, attn_weight, 1.0, high_precision_flag) + attn_weight = g.op("Cast", attn_weight, to_i=onnx_dtype_map["Float"]) + + if dropout_p != 0: + attn_weight = g.op( + "Dropout", + attn_weight, + g.op("Constant", value_t=torch.tensor(dropout_p, dtype=torch.float)), + ) + if not disable_fp8_mha: + value = export_fp8(g, value, v_quantized_scale, high_precision_flag) + value = g.op("Cast", value, to_i=onnx_dtype_map["Float"]) + return g.op( + "Cast", + g.op("MatMul", attn_weight, value), + to_i=onnx_dtype_map[high_precision_flag], + ) + else: + return g.op("MatMul", attn_weight, value) + + +def _fp4_dynamic_quantize( + g: torch.onnx._internal.jit_utils.GraphContext, + inputs: torch.Value, + scale: float, + trt_high_precision_dtype: str, + block_size: int, + axis: int = -1, + scale_type: int = onnx_dtype_map["Float8"], +): + """Helper Function for Dynamic Quantization.""" + # TRT StronglyType only supports FP16 QDQ ops, so cast the input if needed. + input_type = inputs.type().scalarType() + if trt_high_precision_dtype != input_type: + inputs = g.op("Cast", inputs, to_i=onnx_dtype_map[trt_high_precision_dtype]) + + scale = g.op( + "Constant", + value_t=torch.tensor(scale).to(torch_dtype_map["Float"]), + ) + # This is a TensorRT local function, it dynamically quantizes the input tensor to FP4. + xf4, sx_f8 = g.op( + "trt::TRT_FP4DynamicQuantize", + inputs, + scale, + axis_i=axis, + block_size_i=block_size, + scale_type_i=scale_type, + outputs=2, + ) + + return xf4, sx_f8 + + +def _fp4_dequantize( + g: torch.onnx._internal.jit_utils.GraphContext, + inputs: torch.Value, + scale: Union[float, torch.Value], + trt_high_precision_dtype: str, +): + """Helper Function for Dequantization.""" + if isinstance(scale, float): + scale = g.op( + "Constant", + value_t=torch.tensor(scale, dtype=torch_dtype_map["Float"]), + ) + return g.op("trt::DequantizeLinear", inputs, scale) + + +def _fp4_dequantize_2( + g: torch.onnx._internal.jit_utils.GraphContext, + inputs: torch.Value, + dyn_scale: torch.Value, + block_size: int, + axis: int = -1, +): + """Helper Function for Dequantization.""" + return g.op("trt::DequantizeLinear", inputs, dyn_scale, axis_i=axis, block_size_i=block_size) + + +def export_fp4( + g: torch.onnx._internal.jit_utils.GraphContext, + inputs: torch.Value, + block_size: int, + amax: torch.Value, + num_bits: tuple[int, int], + trt_high_precision_dtype: str, + onnx_quantizer_type: str, +): + """Export quantized model to FP4 ONNX.""" + if onnx_quantizer_type == "dynamic": + amax = sym_help._get_const(amax, "f", "amax") + if amax is None or amax == 0: + sx_f32_per_tensor = 1.0 + else: + assert num_bits == (2, 1) + sx_f32_per_tensor = float(amax) / 6.0 + + x_f4, sx_f8 = _fp4_dynamic_quantize( + g, inputs, sx_f32_per_tensor, trt_high_precision_dtype, block_size + ) + dq_scale = _fp4_dequantize(g, sx_f8, sx_f32_per_tensor, trt_high_precision_dtype) + return _fp4_dequantize_2(g, x_f4, dq_scale, block_size) + else: + # This is a dummy custom op, we post-process the exported ONNX model to replace this node + # with two DQ nodes as described in the static double quantization recipe. + return g.op("trt::TRT_FP4QDQ", inputs, block_size_i=block_size) diff --git a/modelopt/torch/quantization/extensions.py b/modelopt/torch/quantization/extensions.py new file mode 100644 index 0000000000..185dded428 --- /dev/null +++ b/modelopt/torch/quantization/extensions.py @@ -0,0 +1,89 @@ +# 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. + +"""Module to load C++ / CUDA extensions.""" + +from pathlib import Path + +from modelopt.torch.utils import load_cpp_extension + +__all__ = ["get_cuda_ext", "get_cuda_ext_fp8", "get_cuda_ext_mx", "precompile"] + +path = Path(__file__).parent + + +def get_cuda_ext(raise_if_failed: bool = False): + """Returns the cuda extention for tensor_quant.""" + if not hasattr(get_cuda_ext, "extension"): + get_cuda_ext.extension = load_cpp_extension( # type:ignore[attr-defined] + name="modelopt_cuda_ext", + sources=[path / "src/tensor_quant.cpp", path / "src/tensor_quant_gpu.cu"], + cuda_version_specifiers=">=11", + raise_if_failed=raise_if_failed, + ) + return get_cuda_ext.extension # type:ignore[attr-defined] + + +def get_cuda_ext_fp8(raise_if_failed: bool = False): + """Returns the cuda extention for tensor_quant_fp8.""" + if not hasattr(get_cuda_ext_fp8, "extension"): + get_cuda_ext_fp8.extension = load_cpp_extension( # type:ignore[attr-defined] + name="modelopt_cuda_ext_fp8", + sources=[path / "src/tensor_quant_fp8.cpp", path / "src/tensor_quant_gpu_fp8.cu"], + cuda_version_specifiers=">=11.8", + fail_msg=( + "CUDA extension for FP8 quantization could not be built and loaded, FP8 simulated" + " quantization will not be available." + ), + raise_if_failed=raise_if_failed, + ) + return get_cuda_ext_fp8.extension # type:ignore[attr-defined] + + +def get_cuda_ext_mx(raise_if_failed: bool = False): + """Returns the cuda extention for tensor_quant_mx.""" + if not hasattr(get_cuda_ext_mx, "extension"): + get_cuda_ext_mx.extension = load_cpp_extension( # type:ignore[attr-defined] + name="modelopt_cuda_ext_mx", + sources=[ + path / "src/tensor_quant_mx.cu", + ], + cuda_version_specifiers=">=11.8", + fail_msg=( + "CUDA extension for MX quantization could not be built and loaded, MX simulated" + " quantization will not be available." + ), + extra_cuda_cflags=["--use_fast_math"], + raise_if_failed=raise_if_failed, + ) + return get_cuda_ext_mx.extension # type:ignore[attr-defined] + + +def __getattr__(name): + if name == "cuda_ext": + return get_cuda_ext() + elif name == "cuda_ext_fp8": + return get_cuda_ext_fp8() + elif name == "cuda_ext_mx": + return get_cuda_ext_mx() + else: + raise AttributeError(f"module {__name__} has no attribute {name}") + + +def precompile(): + """Precompile the CUDA extensions.""" + print(get_cuda_ext()) + print(get_cuda_ext_fp8()) + print(get_cuda_ext_mx()) diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py new file mode 100644 index 0000000000..59eeb6b650 --- /dev/null +++ b/modelopt/torch/quantization/mode.py @@ -0,0 +1,122 @@ +# 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. + +"""This module contains the mode descriptor for the quantization mode.""" + +from typing import Optional + +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.mode import ( + ConvertEntrypoint, + RestoreEntrypoint, + UpdateEntrypoint, + _ModeDescriptor, + _ModeRegistryCls, +) + +from .config import QuantizeConfig, _QuantizeExportConfig +from .conversion import ( + convert_to_quantized_model, + export_quantized_model, + restore_export_quantized_model, + restore_quantized_model, + update_quantize_metadata, +) + +QuantizeModeRegistry = _ModeRegistryCls() + + +# TODO: OMNIML-717 Reuse search infra for quantization calibration algorithms +@QuantizeModeRegistry.register_mode +class QuantizeModeDescriptor(_ModeDescriptor): + """Class to describe the ``"quant"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "quantize" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return QuantizeConfig + + @property + def next_modes(self) -> Optional[set[str]]: + """Modes that must immediately follow this mode.""" + return {"kd_loss"} + + @property + def export_mode(self) -> Optional[str]: + """The mode that corresponds to the export mode of this mode.""" + return "export_quantize" + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return convert_to_quantized_model + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_quantized_model + + @property + def update_for_save(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models state before saving.""" + return update_quantize_metadata + + @property + def update_for_new_mode(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models state before new mode.""" + return update_quantize_metadata + + +@QuantizeModeRegistry.register_mode +class QuantizeExportModeDescriptor(_ModeDescriptor): + """Class to describe the export of quantization mode. + + Note that this mode is just a placeholder to throw an error since we don't support exporting + quantized models right now. It is used to properly indicate that the ``quantize`` mode does + require an export mode if we ever wanted to do chaining/stacking of modes with it. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "quantize_export" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return _QuantizeExportConfig + + @property + def is_export_mode(self) -> bool: + """Specifies whether the mode is an export mode.""" + return True + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return export_quantized_model + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_export_quantized_model diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py new file mode 100644 index 0000000000..38beabf0d9 --- /dev/null +++ b/modelopt/torch/quantization/model_calib.py @@ -0,0 +1,912 @@ +# 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. + +"""Calibration utilities.""" + +import math +import types +import warnings +from functools import partial +from typing import Optional + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F + +from modelopt.torch.opt.searcher import ForwardLoop +from modelopt.torch.utils.distributed import ParallelState + +from .conversion import set_quantizer_by_cfg_context +from .nn import SequentialQuantizer, TensorQuantizer +from .qtensor import QTensorWrapper, pack_real_quantize_weight +from .utils import ( + get_parallel_state, + is_quantized_column_parallel_linear, + is_quantized_layer_with_weight, + is_quantized_linear, + is_quantized_row_parallel_linear, +) + +__all__ = ["max_calibrate", "real_quantize", "awq", "smoothquant"] + + +@torch.no_grad() +def max_calibrate(model: nn.Module, forward_loop: Optional[ForwardLoop] = None): + """Calibrate the model using max.""" + enable_stats_collection(model) + if forward_loop is None: + # Lets do a weight only calibration + def forward_loop(model: nn.Module): + seen_modules = set() + for name, module in model.named_modules(): + if module in seen_modules: + continue + if is_quantized_layer_with_weight(module) and hasattr(module, "weight_quantizer"): + module.weight_quantizer(module.weight) + seen_modules.add(module) + + forward_loop(model) + finish_stats_collection(model) + + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer) and module.amax is not None: + module.sync_amax_across_distributed_group( + get_parallel_state(model, name).data_parallel_group + ) + + # TP sync: + # Objective: the quantization parameters when TP = 8 then changed to TP=4 then back to TP=8 should be the same + + # ColumnParallel: X @ [A_1, A_2] (weights split along Cout) + # activations: TPG should have the same amax if axis in [None, -1] + # weights: TPG should have the same amax if axis in [None, -1] (note: we dont use -1 axis for weights) + + # RowParallel: [X_1, X_2] @ [A_1 + # A_2] (weights split along Cin) + # activations: TPG should have the same amax if axis in [None] + # weights: TPG should have the same amax if axis in [None, 0] + + def sync_quantizer_amax_across_tp( + quantizer: TensorQuantizer, + linear_name: str, + quantizer_type: str, + axes_for_sync: list, + parallel_state: ParallelState, + ): + # sync is not needed for block quantization + if quantizer.block_sizes is not None: + if hasattr(quantizer, "_padding"): + warnings.warn( + f"Found block-quantized padded {quantizer_type} for {linear_name}, amax will" + " not be synced correctly." + ) + # Skip amax sync for INT4 / W4A8 block quantization + # Sync amax for NVFP4 (dynamic per-block, static per-tensor quantized scale) + if not getattr(quantizer.block_sizes, "type", None) != "dynamic": + return + + if quantizer.axis in axes_for_sync and quantizer.amax is not None: + quantizer.sync_amax_across_distributed_group(parallel_state.tensor_parallel_group) + + for name, module in model.named_modules(): + if is_quantized_column_parallel_linear(module): + for input_quantizer in SequentialQuantizer.tensor_quantizer_iterator( + module.input_quantizer + ): + sync_quantizer_amax_across_tp( + input_quantizer, + name, + "input_quantizer", + axes_for_sync=[None, -1], + parallel_state=get_parallel_state(module), + ) + + for weight_quantizer in SequentialQuantizer.tensor_quantizer_iterator( + module.weight_quantizer + ): + sync_quantizer_amax_across_tp( + weight_quantizer, + name, + "weight_quantizer", + axes_for_sync=[None, -1], + parallel_state=get_parallel_state(module), + ) + + if is_quantized_row_parallel_linear(module): + for input_quantizer in SequentialQuantizer.tensor_quantizer_iterator( + module.input_quantizer + ): + sync_quantizer_amax_across_tp( + input_quantizer, + name, + "input_quantizer", + axes_for_sync=[None], + parallel_state=get_parallel_state(module), + ) + + for weight_quantizer in SequentialQuantizer.tensor_quantizer_iterator( + module.weight_quantizer + ): + sync_quantizer_amax_across_tp( + weight_quantizer, + name, + "weight_quantizer", + axes_for_sync=[None, 0], + parallel_state=get_parallel_state(module), + ) + + +def enable_stats_collection(model: nn.Module): + """Enable stats collection for all quantizers in the model.""" + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer) and not module._disabled: + if module._calibrator is not None: + module.disable_quant() + module.enable_calib() + else: + module.disable() + + +def finish_stats_collection(model: nn.Module, method: Optional[str] = None): + """Finish stats collection for all quantizers in the model.""" + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer) and not module._disabled: + if module._calibrator is not None and not module._dynamic: + if method in ["mse", "entropy"]: + if module._calibrator.compute_amax(method) is not None: + module.load_calib_amax(method) + elif module._calibrator.compute_amax() is not None: + module.load_calib_amax() + + module.enable_quant() + module.disable_calib() + + +def real_quantize( + model: nn.Module, + forward_loop: Optional[ForwardLoop] = None, + **kwargs, +): + """Calibrate the model if given additonal_algorithm and then apply real quantization.""" + # calibrate the model before real quantization + # Import here to avoid circular import + + additional_algorithm = ( + kwargs["additional_algorithm"] if "additional_algorithm" in kwargs else None + ) + if additional_algorithm and "method" in additional_algorithm: + calib_algo = additional_algorithm["method"] + with set_quantizer_by_cfg_context(model, {"*weight_quantizer": {"fake_quant": True}}): + if "awq" in calib_algo: + awq(model, calib_algo, forward_loop, **kwargs) + + # Apply real quantize to the model if needed + real_quantize_weight(model) + + +def real_quantize_weight(model): + """Applies real weight-only quantization to the given model using the provided configuration. + + This function first applies a quantization mode to the model using the provided configuration. + It then recursively traverses the model to apply real quantization to the weights of any modules + that have a `weight_quantizer` attribute enabled. The quantized weights are wrapped in a + `QTensorWrapper` to make them compatible with PyTorch's `torch.nn.Parameter`. + + Args: + module (nn.Module): The neural network model to be quantized. + + Note: + This function modifies the input model in-place. + """ + # Real quantize the weights. If the weight quantizer is a `SequentialQuantizer`, + # we only need to real-quantize the first weight quantizer. + with SequentialQuantizer.replace_sequential_quantizer_with_single_quantizer(model): + pack_real_quantize_weight(model) + + # TODO: remove this warnning once the real quantization feature is mature + def _has_qtensorwrapper(module): + if hasattr(module, "weight") and isinstance(module.weight, QTensorWrapper): + return True + for _, submodule in module.named_children(): + if _has_qtensorwrapper(submodule): + return True + return False + + if _has_qtensorwrapper(model): + warnings.warn( + "Real quantization has been applied to the model. This feature is still " + "experimental, and some functionalities may not be supported. For example, " + "converting the model back to its original state or saving and restoring " + "the quantized model may not be available." + ) + + return model + + +@torch.no_grad() +def disable_pre_quant_scale_and_resmooth(linear: nn.Module, delete_pre_quant_scale: bool = False): + """Disable pre_quant_scale and resmooth the quantized linear weights.""" + assert is_quantized_linear(linear), "Only quantized linear modules are supported" + assert ( + linear.input_quantizer._enable_pre_quant_scale + ), "pre_quant_scale should be enabled first!" + assert hasattr( + linear.input_quantizer, "_pre_quant_scale" + ), "pre_quant_scale should be available" + + pre_quant_scale = linear.input_quantizer._pre_quant_scale.to(torch.float32) + + linear.weight.copy_( + (linear.weight * pre_quant_scale.squeeze()[None, :]).to(linear.weight.dtype) + ) + linear.weight_quantizer.reset_amax() + max_calibrate(linear, lambda linear: linear.weight_quantizer(linear.weight)) + + # Lets not delete the _pre_quant_scale, it might useful later; Instead we will disable it + linear.input_quantizer._enable_pre_quant_scale = False + + if linear.input_quantizer.amax is not None: + assert hasattr(linear.input_quantizer, "_amax_for_smoothing") + linear.input_quantizer.amax = linear.input_quantizer._amax_for_smoothing.amax().to( + linear.weight.dtype + ) + + if delete_pre_quant_scale: + delattr(linear.input_quantizer, "_pre_quant_scale") + linear.input_quantizer._enable_pre_quant_scale = False + + +# A global variable used during AutoQuantize to avoid folding pre_quant_scale to weights +_ENABLE_FOLDING_PQS_TO_WEIGHTS = True + + +@torch.no_grad() +def _apply_weight_pre_quant_scale(linear, pre_quant_scale): + if _ENABLE_FOLDING_PQS_TO_WEIGHTS: + linear.weight.data.copy_( + (linear.weight * pre_quant_scale.squeeze()[None, :]).to(linear.weight.dtype) + ) + else: + linear.weight_quantizer._enable_pre_quant_scale = True + linear.weight_quantizer.pre_quant_scale = pre_quant_scale.squeeze()[None, :].to( + linear.weight.dtype + ) + + linear.weight_quantizer.reset_amax() + max_calibrate(linear, lambda linear: linear.weight_quantizer(linear.weight)) + + +@torch.no_grad() +def apply_pre_quant_scale_and_smooth( + linear: nn.Module, pre_quant_scale: Optional[torch.Tensor] = None +): + """Apply pre_quant_scale and smooth the quantized linear weights. + + If pre_quant_scale is not provided, the existing pre_quant_scale of input_quantizer will be used. + """ + assert is_quantized_linear(linear), "Only quantized linear modules are supported" + assert linear.input_quantizer.pre_quant_scale is None, "pre_quant_scale should be None first!" + + if pre_quant_scale is None: + pre_quant_scale = ( + linear.input_quantizer._pre_quant_scale + if hasattr(linear.input_quantizer, "_pre_quant_scale") + else None + ) + + assert pre_quant_scale is not None, "pre_quant_scale should be provided or already set" + + assert torch.all(pre_quant_scale > 0), "pre_quant_scale should be positive" + + # pre_quant_scale should be in fp32 for the scaling math to be numerically safe + pre_quant_scale = pre_quant_scale.to(torch.float32) + + linear.input_quantizer._enable_pre_quant_scale = True + linear.input_quantizer.pre_quant_scale = pre_quant_scale.to(linear.weight.dtype) + + inv_scale = 1.0 / pre_quant_scale + _apply_weight_pre_quant_scale(linear, inv_scale) + + if linear.input_quantizer.amax is not None: + assert hasattr(linear.input_quantizer, "_amax_for_smoothing") + linear.input_quantizer.amax = ( + (linear.input_quantizer._amax_for_smoothing * pre_quant_scale) + .amax() + .to(linear.weight.dtype) + ) + + if is_quantized_column_parallel_linear(linear) or is_quantized_row_parallel_linear(linear): + linear.input_quantizer.sync_amax_across_distributed_group( + get_parallel_state(linear).tensor_parallel_group + ) + + +@torch.no_grad() +def smoothquant(model: nn.Module, forward_loop: Optional[ForwardLoop] = None, alpha=1.0): + """Smooth-Quant variant with per-channel weight scaling. + + The parameters are as described in + :class:`SmoothQuantCalibConfig `. + """ + # distributed synchornization + # max_calibrate performs amax sync for data parallel + + # Column parallel: + # activations: TPG should have the same pre_quant_scale + # This is achieved by syncing act_amax and weight_scale across TPG which is used to + # compute pre_quant_scale + # weights: no-op + + # Row parallel: + # activations: TPG should have same activation amax + # weights: TPG should have the same weight amax + + assert forward_loop is not None, "forward_loop must be provided for smoothquant" + for name, module in model.named_modules(): + if ( + is_quantized_linear(module) + and module.input_quantizer.is_enabled + and module.input_quantizer.axis is None + ): + module.input_quantizer.axis = -1 + + max_calibrate(model, forward_loop) + + smoothed_modules = 0 + for name, module in model.named_modules(): + if is_quantized_linear(module): + if not hasattr(module.input_quantizer, "_amax"): + print(f"Warning: {name} is not calibrated, skip smoothing") + continue + if module.input_quantizer.num_bits != 8 or module.weight_quantizer.num_bits != 8: + print(f"Warning: only int8 smoothing is supported, skip {name}") + continue + if module.input_quantizer.axis != -1: + print(f"Warning: only per-channel smoothing is supported, skip {name}") + continue + + assert ( + module.input_quantizer._amax.numel() > 1 + ), f"Error: {name} has only one channel to smooth" + + # It is important to keep scaling math in fp32 to be numerically safe + act_amax = module.input_quantizer.amax.float() + weight_scale = module.weight.abs().amax(dim=0, keepdim=True) + + parallel_group = get_parallel_state(module).tensor_parallel_group + if is_quantized_column_parallel_linear(module) and parallel_group.is_initialized(): + dist.all_reduce(act_amax, op=dist.ReduceOp.MAX, group=parallel_group.group) + dist.all_reduce(weight_scale, op=dist.ReduceOp.MAX, group=parallel_group.group) + + scale_a = (weight_scale.pow(1 - alpha) / act_amax.pow(alpha)).squeeze() + + # Now that activation per-channel amax have been collected, use per-tensor quantization for activation + module.input_quantizer._amax_for_smoothing = act_amax + module.input_quantizer.reset_amax() + module.input_quantizer.axis = None + module.input_quantizer.amax = act_amax.amax() + + # Some channel could have 0 amax which causes scale_a to overflow. Explicitly mask them out here + epsilon = 1.0 / (1 << 31) + if scale_a.min() <= epsilon: + zero_mask = act_amax <= epsilon + scale_a[zero_mask] = 1 + scale_a = scale_a.clamp(min=1e-4, max=1e4) + apply_pre_quant_scale_and_smooth(module, scale_a) + + smoothed_modules += 1 + print(f"Smoothed {smoothed_modules} modules") + + +def _smoothquant_fasteval(model: nn.Module): + """Hacky implementation of Smooth-Quant. Copied from monkey-quant.""" + smoothed_modules = 0 + for name, module in model.named_modules(): + if is_quantized_linear(module): + if not hasattr(module.input_quantizer, "_amax"): + print(f"Warning: {name} is not calibrated, skip smoothing") + continue + if module.input_quantizer.num_bits != 8 or module.weight_quantizer.num_bits != 8: + print(f"Warning: only int8 smoothing is supported, skip {name}") + continue + if module.input_quantizer.axis != -1: + print(f"Warning: only per-channel smoothing is supported, skip {name}") + continue + + assert module.input_quantizer._amax.numel() > 1 + delattr(module.weight_quantizer, "_amax") + + # It is important to keep scaling math in fp32 to be numerically safe + act_amax = module.input_quantizer.amax.float() + if act_amax.shape[0] == 1: + act_amax = act_amax.squeeze(0) + # If model is split across devices, this tensor may be on wrong one + act_amax = act_amax.to(module.weight.device) + + max_bound = module.input_quantizer.maxbound + scale_a = max_bound / act_amax + # Some channel could have 0 amax which causes scale_a to overflow. Explicitly mask them out here + epsilon = 1.0 / (1 << 31) + if act_amax.min() <= epsilon: + zero_mask = act_amax <= epsilon + scale_a[zero_mask] = 1 + inv_scale_a = act_amax / max_bound + + module.weight.data.copy_( + (module.weight_quantizer(inv_scale_a * module.weight.float()) * scale_a).to( + module.weight.dtype + ) + ) + module.weight_quantizer.disable() + + smoothed_modules += 1 + print(f"Smoothed {smoothed_modules} modules") + + +def awq( + model: nn.Module, + algorithm: str = "awq_lite", + forward_loop: Optional[ForwardLoop] = None, + **kwargs, +): + """Apply AWQ to the model.""" + with SequentialQuantizer.replace_sequential_quantizer_with_single_quantizer(model): + if algorithm in ["awq_full", "awq_lite"]: + awq_lite(model, forward_loop, **kwargs) + + if algorithm in ["awq_full", "awq_clip"]: + awq_clip(model, forward_loop, **kwargs) + + for name, module in model.named_modules(): + if is_quantized_linear(module) and isinstance(module.weight_quantizer, SequentialQuantizer): + max_calibrate(module, lambda linear: linear.weight_quantizer(module.weight)) + # If there are any uncalibrated quantizers, calibrate them using max + if any( + not module._disabled and module.amax is None + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) + ): + from .model_quant import calibrate + + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer) and hasattr(module, "_amax"): + # Already calibrated modules should not be calibrated again + # This is true for the first weight quantizer of the sequential weight quantizers + assert module.is_enabled, "calibrated modules should be enabled" + module.disable() + + calibrate(model, algorithm="max", forward_loop=forward_loop) + + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer) and hasattr(module, "_amax"): + module.enable() + + +@torch.no_grad() +def awq_lite( + model: nn.Module, + forward_loop: ForwardLoop, + alpha_step: float = 0.1, + debug: bool = False, + **kwargs, +): + """Lite version of AWQ. + + The parameters are as described in + :class:`AWQLiteCalibConfig `. + + Args: + model: Model to be calibrated. + forward_loop: A callable which takes the model as argument and + forwards calibration data through the model. + alpha_step: Step size for the searching awq_lite hyper parameter alpha. alpha is in [0, 1]. + debug: If True, module's search metadata will be kept as a module attribute named `awq_lite`. + + """ + assert forward_loop is not None, "forward_loop must be provided for awq_lite" + + class AWQLiteHelper: + cache_mode: bool = False + + def __init__(self, module): + self.act_scale = 0.0 + self.num_cache_steps = 0 + self.num_search_steps = 0 + self.block_size = _get_awq_quantizer_block_size(module.weight, module.weight_quantizer) + self.weight_scale = get_weight_scale(module.weight, self.block_size) + self.loss = {k.item(): 0.0 for k in torch.arange(0, 1.0 + alpha_step, alpha_step)} + self.best_scale = None + self.best_alpha = None + self.is_input_quantized = module.input_quantizer.is_enabled + + def get_weight_scale(weight, block_size=None): + org_shape = weight.shape + slice_after_padding = None + if block_size: + if org_shape[-1] % block_size != 0: + slice_after_padding = slice(org_shape[-1]) + weight = F.pad(weight, (0, block_size - org_shape[-1] % block_size), "constant", 0) + org_shape = weight.shape + weight = weight.contiguous().view(-1, block_size) + weight_abs_amax = weight.abs().amax(dim=1, keepdim=True) + scale = weight.abs() / (weight_abs_amax + torch.finfo(weight.dtype).tiny) + scale = scale.view(org_shape) + if slice_after_padding is not None: + scale = scale[..., slice_after_padding] + scale = scale.mean(0).to(torch.float32) + return scale + + def get_act_scale(x): + return x.abs().contiguous().view(-1, x.shape[-1]).mean(0).to(torch.float32) + + def get_scale(x_max, w_max, alpha, tensor_parallel_group=None, dtype=torch.float32): + scales = ( + (x_max.pow(alpha) / (w_max.pow(1 - alpha) + torch.finfo(w_max.dtype).tiny)) + .clamp(min=1e-4, max=1e4) + .view(-1) + ) + scales = (scales / (scales.max() * scales.min()).sqrt()).view(-1) + if tensor_parallel_group and tensor_parallel_group.is_initialized(): + dist.all_reduce(scales, op=dist.ReduceOp.SUM, group=tensor_parallel_group.group) + scales /= tensor_parallel_group.world_size() + return scales.to(dtype) + + def update_loss(self, out, out_actual, alpha): + out_actual = out_actual[0] if isinstance(out_actual, tuple) else out_actual + out = out[0] if isinstance(out, tuple) else out + loss = (out - out_actual).float().pow(2).mean().item() + self.awq_lite.loss[alpha] += loss + + def update_best_params(self): + self.awq_lite.best_alpha = min(self.awq_lite.loss, key=self.awq_lite.loss.get) + self.awq_lite.best_scale = get_scale( + self.awq_lite.act_scale, + self.awq_lite.weight_scale, + self.awq_lite.best_alpha, + ( + get_parallel_state(self).tensor_parallel_group + if is_quantized_column_parallel_linear(self) + else None + ), + self.weight.dtype, + ) + + def forward(self, input, *args, **kwargs): + # Collect actual output without quantization + self.weight_quantizer.disable() + if hasattr(self.input_quantizer, "_pre_quant_scale"): + delattr(self.input_quantizer, "_pre_quant_scale") + if hasattr(self.weight_quantizer, "_pre_quant_scale"): + delattr(self.weight_quantizer, "_pre_quant_scale") + out_actual = self._forward_no_awq(input, *args, **kwargs) + self.weight_quantizer.enable() + + if AWQLiteHelper.cache_mode: + # NOTE: Input can be on different device than weights + self.awq_lite.act_scale += get_act_scale(input).to(self.awq_lite.weight_scale.device) + self.awq_lite.num_cache_steps += 1 + if self.awq_lite.is_input_quantized: + with set_quantizer_by_cfg_context(self.input_quantizer, {"*": {"enable": True}}): + # TODO: Why? Calling max_calibrate on the linear layer instead of input_quantizer + # is causing memory leakage due to the garbage collection not working as expected + # max_calibrate(self, lambda linear: linear.input_quantizer(input)) + max_calibrate(self.input_quantizer, lambda quantizer: quantizer(input)) + return out_actual + + for alpha in self.awq_lite.loss.keys(): + awq_scale = get_scale( + self.awq_lite.act_scale, + self.awq_lite.weight_scale, + alpha, + ( + get_parallel_state(self).tensor_parallel_group + if is_quantized_column_parallel_linear(self) + else None + ), + self.weight.dtype, + ) + self.input_quantizer.pre_quant_scale = 1 / awq_scale + self.weight_quantizer.pre_quant_scale = awq_scale + out = self._forward_no_awq(input, *args, **kwargs) + + update_loss(self, out, out_actual, alpha) + + self.awq_lite.num_search_steps += 1 + + # Now forward the actual output without any quantization + return out_actual + + for name, module in model.named_modules(): + if is_quantized_linear(module) and module.weight_quantizer.is_enabled: + module._forward_no_awq = module.forward + module.awq_lite = AWQLiteHelper(module) + module.forward = types.MethodType(forward, module) + + if module.input_quantizer.is_enabled: + module.input_quantizer.disable() + if module.input_quantizer.axis not in [None, -1]: + raise NotImplementedError( + "input quantization needs to be per-tensor or None for AWQ algorithm" + ) + module.input_quantizer.axis = -1 + + # Collect activation scale values + AWQLiteHelper.cache_mode = True + print("Caching activation statistics for awq_lite...") + forward_loop(model) + + # If input_quantizers were enabled they would be calibrated in caching mode + # Call max_calibrate to sync amax across distributed groups + max_calibrate(model, lambda model: None) + + for name, module in model.named_modules(): + if ( + is_quantized_linear(module) + and hasattr(module, "awq_lite") + and module.awq_lite.num_cache_steps > 0 + ): + module.awq_lite.act_scale /= module.awq_lite.num_cache_steps + + AWQLiteHelper.cache_mode = False + print("Searching awq_lite parameters...") + forward_loop(model) + + for name, module in model.named_modules(): + if hasattr(module, "awq_lite"): + if module.awq_lite.num_cache_steps > 0: + assert module.awq_lite.num_search_steps > 0, ( + "Calling `forward_loop(model)` the second time did not forward data through the" + " model. Please provide a valid `forward_loop` function that can be used to" + " forward data through the model many times." + ) + + update_best_params(module) + delattr(module.weight_quantizer, "_pre_quant_scale") + delattr(module.input_quantizer, "_pre_quant_scale") + if module.awq_lite.is_input_quantized: + assert module.input_quantizer.amax is not None + module.input_quantizer.to(module.weight.device) + module.input_quantizer._amax_for_smoothing = module.input_quantizer.amax + module.input_quantizer.reset_amax() + module.input_quantizer.axis = None + module.input_quantizer.amax = module.input_quantizer._amax_for_smoothing.amax() + module.input_quantizer.enable() + + apply_pre_quant_scale_and_smooth(module, 1.0 / module.awq_lite.best_scale) + + if not debug: + delattr(module, "awq_lite") + + # Restore forward function; Why is this dangerous? + module.forward = module._forward_no_awq + delattr(module, "_forward_no_awq") + + +@torch.no_grad() +def awq_clip( + model: nn.Module, + forward_loop: ForwardLoop, + max_co_batch_size: int = 1024, + max_tokens_per_batch: int = 64, + min_clip_ratio: float = 0.5, + shrink_step: float = 0.05, + debug: bool = False, + **kwargs, +): + """AWQ-Clip variant. + + The parameters are as described in + :class:`AWQClipCalibConfig `. + + Args: + model: Model to calibrate. + forward_loop: A callable that runs the forward pass of the model. + max_co_batch_size: Maximum output channel batch size while searching clip values. Reduce this number if CUDA Out + of Memory error occurs. + max_tokens_per_batch: Maximum tokens per batch while searching clip values. The total tokens used for + clip search would be max_tokens_per_batch * number of batches. Original AWQ uses a total of 512 tokens to + search for clip values. + min_clip_ratio: Minimum clip ratio to search for. It should be in (0, 1.0). Clip will search for the optimal + clipping value in the range [original block amax * min_clip_ratio, original block amax]. + shrink_step: Step size to search for clip values. It should be in range (0, 1.0]. + debug: If True, module's search metadata will be kept as a module attribute named `awq_clip`. + """ + assert forward_loop is not None, "forward_loop must be provided for awq_clip" + + class AWQClipHelper: + def __init__(self, module): + self.num_tokens = 0 + self.block_size = _get_awq_quantizer_block_size(module.weight, module.weight_quantizer) + + # Cache the original amax + module.weight_quantizer.reset_amax() + enable_stats_collection(module.weight_quantizer) + module.weight_quantizer(module.weight) + finish_stats_collection(module.weight_quantizer) + self.w_amax = module.weight_quantizer.amax.clone() + + co, ci = module.weight.shape + clip_ratios = [ + round(float(k), 2) for k in torch.arange(min_clip_ratio, 1.0, shrink_step) + ] + [1.0] + if self.is_per_tensor_clip(module): + self.loss = {k: torch.tensor(0.0, device=module.weight.device) for k in clip_ratios} + else: + self.loss = { + k: torch.zeros( + (co, math.ceil(ci / self.block_size)), device=module.weight.device + ) + for k in clip_ratios + } + self.best_clip_val = None + self.best_loss = None + + self.is_input_quantized = module.input_quantizer.is_enabled + + def is_per_tensor_clip(self, module): + quantizer = module.weight_quantizer + is_dynamic_w_per_tensor = ( + hasattr(quantizer, "block_sizes") + and quantizer.block_sizes.get("type", None) == "dynamic" + and quantizer.axis is None + ) + is_per_tensor = quantizer.axis is None and quantizer.block_sizes is None + return is_dynamic_w_per_tensor or is_per_tensor + + def update_best_params(self): + self.awq_clip.best_loss = torch.ones_like(self.awq_clip.w_amax) * float("inf") + self.awq_clip.best_clip_val = torch.zeros_like(self.awq_clip.w_amax) + + for shrink, loss in self.awq_clip.loss.items(): + loss = loss.view_as(self.awq_clip.w_amax) + indices = loss < self.awq_clip.best_loss + self.awq_clip.best_loss = torch.where(indices, loss, self.awq_clip.best_loss) + self.awq_clip.best_clip_val = torch.where( + indices, self.awq_clip.w_amax * shrink, self.awq_clip.best_clip_val + ) + + def _clip_search(self, inputs, co_bsz=256, max_tokens=16): + weight = self.weight + self.weight_quantizer.enable() + + if self.awq_clip.is_per_tensor_clip(self): + # In NVFP4, only the per-tensor amax is clipped + out_actual = inputs @ self.weight.T + original_amax = self.weight_quantizer.amax.clone() + self.awq_clip.num_tokens += inputs.shape[0] + for shrink in self.awq_clip.loss.keys(): + self.weight_quantizer.amax = original_amax * shrink + out = inputs @ self.weight_quantizer(self.weight).T + loss = (out - out_actual).float().pow(2).mean() + self.awq_clip.loss[shrink] += loss + else: + # weight [co, ci] -> [co, 1, n_block, block_size] + # inputs [..., ci] -> [1, max_tokens, n_block, block_size] + + inputs = inputs.view(-1, inputs.shape[-1]) # _, ci + # Select max_tokens from the total input tokens of count batch * n_token + inputs = inputs[0 :: max(1, inputs.shape[0] // max_tokens)] # max_tokens, ci + self.awq_clip.num_tokens += inputs.shape[0] + + block_size = self.awq_clip.block_size + co, ci = weight.shape + if ci % block_size != 0: + weight = F.pad(weight, (0, block_size - ci % block_size), "constant", 0) + inputs = F.pad(inputs, (0, block_size - ci % block_size), "constant", 0) + ci = weight.shape[-1] + + weight = weight.reshape(co, 1, -1, block_size) # co, 1, n_block, block_size + + # 1, max_tokens, n_block, block_size + inputs = inputs.reshape(1, inputs.shape[0], -1, block_size) + + for co_batch in range(math.ceil(co / co_bsz)): + w = weight[co_batch * co_bsz : min((co_batch + 1) * co_bsz, co)] + + org_out = (inputs * w).sum(dim=-1) # co_bsz, max_tokens, n_block + + for shrink in self.awq_clip.loss.keys(): + self.weight_quantizer.amax = self.awq_clip.w_amax * shrink + quantized_clipped_weight = self.weight_quantizer(self.weight) + cur_w = quantized_clipped_weight[ + co_batch * co_bsz : min((co_batch + 1) * co_bsz, co) + ] + if cur_w.shape[-1] % block_size != 0: + cur_w = F.pad( + cur_w, (0, block_size - cur_w.shape[-1] % block_size), "constant", 0 + ) + cur_w = cur_w.reshape(w.shape) + cur_out = (inputs * cur_w).sum(dim=-1) # co_bsz, max_tokens, n_block + + # co_bsz, n_block + loss = (cur_out - org_out).float().pow(2).mean(dim=1) + + parallel_group = get_parallel_state(self).data_parallel_group + if parallel_group.is_initialized(): + dist.all_reduce(loss, op=dist.ReduceOp.SUM, group=parallel_group.group) + loss /= parallel_group.world_size() + + del cur_out, cur_w + self.awq_clip.loss[shrink][ + co_batch * co_bsz : min((co_batch + 1) * co_bsz, co) + ] += loss + del org_out + + def forward(name, self, input, *args, **kwargs): + # input shape : (..., cin) + # weight shape : (cout, cin) + # Disable quantization + self.weight_quantizer.disable() + if self.awq_clip.is_input_quantized: + self.input_quantizer.enable() + max_calibrate(self.input_quantizer, lambda input_quantizer: input_quantizer(input)) + self.input_quantizer.disable() + try: + _clip_search(self, self.input_quantizer(input), max_co_batch_size, max_tokens_per_batch) + except RuntimeError as e: + if "CUDA out of memory" in str(e): + raise RuntimeError( + f"Clip search on {name} failed due to CUDA out of memory, try reducing" + " max_co_batch_size" + ) from e + raise RuntimeError(e) + + # Disable quantization + self.weight_quantizer.disable() + return self._forward_no_awq(input, *args, **kwargs) + + for name, module in model.named_modules(): + if ( + is_quantized_linear(module) + and module.weight_quantizer.is_enabled + and module.weight_quantizer.block_sizes is not None + ): + module._forward_no_awq = module.forward + module.forward = types.MethodType(partial(forward, name), module) + module.awq_clip = AWQClipHelper(module) + + print("Estimating awq_clip parameters...") + forward_loop(model) + + for name, module in model.named_modules(): + if is_quantized_linear(module) and hasattr(module, "awq_clip"): + if module.awq_clip.num_tokens > 0: + update_best_params(module) + + # Load the best clip value (amax) + module.weight_quantizer.amax = module.awq_clip.best_clip_val + module.weight_quantizer.enable() + if module.awq_clip.is_input_quantized: + module.input_quantizer.enable() + + if not debug: + delattr(module, "awq_clip") + + # Restore forward function + module.forward = module._forward_no_awq + delattr(module, "_forward_no_awq") + + +def _get_awq_quantizer_block_size(tensor: torch.Tensor, quantizer: TensorQuantizer): + if quantizer.block_sizes is None: + return None + if -1 in quantizer.block_sizes: + blocksize = quantizer.block_sizes[-1] + elif 1 in quantizer.block_sizes: + blocksize = quantizer.block_sizes[1] + else: + raise ValueError("AWQ requires block quantization along -1 axis") + return blocksize diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py new file mode 100644 index 0000000000..a5a7f63050 --- /dev/null +++ b/modelopt/torch/quantization/model_quant.py @@ -0,0 +1,448 @@ +# 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. + +"""User-facing quantization API.""" + +import fnmatch +import inspect +import warnings +from typing import Any, Callable, Iterable, Optional, Union + +import torch +import torch.nn as nn + +from modelopt.core.torch.quantization.algorithms import AutoQuantizeSearcher +from modelopt.torch.opt import apply_mode +from modelopt.torch.opt.searcher import ForwardLoop + +from .config import CalibAlgorithmCfgType +from .conversion import set_quantizer_attribute +from .mode import QuantizeModeRegistry +from .model_calib import awq, max_calibrate, real_quantize, smoothquant +from .nn import TensorQuantizer + +__all__ = [ + "calibrate", + "postprocess_amax", + "quantize", + "auto_quantize", + "disable_quantizer", + "enable_quantizer", + "print_quant_summary", + "fold_weight", +] + + +# TODO: Descriptors for the supported algorithms +def calibrate( + model: nn.Module, + algorithm: Union[str, CalibAlgorithmCfgType, None] = "max", + forward_loop: Optional[ForwardLoop] = None, +) -> nn.Module: + """Adjusts weights and scaling factors based on selected algorithms. + + Args: + model: A pytorch model with quantizer modules. + algorithm: A string or dictionary specifying the calibration algorithm to use. Supported + algorithms are ``"max"``, ``"smoothquant"``, ``"awq_lite"``, ``"awq_full"``, and + ``"awq_clip"``. If a dictionary is passed, the key ``"method"`` should specify the + calibration algorithm to use. Other key-value pairs in this dictionary will be passed + as kwargs to the algorithm. An example dictionary argument: + ``{"method": "awq_clip", "max_co_batch_size": 4096}``. If ``None``, no calibration is + performed. For real quantization, the key ``method`` should be ``real_quantize``, and + the calibration algorithm used should be specified in ``additional_algorithm``. + forward_loop: A callable which takes the model as argument and forwards calibration data + through the model. This is not required for weight-only quantization with the ``"max"`` + algorithm. + + Returns: The calibrated pytorch model. + """ + if forward_loop is not None: + # get the number of arguments of forward_loop + num_args = len(inspect.signature(forward_loop).parameters) + if num_args == 0: + warnings.warn( + ( + "forward_loop should take model as argument, but got forward_loop without any" + " arguments. This usage will be deprecated in future versions." + ), + DeprecationWarning, + ) + original_forward_loop = forward_loop + + def forward_loop(model): + return original_forward_loop() + + if algorithm is None: + return model + + algorithm = algorithm if isinstance(algorithm, (str, dict)) else algorithm.model_dump() + + if isinstance(algorithm, str): + kwargs = {} + elif isinstance(algorithm, dict): + kwargs = algorithm.copy() + algorithm = kwargs.pop("method") + else: + raise TypeError(f"Unsupported type for algorithm: {type(algorithm)}") + + # move the model to eval mode + is_training = model.training + model.eval() + + if algorithm.startswith("awq"): # type: ignore[union-attr] + awq(model, algorithm, forward_loop, **kwargs) # type: ignore[arg-type] + elif algorithm == "smoothquant": + smoothquant(model, forward_loop, **kwargs) + elif algorithm == "max": + max_calibrate(model, forward_loop) + elif algorithm == "real_quantize": + real_quantize(model, forward_loop, **kwargs) + else: + raise ValueError(f"Unsupported calibration algorithm: {algorithm}") + + # TODO: Re-enable when the CUDA error: unspecified launch failure is fixed. + # clear_cuda_cache() + + model.train(is_training) + + return model + + +def postprocess_amax(model: nn.Module, key: str, post_process_fn) -> nn.Module: + """Experimental API to postprocess the amax values after calibration.""" + assert isinstance(key, str), "key should be a string" + for name, module in model.named_modules(): + if not isinstance(module, TensorQuantizer): + continue + if not hasattr(module, "_amax"): + continue + if not fnmatch.fnmatch(name, key): + continue + module.amax = post_process_fn(module.amax) + + return model + + +def quantize( + model: nn.Module, + config: dict[str, Any], + forward_loop: Optional[ForwardLoop] = None, +) -> nn.Module: + """Quantizes and calibrates the model in-place. + + This method performs replacement of modules with their quantized counterparts and + performs calibration as specified by ``quant_cfg``. + ``forward_loop`` is used to forward data through the model and gather statistics for calibration. + + Args: + model: A pytorch model + config: A dictionary or an instance of + :class:`QuantizeConfig ` specifying the + values for keys ``"quant_cfg"`` and ``"algorithm"``. + It is basically a dictionary specifying the values for keys ``"quant_cfg"`` and ``"algorithm"``. + The ``"quant_cfg"`` key specifies the quantization configurations. + The ``"algorithm"`` key specifies the ``algorithm`` argument to + :meth:`calibrate `. + + Quantization configurations is a dictionary mapping wildcards or filter functions + to its quantizer attributes. The wildcards or filter functions are matched + against the quantizer module names. The quantizer modules have names ending with + ``weight_quantizer`` and ``input_quantizer`` and they perform weight quantization and + input quantization (or activation quantization) respectively. The quantizer modules + are instances of + :class:`TensorQuantizer `. + The quantizer attributes are defined by :class:`QuantizerAttributeConfig`. See + :class:`QuantizerAttributeConfig` for details on the quantizer attributes and their values. + + An example ``config`` dictionary is given below: + + .. code-block::python + + config = { + + "quant_cfg": { + # "num_bits" specifies the number of bits for quantization + # "axis" specifies the axis for quantization + "*weight_quantizer": {"num_bits": 8, "axis": 0}, + "*input_quantizer": {"num_bits": 8, "axis": -1}, + + # Default quantization settings + "default": {"num_bits": 8, "axis": None}, + } + "algorithm": "max" + } + + See :ref:`Quantization Formats ` to learn more about the supported + quantization formats. See :ref:`Quantization Configs ` for more details on + ``config`` dictionary. + + forward_loop: A callable that forwards all calibration data through the model. This is used + to gather statistics for calibration. It should take model as the argument. It does not need + to return anything. + + This argument is not required for weight-only quantization with the ``"max"`` + algorithm. + + Here are a few examples for correct ``forward_loop`` definitions: + Example 1: + + .. code-block:: + + def forward_loop(model) -> None: + # iterate over the data loader and forward data through the model + for batch in data_loader: + model(batch) + + Example 2: + + .. code-block:: + + def forward_loop(model) -> float: + # evaluate the model on the task + return evaluate(model, task, ....) + + Example 3: + + .. code-block:: + + def forward_loop(model) -> None: + # run evaluation pipeline + evaluator.model = model + evaluator.evaluate() + + .. note:: + + Calibration does not require forwarding the entire dataset through the model. + Please subsample the dataset or reduce the number of batches if needed. + + Returns: A pytorch model which has been quantized and calibrated. + """ + model = apply_mode(model, mode=[("quantize", config)], registry=QuantizeModeRegistry) + return calibrate(model, config["algorithm"], forward_loop=forward_loop) + + +def auto_quantize( + model: nn.Module, + constraints: dict[str, Union[float, str]] = {"effective_bits": 4.8}, + quantization_formats: list[Optional[str]] = ["NVFP4_DEFAULT_CFG", "FP8_DEFAULT_CFG", None], + data_loader: Iterable = None, + forward_step: Callable[[nn.Module, Any], Union[Any, torch.Tensor]] = None, + loss_func: Callable[[Any, Any], torch.Tensor] = None, + forward_backward_step: Optional[Callable[[nn.Module, Any], Any]] = None, + disabled_layers: Optional[Union[list[str], str]] = None, + num_calib_steps: int = 512, + num_score_steps: int = 128, + verbose: bool = False, +): + r"""API for ``AutoQuantize`` which quantizes a model by searching for the best quantization formats per-layer. + + ``auto_quantize`` uses a gradient based sensitivity score to rank the per-layer quantization formats and search + for the best quantization formats per-layer. + + Args: + model: A pytorch model with quantizer modules. + constraints: Constraints for the search. Currently we support only ``effective_bits``. + ``effective_bits`` specifies the effective number of bits for the quantized model. + + Here is an example for valid ``effective_bits`` argument: + + .. code-block:: python + + # For an effective quantization bits of 4.8 + constraints = {"effective_bits": 4.8} + + quantization_formats: A list of the string names of the quantization formats to search for. + The supported quantization formats are as listed by :attr:`modelopt.torch.quantization.config.choices`. + + In addition, the quantization format can also be ``None`` which implies skipping quantization for + the layer. + + .. note:: + + The quantization formats will be applied on a per-layer match basis. The global model level name + based quantizer attribute setting will be ignored. For example, in ``FP8_DEFAULT_CFG`` quantizer + configuration the key ``"*lm_head*": {"enable": False}`` disables quantization for the ``lm_head`` + layer. However in ``auto_quantize``, the quantization format for the ``lm_head`` layer will be searched. + This is because the key ``"*lm_head*"`` sets the quantizer attributes based on the global model level + name, not per-layer basis. The keys ``"*input_quantizer"``, ``"*weight_quantizer"`` etc. in + ``FP8_DEFAULT_CFG`` match on a per-layer basis - hence the corresponding quantizers + will be set as specified. + + Here is an example `quantization_formats` argument: + + .. code-block:: python + + # A valid `quantization_formats` argument + # This will search for the best per-layer quantization from FP8, W4A8_AWQ or No quantization + quantization_formats = ["FP8_DEFAULT_CFG", "W4A8_AWQ", None] + + data_loader: An iterator that yields data that is to be used for calibrating quantized layers and estimating + ``auto_quantize`` scores. + forward_step: A callable that takes the model and a batch of data from ``data_loader`` as input, forwards + the data through the model and returns the model output. + This is a required argument. + + Here is an example for a valid ``forward_step``: + + .. code-block:: python + + # Takes the model and a batch of data as input and returns the model output + def forward_step(model, batch) -> torch.Tensor: + output = model(batch) + return output + + loss_func: (Optional) A callable that takes the model output and the batch of data as input and computes the + loss. The model output is the output given by ``forward_step``. `.backward()` will be called on the loss. + + Here is an example for a valid ``loss_func``: + + .. code-block:: python + + # Takes the model output and a batch of data as input and returns the loss + def loss_func(output, batch) -> torch.Tensor: + ... + return loss + + + # loss should be a scalar tensor such that loss.backward() can be called + loss = loss_func(output, batch) + loss.backward() + + If this argument is not provided, ``forward_backward_step`` should be provided. + + forward_backward_step: (Optional) A callable that takes batch of data from ``data_loader``, forwards it + through the model, computes the loss and runs backward on the loss. + + Here is an example for a valid ``forward_backward_step`` argument: + + .. code-block:: python + + # Takes the model and a batch of data as input and runs forward and backward pass + def forward_backward_step(model, batch) -> None: + output = model(batch) + loss = my_loss_func(output, batch) + run_custom_backward(loss) + + If this argument is not provided, ``loss_func`` should be provided. + + disabled_layers: (Optional) One or a list of wildcard strings to disable quantization for the layers. Example: + + .. code-block:: python + + disabled_layers = "*lm_head*" + disabled_layers = ["*lm_head*", "*mlp*"] + + num_calib_steps: Number of batches to use for calibrating the quantized model. Suggested value is 512. + num_score_steps: Number of batches to use for estimating ``auto_quantize`` scores. Suggested value is 128. + A higher value could increase the time taken for performing ``auto_quantize``. + verbose: If True, prints the search progress/intermediate results. + + Returns: A tuple (model, state_dict) where ``model`` is the searched and quantized model and + ``state_dict`` contains the history and detailed stats of the search procedure. + + .. note:: + + ``auto_quantize`` groups certain layers and restricts the quantization formats for them to be same. For example, + Q, K, V linear layers belonging to the same transformer layer will have the same quantization format. + This is to ensure compatibility with TensorRT-LLM which fuses these three linear layers into a single linear + layer. + + A list of regex pattern rules as defined in :attr:`rules <.algorithms.AutoQuantizeSearcher.rules>` + are used to specify the group of layers. The first captured group + in the regex pattern (i.e, ``pattern.match(name).group(1)``) is used to group the layers. All the layers + that share the same first captured group will have the same quantization format.. + + For example, the rule ``r"^(.*?)\.(q_proj|k_proj|v_proj)$"`` + groups the `q_proj`, `k_proj`, `v_proj` linear layers belonging to the same transformer layer. + + You may modify the rules to group the layers as per your requirement. + + .. code-block:: python + + from modelopt.torch.quantization.algorithms import AutoQuantizeSearcher + + # To additionally group the layers belonging to same `mlp` layer, + # add the following rule + AutoQuantizeSearcher.rules.append(r"^(.*?)\.mlp") + + # Perform `auto_quantize` + model, state_dict = auto_quantize(model, ...) + + .. note:: + + The ``auto_quantize`` API and algorithm is experimental and subject to change. ``auto_quantize`` searched models + might not be readily deployable to TensorRT-LLM yet. + + """ + model = apply_mode( + model, + mode=[("quantize", {"quant_cfg": {}})], + registry=QuantizeModeRegistry, + ) + searcher = AutoQuantizeSearcher() + search_config = { + "quantization_formats": quantization_formats, + "data_loader": data_loader, + "forward_step": forward_step, + "loss_func": loss_func, + "forward_backward_step": forward_backward_step, + "num_calib_steps": num_calib_steps, + "num_score_steps": num_score_steps, + "verbose": verbose, + } + searcher.search(model, constraints, config=search_config) # type: ignore[arg-type] + + if disabled_layers: + if isinstance(disabled_layers, str): + disabled_layers = [disabled_layers] + for layer_pattern in disabled_layers: + disable_quantizer(model, layer_pattern) + + return model, searcher.state_dict() + + +def disable_quantizer(model: nn.Module, wildcard_or_filter_func: Union[str, Callable]): + """Disable quantizer by wildcard or filter function.""" + set_quantizer_attribute(model, wildcard_or_filter_func, {"enable": False}) + + +def enable_quantizer(model: nn.Module, wildcard_or_filter_func: Union[str, Callable]): + """Enable quantizer by wildcard or filter function.""" + set_quantizer_attribute(model, wildcard_or_filter_func, {"enable": True}) + + +def print_quant_summary(model: nn.Module): + """Print summary of all quantizer modules in the model.""" + count = 0 + for name, mod in model.named_modules(): + if isinstance(mod, TensorQuantizer): + print(f"{name:80} {mod}") + count += 1 + print(f"{count} TensorQuantizers found in model") + + +def fold_weight(model: nn.Module): + """Fold weight quantizer for fast evaluation.""" + for name, module in model.named_modules(): + if hasattr(module, "weight_quantizer") and hasattr(module, "weight"): + module.weight.data.copy_( + (module.weight_quantizer(module.weight.float())).to(module.weight.dtype) + ) + module.weight_quantizer.disable() + if hasattr(module.weight_quantizer, "_pre_quant_scale"): + delattr(module.weight_quantizer, "_pre_quant_scale") + if hasattr(module.weight_quantizer, "_amax"): + delattr(module.weight_quantizer, "_amax") diff --git a/modelopt/torch/quantization/nn/__init__.py b/modelopt/torch/quantization/nn/__init__.py new file mode 100644 index 0000000000..ae2ba5cd78 --- /dev/null +++ b/modelopt/torch/quantization/nn/__init__.py @@ -0,0 +1,27 @@ +# 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. + +"""Modules with quantization support.""" + +from .modules.clip import * +from .modules.quant_activations import * +from .modules.quant_batchnorm import * +from .modules.quant_conv import * +from .modules.quant_instancenorm import * +from .modules.quant_linear import * +from .modules.quant_module import * +from .modules.quant_pooling import * +from .modules.quant_rnn import * +from .modules.tensor_quantizer import * diff --git a/modelopt/torch/quantization/nn/functional.py b/modelopt/torch/quantization/nn/functional.py new file mode 100644 index 0000000000..7f086a60f3 --- /dev/null +++ b/modelopt/torch/quantization/nn/functional.py @@ -0,0 +1,72 @@ +# 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. + +"""Some supportive functions.""" + +import warnings + +import torch +from torch.autograd import Function + + +class ClipFunction(Function): + """An universal tensor clip function. + + Pytorch's clamp() only supports scalar range and doesn't support broadcast. This implementation uses min/max which + is more genaral. The gradient is defined according to IBM's PACT paper https://arxiv.org/abs/1805.06085, which is + also the behavior of Tensorflow's clip_by_value() + """ + + @staticmethod + def forward(ctx, input, clip_value_min, clip_value_max): + """Forward pass for the clip function.""" + output = torch.min(input, clip_value_max) + output = torch.max(output, clip_value_min) + ctx.save_for_backward(input, clip_value_min, clip_value_max) + return output + + @staticmethod + def backward(ctx, grad_output): + """Backward pass for the clip function.""" + input, clip_value_min, clip_value_max = ctx.saved_tensors + min_mask = (input > clip_value_min).to(grad_output.dtype) + max_mask = (input < clip_value_max).to(grad_output.dtype) + grad_input = grad_output * min_mask * max_mask + + if clip_value_min.requires_grad or clip_value_max.requires_grad: + warnings.warn("Learning enabled for clip min/max. This is an experimental feature.") + if clip_value_min.numel() != 1 or clip_value_max.numel() != 1: + raise ValueError( + "Learnable min/max can only be scalar, got size %s and %s." + % (clip_value_min.size(), clip_value_max.size()) + ) + + # Ensure the dtypes of min/max grads matches the input dtype + # This might be necessary if running w/ AMP which will cast to fp32 before `sum()` + grad_clip_value_min = ( + (grad_output * (1.0 - min_mask)).sum().to(clip_value_min.dtype) + if clip_value_min.requires_grad + else None + ) + grad_clip_value_max = ( + (grad_output * (1.0 - max_mask)).sum().to(clip_value_min.dtype) + if clip_value_max.requires_grad + else None + ) + + return grad_input, grad_clip_value_min, grad_clip_value_max + + +clip = ClipFunction.apply diff --git a/modelopt/torch/quantization/nn/modules/__init__.py b/modelopt/torch/quantization/nn/modules/__init__.py new file mode 100644 index 0000000000..52b069c58b --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/__init__.py @@ -0,0 +1,16 @@ +# 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. + +"""Modules with quantization support.""" diff --git a/modelopt/torch/quantization/nn/modules/clip.py b/modelopt/torch/quantization/nn/modules/clip.py new file mode 100644 index 0000000000..a5bf3057d1 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/clip.py @@ -0,0 +1,56 @@ +# 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. + +"""Implement a clip module as pytorch only has a simple clamp function.""" + +import torch +from torch import nn +from torch.nn.parameter import Parameter + +from .. import functional as QF + +__all__ = ["Clip"] + + +class Clip(nn.Module): + """Clip tensor. + + Args: + clip_value_min: A number or tensor of lower bound to clip + clip_value_max: A number of tensor of upper bound to clip + learn_min: A boolean. If True, learn min. clip_value_min will be used to initialize. Default False + learn_max: A boolean. Similar as learn_min but for max. + + Raises: + ValueError: + """ + + def __init__(self, clip_value_min, clip_value_max, learn_min=False, learn_max=False): + """Initialize.""" + super(Clip, self).__init__() + if learn_min: + self.clip_value_min = Parameter(torch.tensor(clip_value_min)) + else: + self.clip_value_min = clip_value_min + + if learn_max: + self.clip_value_max = Parameter(torch.tensor(clip_value_max)) + else: + self.clip_value_max = clip_value_max + + def forward(self, inputs): + """Clip input tensor.""" + outputs = QF.clip(inputs, self.clip_value_min, self.clip_value_max) + return outputs diff --git a/modelopt/torch/quantization/nn/modules/quant_activations.py b/modelopt/torch/quantization/nn/modules/quant_activations.py new file mode 100644 index 0000000000..0448f94b34 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_activations.py @@ -0,0 +1,22 @@ +# 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. + +"""Quantized activations module.""" + +import torch.nn as nn + +from .quant_module import QuantInputBase, QuantModuleRegistry + +QuantModuleRegistry.register({nn.LeakyReLU: "nn.LeakyReLU"})(QuantInputBase) diff --git a/modelopt/torch/quantization/nn/modules/quant_batchnorm.py b/modelopt/torch/quantization/nn/modules/quant_batchnorm.py new file mode 100644 index 0000000000..21eed5b82e --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_batchnorm.py @@ -0,0 +1,24 @@ +# 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. + +"""Quantized batch normalization module.""" + +import torch.nn as nn + +from .quant_module import QuantInputBase, QuantModuleRegistry + +QuantModuleRegistry.register({nn.BatchNorm1d: "nn.BatchNorm1d"})(QuantInputBase) +QuantModuleRegistry.register({nn.BatchNorm2d: "nn.BatchNorm2d"})(QuantInputBase) +QuantModuleRegistry.register({nn.BatchNorm3d: "nn.BatchNorm3d"})(QuantInputBase) diff --git a/modelopt/torch/quantization/nn/modules/quant_conv.py b/modelopt/torch/quantization/nn/modules/quant_conv.py new file mode 100644 index 0000000000..fedd3e8d36 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_conv.py @@ -0,0 +1,123 @@ +# 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. + +"""Quantized convolution.""" + +import torch.nn as nn + +from ... import tensor_quant +from .quant_module import QuantLinearConvBase, QuantModuleRegistry, _LegacyQuantLinearConvBaseMixin + +__all__ = [ + "Conv2d", + "QuantConv2d", + "Conv3d", + "QuantConv3d", + "Conv1d", + "QuantConv1d", + "ConvTranspose1d", + "ConvTranspose2d", + "ConvTranspose3d", + "QuantConvTranspose1d", + "QuantConvTranspose2d", + "QuantConvTranspose3d", +] + + +@QuantModuleRegistry.register({nn.Conv1d: "nn.Conv1d"}) +class _QuantConv1d(QuantLinearConvBase): + """Quantized 1D convolution.""" + + default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL + + +class QuantConv1d(_LegacyQuantLinearConvBaseMixin, nn.Conv1d): + """Quantized 1D convolution.""" + + default_quant_desc_weight = _QuantConv1d.default_quant_desc_weight + + +@QuantModuleRegistry.register({nn.Conv2d: "nn.Conv2d"}) +class _QuantConv2d(QuantLinearConvBase): + """Quantized 2D convolution.""" + + default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL + + +class QuantConv2d(_LegacyQuantLinearConvBaseMixin, nn.Conv2d): + """Quantized 2D convolution.""" + + default_quant_desc_weight = _QuantConv2d.default_quant_desc_weight + + +@QuantModuleRegistry.register({nn.Conv3d: "nn.Conv3d"}) +class _QuantConv3d(QuantLinearConvBase): + """Quantized 3D convolution.""" + + default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL + + +class QuantConv3d(_LegacyQuantLinearConvBaseMixin, nn.Conv3d): + """Quantized 3D convolution.""" + + default_quant_desc_weight = _QuantConv3d.default_quant_desc_weight + + +@QuantModuleRegistry.register({nn.ConvTranspose1d: "nn.ConvTranspose1d"}) +class _QuantConvTranspose1d(QuantLinearConvBase): + """Quantized 1D transposed convolution.""" + + default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL + + +class QuantConvTranspose1d(_LegacyQuantLinearConvBaseMixin, nn.ConvTranspose1d): + """Quantized 1D transposed convolution.""" + + default_quant_desc_weight = _QuantConvTranspose1d.default_quant_desc_weight + + +@QuantModuleRegistry.register({nn.ConvTranspose2d: "nn.ConvTranspose2d"}) +class _QuantConvTranspose2d(QuantLinearConvBase): + """Quantized 2D transposed convolution.""" + + default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL + + +class QuantConvTranspose2d(_LegacyQuantLinearConvBaseMixin, nn.ConvTranspose2d): + """Quantized 2D transposed convolution.""" + + default_quant_desc_weight = _QuantConvTranspose2d.default_quant_desc_weight + + +@QuantModuleRegistry.register({nn.ConvTranspose3d: "nn.ConvTranspose3d"}) +class _QuantConvTranspose3d(QuantLinearConvBase): + """Quantized 3D transposed convolution.""" + + default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL + + +class QuantConvTranspose3d(_LegacyQuantLinearConvBaseMixin, nn.ConvTranspose3d): + """Quantized 3D transposed convolution.""" + + default_quant_desc_weight = _QuantConvTranspose3d.default_quant_desc_weight + + +# Define alias with Quant prefix +Conv1d = QuantConv1d +Conv2d = QuantConv2d +Conv3d = QuantConv3d +ConvTranspose1d = QuantConvTranspose1d +ConvTranspose2d = QuantConvTranspose2d +ConvTranspose3d = QuantConvTranspose3d diff --git a/modelopt/torch/quantization/nn/modules/quant_instancenorm.py b/modelopt/torch/quantization/nn/modules/quant_instancenorm.py new file mode 100644 index 0000000000..e5647c2ea4 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_instancenorm.py @@ -0,0 +1,45 @@ +# 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. + +"""Quantized instance normalization module.""" + +import torch.nn as nn + +from .quant_module import QuantInputBase, QuantModuleRegistry, _LegacyQuantInputBaseMixin + +__all__ = ["QuantInstanceNorm1d", "QuantInstanceNorm2d", "QuantInstanceNorm3d"] + + +class QuantInstanceNorm1d(_LegacyQuantInputBaseMixin, nn.InstanceNorm1d): + """Applies Quantized Instance Normalization over a 3D input.""" + + pass + + +class QuantInstanceNorm2d(_LegacyQuantInputBaseMixin, nn.InstanceNorm2d): + """Applies Quantized Instance Normalization over a 4D input.""" + + pass + + +class QuantInstanceNorm3d(_LegacyQuantInputBaseMixin, nn.InstanceNorm3d): + """Applies Quantized Instance Normalization over a 5D input.""" + + pass + + +QuantModuleRegistry.register({nn.InstanceNorm1d: "nn.InstanceNorm1d"})(QuantInputBase) +QuantModuleRegistry.register({nn.InstanceNorm2d: "nn.InstanceNorm2d"})(QuantInputBase) +QuantModuleRegistry.register({nn.InstanceNorm3d: "nn.InstanceNorm3d"})(QuantInputBase) diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py new file mode 100644 index 0000000000..74e627d8b1 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_linear.py @@ -0,0 +1,50 @@ +# 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. + +"""Quantized Linear.""" + +import torch.nn as nn + +from ... import tensor_quant +from .quant_module import QuantLinearConvBase, QuantModuleRegistry, _LegacyQuantLinearConvBaseMixin + +__all__ = ["Linear", "QuantLinear"] + + +@QuantModuleRegistry.register({nn.Linear: "nn.Linear"}) +class _QuantLinear(QuantLinearConvBase): + """Quantized base class for nn.Linear type classes.""" + + default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW + + @staticmethod + def quantized_linear_fn(package, func_name, self, input, weight, *args, **kwargs): + """Quantized version of a generic linear functional.""" + output = getattr(package, func_name)( + self.input_quantizer(input), + self.weight_quantizer(weight), + *args, + **kwargs, + ) + return self.output_quantizer(output) + + +class QuantLinear(_LegacyQuantLinearConvBaseMixin, nn.Linear): + """Quantized version of nn.Linear.""" + + default_quant_desc_weight = _QuantLinear.default_quant_desc_weight + + +Linear = QuantLinear diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py new file mode 100644 index 0000000000..fcd94f3722 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -0,0 +1,186 @@ +# 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. + +"""Base class for quantization modules.""" + +import contextlib +from typing import Union + +import torch + +from modelopt.torch.opt.dynamic import DynamicModule, _DMRegistryCls + +from ...qtensor import pack_real_quantize_weight +from ...tensor_quant import QUANT_DESC_8BIT_PER_TENSOR +from ...utils import is_torch_export_mode +from .tensor_quantizer import SequentialQuantizer, TensorQuantizer + +__all__ = ["QuantInputBase", "QuantLinearConvBase", "QuantModuleRegistry"] + +QuantModuleRegistry = _DMRegistryCls("Quant") + + +class QuantInputBase(DynamicModule): + """Base class for modules where the input is quantized.""" + + input_quantizer: Union[TensorQuantizer, SequentialQuantizer] + output_quantizer: Union[TensorQuantizer, SequentialQuantizer] + default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR + default_quant_desc_output = QUANT_DESC_8BIT_PER_TENSOR + + def forward(self, input, *args, **kwargs): + """Quantize the input before calling the original forward method.""" + input = self.input_quantizer(input) + output = super().forward(input, *args, **kwargs) + if isinstance(output, tuple): + return (self.output_quantizer(output[0]), *output[1:]) + return self.output_quantizer(output) + + def _setup(self): + """Patch the module's forward method to quantize the input.""" + self._register_temp_attribute( + "input_quantizer", TensorQuantizer(self.default_quant_desc_input) + ) + self._register_temp_attribute( + "output_quantizer", TensorQuantizer(self.default_quant_desc_output) + ) + self.output_quantizer.disable() + + +class QuantLinearConvBase(QuantInputBase): + """Base class for quantized linear modules. + + Quantized linear modules are modules where both the input and the weight are quantized. + """ + + weight_quantizer: Union[TensorQuantizer, SequentialQuantizer] + _enable_weight_quantization: bool + default_quant_desc_weight = QUANT_DESC_8BIT_PER_TENSOR + + @contextlib.contextmanager + def quantize_weight(self): + """Context in which `self.weight` is quantized.""" + self._enable_weight_quantization = True + yield + self._enable_weight_quantization = False + + @staticmethod + def _get_quantized_weight(module: "QuantLinearConvBase", weight: torch.Tensor) -> torch.Tensor: + if module._enable_weight_quantization or is_torch_export_mode(): + return module.weight_quantizer(weight) + return weight + + def forward(self, input, *args, **kwargs): + """Quantize the input and the weight before calling the original forward method.""" + # self.quntize_weight() setting attributes is not allowed for torch.export. + if is_torch_export_mode(): + return super().forward(input, *args, **kwargs) + with self.quantize_weight(): + return super().forward(input, *args, **kwargs) + + def _setup(self): + super()._setup() + self._register_temp_attribute( + "weight_quantizer", TensorQuantizer(self.default_quant_desc_weight) + ) + self._register_temp_attribute("_enable_weight_quantization", False) + self._register_dynamic_attribute("weight", self._get_quantized_weight) + + @staticmethod + def initialize_quantizer_with_dummy_states(module): + """Initialize the quantizer states with dummy values with the correct type and device.""" + # # Import in the function to avoid circular imports + from ...model_calib import max_calibrate + + def _initialize_activation_quantizer_amax(quantizer, device, dtype): + if not getattr(quantizer, "_has_amax", False): + return + # We need the outputs to initialize the amax in this case; Weights alone are not enough + if ( + quantizer.block_sizes is not None + and quantizer.block_sizes.get("type", None) != "dynamic" + ): + return + if quantizer.axis is not None: + return + quantizer.amax = torch.tensor(1, device=device, dtype=dtype) + + device, dtype = module.weight.device, module.weight.dtype + + for input_quantizer in SequentialQuantizer.tensor_quantizer_iterator( + getattr(module, "input_quantizer", None) + ): + if getattr(input_quantizer, "_has_pre_quant_scale", False): + input_quantizer.pre_quant_scale = torch.ones( + module.weight.shape[1], device=device, dtype=dtype + ) + + _initialize_activation_quantizer_amax(input_quantizer, device, dtype) + + for output_quantizer in SequentialQuantizer.tensor_quantizer_iterator( + getattr(module, "output_quantizer", None) + ): + _initialize_activation_quantizer_amax(output_quantizer, device, dtype) + + for weight_quantizer in SequentialQuantizer.tensor_quantizer_iterator( + getattr(module, "weight_quantizer", None) + ): + if getattr(weight_quantizer, "_has_amax", False): + max_calibrate( + weight_quantizer, lambda weight_quantizer: weight_quantizer(module.weight) + ) + + @staticmethod + def initialize_real_qtensor_with_dummy_weight(module): + """Initalize the real qunatized tensors.""" + pack_real_quantize_weight(module, force_quantize=True) + + @staticmethod + def sanitize_dummy_weight(module): + """Replace nan values with ones in dummy tensors.""" + with torch.no_grad(): + for _, m in module.named_modules(): + if ( + hasattr(m, "weight") + and hasattr(m, "weight_quantizer") + and m.weight_quantizer.is_enabled + ): + m.weight[torch.isnan(m.weight)] = 1 + + +class _LegacyQuantInputBaseMixin: + """A mixin to support legacy quantized modules which needs to have an __init__ method.""" + + _quantized_cls = QuantInputBase + default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR + default_quant_desc_output = QUANT_DESC_8BIT_PER_TENSOR + + def __init__(self, *args, quant_desc_input=None, **kwargs): + """Initialize the module with its original __init__ and patch its forward.""" + self.default_quant_desc_input = quant_desc_input or self.default_quant_desc_input + super().__init__(*args, **kwargs) + QuantModuleRegistry.convert(self) + + +class _LegacyQuantLinearConvBaseMixin(_LegacyQuantInputBaseMixin): + """A mixin to support legacy quantized modules which needs to have an __init__ method.""" + + _quantized_cls = QuantLinearConvBase + default_quant_desc_weight = QUANT_DESC_8BIT_PER_TENSOR + + def __init__(self, *args, quant_desc_input=None, quant_desc_weight=None, **kwargs): + """Initialize the module with its original __init__ and patch its forward.""" + self.default_quant_desc_weight = quant_desc_weight or self.default_quant_desc_weight + super().__init__(*args, quant_desc_input=quant_desc_input, **kwargs) diff --git a/modelopt/torch/quantization/nn/modules/quant_pooling.py b/modelopt/torch/quantization/nn/modules/quant_pooling.py new file mode 100644 index 0000000000..8270a24076 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_pooling.py @@ -0,0 +1,117 @@ +# 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. + +"""Quantized Pooling modules.""" + +from torch.nn.modules import pooling + +from .quant_module import QuantInputBase, QuantModuleRegistry, _LegacyQuantInputBaseMixin + +__all__ = [ + "MaxPool1d", + "QuantMaxPool1d", + "MaxPool2d", + "QuantMaxPool2d", + "MaxPool3d", + "QuantMaxPool3d", + "AvgPool1d", + "QuantAvgPool1d", + "AvgPool2d", + "QuantAvgPool2d", + "AvgPool3d", + "QuantAvgPool3d", + "AdaptiveAvgPool1d", + "QuantAdaptiveAvgPool1d", + "AdaptiveAvgPool2d", + "QuantAdaptiveAvgPool2d", + "AdaptiveAvgPool3d", + "QuantAdaptiveAvgPool3d", +] + + +class QuantMaxPool1d(_LegacyQuantInputBaseMixin, pooling.MaxPool1d): + """Quantized version of nn.MaxPool1d.""" + + pass + + +class QuantMaxPool2d(_LegacyQuantInputBaseMixin, pooling.MaxPool2d): + """Quantized version of nn.MaxPool2d.""" + + pass + + +class QuantMaxPool3d(_LegacyQuantInputBaseMixin, pooling.MaxPool3d): + """Quantized version of nn.MaxPool3d.""" + + pass + + +class QuantAvgPool1d(_LegacyQuantInputBaseMixin, pooling.AvgPool1d): + """Quantized version of nn.AvgPool1d.""" + + pass + + +class QuantAvgPool2d(_LegacyQuantInputBaseMixin, pooling.AvgPool2d): + """Quantized version of nn.AvgPool2d.""" + + pass + + +class QuantAvgPool3d(_LegacyQuantInputBaseMixin, pooling.AvgPool3d): + """Quantized version of nn.AvgPool3d.""" + + pass + + +class QuantAdaptiveAvgPool1d(_LegacyQuantInputBaseMixin, pooling.AdaptiveAvgPool1d): + """Quantized version of nn.AdaptiveAvgPool1d.""" + + pass + + +class QuantAdaptiveAvgPool2d(_LegacyQuantInputBaseMixin, pooling.AdaptiveAvgPool2d): + """Quantized version of nn.AdaptiveAvgPool2d.""" + + pass + + +class QuantAdaptiveAvgPool3d(_LegacyQuantInputBaseMixin, pooling.AdaptiveAvgPool3d): + """Quantized version of nn.AdaptiveAvgPool3d.""" + + pass + + +# Define alias with Quant prefix +MaxPool1d = QuantMaxPool1d +MaxPool2d = QuantMaxPool2d +MaxPool3d = QuantMaxPool3d +AvgPool1d = QuantAvgPool1d +AvgPool2d = QuantAvgPool2d +AvgPool3d = QuantAvgPool3d +AdaptiveAvgPool1d = QuantAdaptiveAvgPool1d +AdaptiveAvgPool2d = QuantAdaptiveAvgPool2d +AdaptiveAvgPool3d = QuantAdaptiveAvgPool3d + +QuantModuleRegistry.register({pooling.MaxPool1d: "nn.MaxPool1d"})(QuantInputBase) +QuantModuleRegistry.register({pooling.MaxPool2d: "nn.MaxPool2d"})(QuantInputBase) +QuantModuleRegistry.register({pooling.MaxPool3d: "nn.MaxPool3d"})(QuantInputBase) +QuantModuleRegistry.register({pooling.AvgPool1d: "nn.AvgPool1d"})(QuantInputBase) +QuantModuleRegistry.register({pooling.AvgPool2d: "nn.AvgPool2d"})(QuantInputBase) +QuantModuleRegistry.register({pooling.AvgPool3d: "nn.AvgPool3d"})(QuantInputBase) +QuantModuleRegistry.register({pooling.AdaptiveAvgPool1d: "nn.AdaptiveAvgPool1d"})(QuantInputBase) +QuantModuleRegistry.register({pooling.AdaptiveAvgPool2d: "nn.AdaptiveAvgPool2d"})(QuantInputBase) +QuantModuleRegistry.register({pooling.AdaptiveAvgPool3d: "nn.AdaptiveAvgPool3d"})(QuantInputBase) diff --git a/modelopt/torch/quantization/nn/modules/quant_rnn.py b/modelopt/torch/quantization/nn/modules/quant_rnn.py new file mode 100644 index 0000000000..b29124d7d3 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/quant_rnn.py @@ -0,0 +1,532 @@ +# 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. + +"""Quantized RNN.""" + +import contextlib +from types import ModuleType +from typing import Any, Callable, Iterator, Optional, Union + +import torch +import torch._VF as _VF +import torch.nn as nn +import torch.nn.functional as F + +from modelopt.torch.opt.dynamic import DynamicModule + +from ...tensor_quant import QUANT_DESC_8BIT_PER_TENSOR +from ...utils import is_torch_export_mode, multi_context, replace_function +from .quant_module import QuantModuleRegistry +from .tensor_quantizer import SequentialQuantizer, TensorQuantizer + +_cell_call_map = { + "RNN_TANH": _VF.rnn_tanh_cell, + "RNN_RELU": _VF.rnn_relu_cell, + "LSTM": _VF.lstm_cell, + "GRU": _VF.gru_cell, +} +_layer_call_name_map = { + "RNN_TANH": "rnn_tanh", + "RNN_RELU": "rnn_relu", + "LSTM": "lstm", + "GRU": "gru", +} + + +class QuantRNNBase(DynamicModule): + """Base class for quantized RNN modules.""" + + weight_quantizer: Union[TensorQuantizer, SequentialQuantizer] + _enable_weight_quantization: bool + default_quant_desc_weight = QUANT_DESC_8BIT_PER_TENSOR + default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR + _functionals_to_replace: list[tuple[ModuleType, str, Callable]] = [] + + @property + def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]: + """Replace functions of packages on the fly.""" + return ( + (package, func_name, quantized_func) + for package, func_name, quantized_func in self._functionals_to_replace + if hasattr(package, func_name) + ) + + @property + def all_input_quantizers_disabled(self): + """Check if all input quantizer are disabled.""" + for iq in self._input_quantizers + self._proj_input_quantizers: + if iq.is_enabled: + return False + + return True + + @contextlib.contextmanager + def quantize_weight(self): + """Context in which ``self.weight`` is quantized.""" + self._enable_weight_quantization = True + yield + self._enable_weight_quantization = False + + @staticmethod + def _get_quantized_weight_handler(weight_quantizer_name: str): + def _get_quantized_weight(module: "QuantRNNBase", weight: torch.Tensor): + if module._enable_weight_quantization: + weight_quantizer = getattr(module, weight_quantizer_name) + return weight_quantizer(weight) + return weight + + return _get_quantized_weight + + def forward(self, input, *args, **kwargs): + """Quantize the input and the weight before calling the original forward method.""" + contexts = ( + replace_function(package, func_name, quantized_func) + for package, func_name, quantized_func in self.functionals_to_replace + ) + if is_torch_export_mode(): + return super().forward(input, *args, **kwargs) + with multi_context( + self.quantize_weight(), + *contexts, + ): + return super().forward(input, *args, **kwargs) + + def _setup(self): + for name, _ in self.named_parameters(): + if name.startswith("weight"): + # to be compatible with our current config, the name is some what weired + # it would be weight_xxx_weight_quantizer + weight_quantizer_name = name + "_weight_quantizer" + self._register_temp_attribute( + weight_quantizer_name, TensorQuantizer(self.default_quant_desc_weight) + ) + self._register_dynamic_attribute( + name, self._get_quantized_weight_handler(weight_quantizer_name) + ) + # for cells + self._register_temp_attribute("_input_quantizers", []) + num_directions = 2 if self.bidirectional else 1 + # for projection layer if exists + self._register_temp_attribute("_proj_input_quantizers", []) + + # the input quantizer is per cell (or per layer) based + for layer in range(self.num_layers * num_directions): + for direction in range(num_directions): + suffix = "_reverse" if direction == 1 else "" + name = f"layer_{layer}{suffix}_input_quantizer" + quantizer = TensorQuantizer(self.default_quant_desc_input) + self._register_temp_attribute(name, quantizer) + self._input_quantizers.append(quantizer) + if self.proj_size > 0: + name = f"proj_{layer}{suffix}_input_quantizer" + quantizer = TensorQuantizer(self.default_quant_desc_input) + self._register_temp_attribute(name, quantizer) + self._proj_input_quantizers.append(quantizer) + self._register_temp_attribute("_enable_weight_quantization", False) + + +class QuantRNNFullBase(QuantRNNBase): + """Quantized RNN with input quantizer.""" + + def _disable_input_quantizers(self): + for iq in self._input_quantizers + self._proj_input_quantizers: + iq.disable() + + def _enable_input_quantizers(self): + for iq in self._input_quantizers + self._proj_input_quantizers: + iq.enable() + + def _disable_weight_quantizers(self): + for name, module in self.named_modules(): + if name.endswith("weight_quantizer"): + module.disable() + + def _enable_weight_quantizer(self): + for name, module in self.named_modules(): + if name.endswith("weight_quantizer"): + module.enable() + + def _setup(self): + super()._setup() + vf_rnn = VFRNNForward( + self.mode, + self.bidirectional, + self.num_layers, + self.proj_size > 0, + self.bias, + self._input_quantizers, + self._proj_input_quantizers if self.proj_size > 0 else None, + self.batch_first, + ) + self._functionals_to_replace = [(_VF, _layer_call_name_map[self.mode], vf_rnn)] + + +class VFRNNForward: + """Reimplement the _VF rnn calls with python to enable input quantizers. + + It's less efficient compared to oringinal calls. + """ + + def __init__( + self, + mode: str, + bidirectional: bool, + num_layers: int, + has_proj: bool, + has_bias: bool, + input_quantizers: list[TensorQuantizer], + proj_input_quantizers: Optional[list[TensorQuantizer]] = None, + batch_first: Optional[bool] = False, + ): + """Pre-construct necessary parameters for vf calls to reduce overhead. + + Refer to torch RNN modules for parameter informations. + """ + cell = _cell_call_map[mode] + if has_proj: + cell = lstm_cell_with_proj + self.layer_forwards_const = (RNNLayerForward(cell, variable_len=False),) + self.layer_forwards_variable = (RNNLayerForward(cell, variable_len=True),) + if bidirectional: + self.layer_forwards_const += (RNNLayerForward(cell, reverse=True, variable_len=False),) + self.layer_forwards_variable += ( + RNNLayerForward(cell, reverse=True, variable_len=True), + ) + self.num_directions = 2 if bidirectional else 1 + self.total_layers = num_layers * self.num_directions + self.num_layers = num_layers + self.is_lstm = mode == "LSTM" + self.has_proj = has_proj + + self.num_weights_per_layer = self.num_cell_weights_per_layer = 4 if has_bias else 2 + if has_proj: + self.num_weights_per_layer += 1 + + self.input_quantizers = input_quantizers + + self.proj_input_quantizers = proj_input_quantizers + self.batch_first = batch_first + + def forward( + self, + layer_forwards: tuple[Callable], + input: torch.Tensor, + flat_weights: list[torch.Tensor], + hidden: Union[torch.Tensor, tuple[torch.Tensor]], + dropout: Optional[float] = 0, + training: Optional[bool] = True, + batch_sizes: Optional[torch.Tensor] = None, + ): + """This this the core implementation of vf rnn calls.""" + all_hiddens = [] + + if self.is_lstm: + hidden = list(zip(*hidden)) + if self.batch_first and batch_sizes is None: + input = input.transpose(0, 1) + + for i in range(self.num_layers): + all_output = [] + for j, layer_forward in enumerate(layer_forwards): + l_i = i * self.num_directions + j + # the flat_weights is a list of tensors, we need extract weights of each layer + w_start = l_i * self.num_weights_per_layer + w_end = w_start + self.num_weights_per_layer + proj_input_quantizer = ( + self.proj_input_quantizers[l_i] if self.proj_input_quantizers else None + ) + + hx, output = layer_forward( + input, + hidden[l_i], + flat_weights[w_start:w_end], + self.input_quantizers[l_i], + batch_sizes, + proj_input_quantizer=proj_input_quantizer, + ) + + all_hiddens.append(hx) + all_output.append(output) + + input = torch.cat(all_output, -1) + + if dropout != 0 and i < self.num_layers - 1: + input = F.dropout(input, p=dropout, training=training, inplace=False) + + if self.batch_first and batch_sizes is None: + input = input.transpose(0, 1) + + if self.is_lstm: + hn, cn = zip(*all_hiddens) + all_hiddens = (torch.stack(hn, 0), torch.stack(cn, 0)) + ouput_tuple = (input, *all_hiddens) + else: + all_hiddens = torch.stack(all_hiddens, 0) + ouput_tuple = (input, all_hiddens) + + return ouput_tuple + + def __call__(self, *args) -> tuple[torch.Tensor, torch.Tensor]: + """Entry of vf calls. + + Original vf funcs are overloaded cpp funcs. Each has two different signatures, + one accepts inputs with variable length (packed sequence), the other accepts constant length. + The difference is that the fourth arg is a bool for constant length. + """ + if isinstance(args[3], bool): + # constant + ( + input, + hidden, + flat_weights, + bias, + num_layers, + dropout, + training, + bidirectional, + batch_first, + ) = args + batch_sizes = None + layer_forwards = self.layer_forwards_const + else: + # variable + ( + input, + batch_sizes, + hidden, + flat_weights, + bias, + num_layers, + dropout, + training, + bidirectional, + ) = args + layer_forwards = self.layer_forwards_variable + + return self.forward( + layer_forwards, + input, + flat_weights, + hidden, + dropout, + training, + batch_sizes, + ) + + +class RNNLayerForward: + """A single layer of rnn modules.""" + + def __init__(self, cell, reverse=False, variable_len=False): + """Init the layer forward for different cells, directions, and inputs.""" + self.cell = cell + self.reverse = reverse + self.variable_len = variable_len + if variable_len: + if reverse: + self.forward = get_quantized_rnn_layer_variable_len_reverse_forward(cell) + else: + self.forward = get_quantized_rnn_layer_variable_len_forward(cell) + else: + self.forward = get_quantized_rnn_layer_forward(cell, reverse=reverse) + + def __call__( + self, input, hidden, weights, input_quantizer, batch_sizes=None, proj_input_quantizer=None + ) -> Any: + """Layer forward.""" + return self.forward( + input, + hidden, + weights, + input_quantizer, + batch_sizes=batch_sizes, + proj_input_quantizer=proj_input_quantizer, + ) + + +def lstm_cell_with_proj(input, hidden, *weights, proj_input_quantizer=None): + """Currently the _VF.lstm_cell doesn't accept projected inputs. i.e. h_n and c_n must be same shape. + + This implementation is not optimized for cuda compared to _VF.lstm_cell, so we only use it when projection exists. + """ + if len(weights) == 3: + weight_ih, weight_hh, weight_hr = weights + bias_ih = bias_hh = None + else: + weight_ih, weight_hh, bias_ih, bias_hh, weight_hr = weights + hn, cn = hidden + + gates = F.linear(input, weight_ih, bias_ih) + F.linear(hn, weight_hh, bias_hh) + + ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1) + + ingate = torch.sigmoid(ingate) + forgetgate = torch.sigmoid(forgetgate) + cellgate = torch.tanh(cellgate) + outgate = torch.sigmoid(outgate) + + cy = (forgetgate * cn) + (ingate * cellgate) + hy = outgate * torch.tanh(cy) + if proj_input_quantizer is not None: + hy = proj_input_quantizer(hy) + hy = F.linear(hy, weight_hr) + return hy, cy + + +def quantized_cell_forward( + cell, input, hidden, weights, input_quantizer, proj_input_quantizer=None +): + """Call input quantizer before calling cell.""" + quant_input = input_quantizer(input) + quant_hidden = ( + (input_quantizer(hidden[0]), input_quantizer(hidden[1])) + if isinstance(hidden, tuple) + else input_quantizer(hidden) + ) + kwargs = {} if proj_input_quantizer is None else {"proj_input_quantizer": proj_input_quantizer} + hidden = cell(quant_input, quant_hidden, *weights, **kwargs) + return hidden + + +def get_quantized_rnn_layer_forward(cell, reverse=False): + """Construct the forward call for different rnn cells. + + Note that batch_sizes is here for keeping a consistant signature with the forward of variable length. + """ + + def forward( + input, hidden, weights, input_quantizer, batch_sizes=None, proj_input_quantizer=None + ): + seq_len = input.shape[0] + steps = reversed(range(seq_len)) if reverse else range(seq_len) + + output = [] + for i in steps: + hidden = quantized_cell_forward( + cell, + input[i], + hidden, + weights, + input_quantizer, + proj_input_quantizer=proj_input_quantizer, + ) + output.append(hidden[0] if isinstance(hidden, tuple) else hidden) + + if reverse: + output.reverse() + output = torch.stack(output, 0) + + return hidden, output + + return forward + + +def get_quantized_rnn_layer_variable_len_forward(cell): + """Construct the forward call for packed sequence.""" + + def forward(input, hidden, weights, input_quantizer, batch_sizes, proj_input_quantizer=None): + output = [] + input_offset = 0 + last_batch_size = batch_sizes[0] + hiddens = [] + flat_hidden = not isinstance(hidden, tuple) + if flat_hidden: + hidden = (hidden,) + for batch_size in batch_sizes: + step_input = input[input_offset : input_offset + batch_size] + input_offset += batch_size + + dec = last_batch_size - batch_size + if dec > 0: + hiddens.append(tuple(h[-dec:] for h in hidden)) + hidden = tuple(h[:-dec] for h in hidden) + last_batch_size = batch_size + hx = hidden[0] if flat_hidden else hidden + hidden = quantized_cell_forward( + cell, + step_input, + hx, + weights, + input_quantizer, + proj_input_quantizer=proj_input_quantizer, + ) + if flat_hidden: + hidden = (hidden,) + + output.append(hidden[0]) + hiddens.append(hidden) + hiddens.reverse() + + hidden = tuple(torch.cat(h, 0) for h in zip(*hiddens)) + assert hidden[0].shape[0] == batch_sizes[0] + if flat_hidden: + hidden = hidden[0] + output = torch.cat(output, 0) + + return hidden, output + + return forward + + +def get_quantized_rnn_layer_variable_len_reverse_forward(cell): + """Construct the forward call for packed sequence in the reversed direction.""" + + def forward(input, hidden, weights, input_quantizer, batch_sizes, proj_input_quantizer=None): + output = [] + input_offset = input.shape[0] + last_batch_size = batch_sizes[-1] + initial_hidden = hidden + flat_hidden = not isinstance(hidden, tuple) + if flat_hidden: + hidden = (hidden,) + initial_hidden = (initial_hidden,) + hidden = tuple(h[: batch_sizes[-1]] for h in hidden) + for i in reversed(range(len(batch_sizes))): + batch_size = batch_sizes[i] + inc = batch_size - last_batch_size + if inc > 0: + hidden = tuple( + torch.cat((h, ih[last_batch_size:batch_size]), 0) + for h, ih in zip(hidden, initial_hidden) + ) + last_batch_size = batch_size + step_input = input[input_offset - batch_size : input_offset] + input_offset -= batch_size + + hx = hidden[0] if flat_hidden else hidden + hidden = quantized_cell_forward( + cell, + step_input, + hx, + weights, + input_quantizer, + proj_input_quantizer=proj_input_quantizer, + ) + if flat_hidden: + hidden = (hidden,) + output.append(hidden[0]) + + output.reverse() + output = torch.cat(output, 0) + if flat_hidden: + hidden = hidden[0] + return hidden, output + + return forward + + +QuantModuleRegistry.register({nn.RNN: "nn.RNN"})(QuantRNNFullBase) +QuantModuleRegistry.register({nn.LSTM: "nn.LSTM"})(QuantRNNFullBase) +QuantModuleRegistry.register({nn.GRU: "nn.GRU"})(QuantRNNFullBase) diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py new file mode 100644 index 0000000000..957b0f3eb2 --- /dev/null +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -0,0 +1,984 @@ +# 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. + +"""TensorQuantizer Module.""" + +import contextlib +import math +import warnings +from typing import Any, Callable, Optional, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import nn +from torch.onnx._globals import GLOBALS + +from modelopt.torch.utils import standardize_constructor_args +from modelopt.torch.utils.distributed import DistributedProcessGroup + +from ... import calib +from ... import utils as quant_utils +from ...config import QuantizerAttributeConfig +from ...qtensor import BaseQuantizedTensor, INT4QTensor, NF4QTensor, NVFP4QTensor, QTensorWrapper +from ...tensor_quant import dynamic_block_quant, fake_tensor_quant, scaled_e4m3 +from ...utils import is_torch_export_mode +from .clip import Clip + +__all__ = ["TensorQuantizer", "SequentialQuantizer"] + + +class TensorQuantizer(nn.Module): + """Tensor quantizer module. + + This module manages quantization and calibration of input tensor. It can perform fake (simulated quantization) + or real quantization for various precisions and formats such as FP8 per-tensor, INT8 per-channel, + INT4 per-block etc. + + If quantization is enabled, it calls the appropriate quantization functional and + returns the quantized tensor. The quantized tensor data type will be same as the input tensor data type for + fake quantization. During calibration mode, the module collects the statistics using its calibrator. + + The quantization parameters are as described in + :class:`QuantizerAttributeConfig `. They can be set + at initialization using ``quant_attribute_cfg`` or later by calling :meth:`set_from_attribute_config`. + + Args: + quant_attribute_cfg: An instance of + :class:`QuantizerAttributeConfig ` or None. + If None, default values are used. + if_quant: A boolean. If True, quantization is enabled in the forward path. + if_clip: A boolean. If True, clipping (with ``_learn_amax``) is enabled in the forward path. + if_calib: A boolean. If True, calibration is enabled in the forward path. + amax: None or an array like object such as list, tuple, numpy array, scalar + which can be used to construct amax tensor. + """ + + def __init__( + self, + quant_attribute_cfg=None, + if_quant=True, + if_clip=False, + if_calib=False, + amax=None, + ): + """Initialize quantizer and set up required variables.""" + super(TensorQuantizer, self).__init__() + quant_attribute_cfg = ( + quant_attribute_cfg if quant_attribute_cfg is not None else QuantizerAttributeConfig() + ) + if amax is not None: + self.amax = amax + + self.set_from_attribute_config(quant_attribute_cfg) + + self._if_quant = if_quant + self._if_clip = if_clip + self._if_calib = if_calib + self._enable_pre_quant_scale = True + self._dequantize = False + self._input_dtype = None + + def set_from_attribute_config(self, attribute_cfg: Union[QuantizerAttributeConfig, dict]): + """Set quantizer attributes from attribute_dict. + + The attributes are defined in + :class:`QuantizerAttributeConfig `. + """ + + def _calibrator_setter(val): + if val in ["max", "histogram"]: + calib_cls = calib.MaxCalibrator if val == "max" else calib.HistogramCalibrator + args, kwargs = (self._num_bits, self._axis, self._unsigned), {} + else: + calib_cls, args, kwargs = standardize_constructor_args(val) + return calib_cls(*args, **kwargs) + + # Some attributes need custom handling. + # By default, attributes from config are mapped to a name ``f"_{attribute}"`` + _custom_setters: dict[str, tuple[str, Callable]] = { + "enable": ("_disabled", lambda val: val is False), + "type": ("_dynamic", lambda val: val == "dynamic"), + "calibrator": ("_calibrator", _calibrator_setter), + } + + for attribute, val in attribute_cfg.items(): + assert ( + attribute in QuantizerAttributeConfig.model_fields + ), f"{attribute} is not a valid `TensorQuantizer` attribute" + _tq_attribute_name, _setter = _custom_setters.get( + attribute, (f"_{attribute}", lambda v: v) + ) + setattr(self, _tq_attribute_name, _setter(val)) + + # Clip module consumes a lot of memory, so only create it if learn_amax is True + if "learn_amax" in attribute_cfg and self._learn_amax: + init_amax = self.amax if hasattr(self, "_amax") else 1.0 + self.clip = Clip(-init_amax, init_amax, learn_min=True, learn_max=True) + self.enable_clip() + + def dequantize(self, qtensor: BaseQuantizedTensor): + """De-quantize a real quantized tensor to a given dtype.""" + kwarg = { + "scale": self._scale, + "block_sizes": self.block_sizes, + "double_scale": getattr(self, "_double_scale", None), + "scale_zeros": getattr(self, "_scale_zeros", None), + } + return qtensor.dequantize(**kwarg) + + @property + def num_bits(self): + """Return num_bits for quantization.""" + return self._num_bits + + @num_bits.setter + def num_bits(self, value): + self._num_bits = value + self._calibrator._num_bits = value + + @property + def maxbound(self): + """Return maxbound for quantization.""" + if self._num_bits == (4, 3): + return 448.0 + if self._num_bits == (2, 1) and self._block_sizes.get("scale_bits") == (4, 3): + return 6.0 + return (1 << (self._num_bits - 1 + int(self._unsigned))) - 1 + + @property + def unsigned(self): + """Return True if unsigned quantization is used.""" + return self._unsigned + + @unsigned.setter + def unsigned(self, value): + self._unsigned = value + self._calibrator._unsigned = value + + @property + def pre_quant_scale(self): + """Return pre_quant_scale used for smoothquant.""" + if not hasattr(self, "_pre_quant_scale") or not self._enable_pre_quant_scale: + return None + return self._pre_quant_scale + + @pre_quant_scale.setter + def pre_quant_scale(self, value): + assert value is not None, "pre_quant_scale cannot be set to None." + assert ( + self._enable_pre_quant_scale + ), "pre_quant_scale cannot be set when forward_with_pre_quant_scale is False." + if not isinstance(value, torch.Tensor): + value = torch.tensor(value) + if not hasattr(self, "_pre_quant_scale"): + self.register_buffer("_pre_quant_scale", value.clone().detach()) + else: + if self._pre_quant_scale.shape != value.shape: + raise RuntimeError("Changing shape when setting pre_quant_scale is not allowed.") + self._pre_quant_scale.data.copy_( + value.clone().detach().to(self._pre_quant_scale.device) + ) + + @property + def amax(self): + """Return amax for quantization.""" + if not hasattr(self, "_amax") or self.is_mx_format: + return None + assert not self._dynamic, "Dynamic quantization does not have fixed amax" + return self._amax + + @amax.setter + def amax(self, value): + assert value is not None, "amax cannot be set to None." + + if not isinstance(value, torch.Tensor): + value = torch.tensor(value) + + if not hasattr(self, "_amax"): + self.register_buffer("_amax", value.clone().detach()) + else: + if self._amax.shape != value.shape: + raise RuntimeError("Changing shape when setting amax is not allowed.") + self._amax.data.copy_(value.clone().detach().to(self._amax.device)) + + def reset_amax(self): + """Reset amax to None.""" + if hasattr(self, "_amax"): + delattr(self, "_amax") + self._calibrator.reset() + + @property + def step_size(self): + """Return step size for integer quantization.""" + if not hasattr(self, "_amax"): + warnings.warn("step_size is undefined under dynamic amax mode!") + return None + assert isinstance( + self._num_bits, int + ), "Step size is not defined for non-integer quantization." + return self._amax / (2.0 ** (self._num_bits - 1 + int(self._unsigned)) - 1.0) + + @property + def axis(self): + """Return axis for quantization.""" + return self._axis + + @axis.setter + def axis(self, value): + self._axis = value + self._calibrator._axis = value + + @property + def block_sizes(self): + """Return block_sizes for quantization.""" + return self._block_sizes + + @block_sizes.setter + def block_sizes(self, value): + self._axis = None + self._block_sizes = value + + @property + def fake_quant(self): + """Return True if fake quantization is used.""" + return self._fake_quant + + @property + def narrow_range(self): + """Return True if symmetric integer range for signed quantization is used.""" + return self._narrow_range + + @narrow_range.setter + def narrow_range(self, value): + self._narrow_range = value + + @property + def is_enabled(self): + """Return true if the modules is not disabled.""" + return not self._disabled + + def disable(self): + """Bypass the module. + + Neither of calibration, clipping and quantization will be performed if the module is disabled. + """ + self._disabled = True + + def enable(self): + """Enable the module.""" + self._disabled = False + + @property + def trt_high_precision_dtype(self): + """Return True if FP16 AMAX is used when exporting the model.""" + return self._trt_high_precision_dtype + + @trt_high_precision_dtype.setter + def trt_high_precision_dtype(self, value): + self._trt_high_precision_dtype = value + + @property + def is_mx_format(self): + """Check if is MX formats.""" + return ( + self.block_sizes is not None + and self.block_sizes.get("type", None) == "dynamic" + and self.block_sizes.get("scale_bits", None) == (8, 0) + ) + + def disable_clip(self): + """Disable clip stage.""" + self._if_clip = False + self.clip.clip_value_min.requires_grad = False + self.clip.clip_value_max.requires_grad = False + + def enable_clip(self): + """Enable clip stage.""" + if not self._learn_amax: + raise ValueError("learn_amax is False. Cannot enable clip.") + self.clip.clip_value_min.requires_grad = True + self.clip.clip_value_max.requires_grad = True + self._if_clip = True + + def disable_calib(self): + """Disable calibration.""" + self._if_calib = False + + def enable_calib(self): + """Enable calibration.""" + if self._calibrator is None: + raise ValueError("Calibrator was not created, cannot enable calibration.") + + # Dynamic quantization does not need calibration. + if self._dynamic: + return + self._if_calib = True + + def disable_quant(self): + """Disable quantization.""" + self._if_quant = False + + def enable_quant(self): + """Enable quantization.""" + self._if_quant = True + + def load_calib_amax(self, *args, **kwargs): + """Load amax from calibrator. + + Updates the amax buffer with value computed by the calibrator, creating it if necessary. + ``*args`` and ``**kwargs`` are directly passed to ``compute_amax``, except ``"strict"`` in + ``kwargs``. Refer to ``compute_amax`` for more details. + """ + assert not self._dynamic, "Dynamic quantization does not need calibration." + + strict = kwargs.pop("strict", True) + if getattr(self, "_calibrator", None) is None: + raise RuntimeError("Calibrator not created.") + calib_amax = self._calibrator.compute_amax(*args, **kwargs) + if calib_amax is None: + err_msg = ( + "Calibrator returned None. This usually happens when calibrator hasn't seen any" + " tensor." + ) + if not strict: + warnings.warn(err_msg) + warnings.warn("Set amax to NaN!") + calib_amax = torch.tensor(math.nan) + else: + raise RuntimeError( + err_msg + + " Passing 'strict=False' to `load_calib_amax()` will ignore the error." + ) + if not hasattr(self, "_amax"): + self.register_buffer("_amax", calib_amax.clone().detach()) + else: + self._amax.data.copy_(calib_amax.clone().detach()) + + def init_learn_amax(self): + """Initialize learned amax from fixed amax.""" + if self._learn_amax is False: + raise RuntimeError("Called init_learn_amax with learn_amax=False.") + + if self._amax.numel() != 1: + warnings.warn("Per channel learned amax not supported. Initializing with max(amax).") + init_amax = torch.max(self._amax) + else: + init_amax = self._amax + self.clip.clip_value_min.data.copy_(-init_amax.clone().detach()) + self.clip.clip_value_max.data.copy_(init_amax.clone().detach()) + + def _get_amax(self, inputs): + """Get amax from buffer or compute it dynamically.""" + if hasattr(self, "_amax"): + amax = self._amax + else: + if self._axis is None: + reduce_axis = None + else: + reduce_axis = [] + # Swap axis to reduce + axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis] + for i in range(inputs.dim()): + if not (i in axis or (i - inputs.dim()) in axis): + reduce_axis.append(i) + amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach() + + amax = amax.detach() if is_torch_export_mode() else amax.data + return amax + + def _validate_amax(self, amax): + # Dynamic control flow is not supported by torch dynamo + if not is_torch_export_mode(): + assert torch.all(amax >= 0) and not torch.any( + torch.isinf(amax) + ), f"Got invalid amax: {amax}" + + def _real_quantize(self, inputs): + assert ( + self._num_bits == 4 or self._num_bits == (2, 1) + ) and self._block_sizes, "Real quantization not supported for this format." + + buffer_to_register = {} + if self._block_sizes.get("scale_bits", 0) == 8 and self._block_sizes.get( + "scale_block_sizes", None + ): + # NF4 double quantization + # Return real quantized tensor class and store scales inside the TensorQuantizer + outputs, scales = NF4QTensor.quantize(inputs, self._block_sizes[-1]) + _scale, _double_scale, _scale_zeros = NF4QTensor.double_quantization( + scales, + self._block_sizes["scale_block_sizes"][-1], + self._block_sizes["scale_bits"], + ) + buffer_to_register["_scale"] = _scale + buffer_to_register["_double_scale"] = _double_scale + buffer_to_register["_scale_zeros"] = _scale_zeros + elif self._block_sizes.get("scale_bits") == (4, 3): + # NVFP4 default quantization + # Return real quantized tensor and store scales inside TensorQuantizer + outputs, _weights_scaling_factor, _weights_scaling_factor_2 = NVFP4QTensor.quantize( + inputs, + self._block_sizes[-1], + ) + buffer_to_register["_scale"] = _weights_scaling_factor.to(torch.float32) + buffer_to_register["_double_scale"] = _weights_scaling_factor_2 + else: + outputs, _scale = INT4QTensor.quantize(inputs, self._block_sizes[-1]) + buffer_to_register["_scale"] = _scale + for k, v in buffer_to_register.items(): + self.register_buffer(k, v) + + return outputs + + def _quant_forward(self, inputs): + """Quantized forward pass.""" + amax = None + if not self._dequantize and not self.is_mx_format: + if self._learn_amax: + inputs = self.clip(inputs) + amax = torch.max(-self.clip.clip_value_min, self.clip.clip_value_max).detach() + else: + amax = self._get_amax(inputs) + + self._validate_amax(amax) + + if self._fake_quant: + if self.block_sizes is not None and self.block_sizes.get("type", None) == "dynamic": + block_size = self.block_sizes.get(-1, None) or self.block_sizes.get( + inputs.dim() - 1, None + ) + assert block_size is not None, "block size for dynamic quantization not found." + outputs = dynamic_block_quant( + inputs, + block_size, + amax, + self._num_bits, + self.block_sizes.get("scale_bits", None), + getattr(self, "_trt_high_precision_dtype", None), + getattr(self, "_onnx_quantizer_type", None), + ) + + elif isinstance(self._num_bits, tuple): + E, M = self._num_bits # noqa: N806 + outputs = scaled_e4m3( + inputs, self._get_amax(inputs), E, M, self._trt_high_precision_dtype + ) + else: + outputs = fake_tensor_quant( + inputs, + amax, + self._num_bits, + self._unsigned, + self._narrow_range, + self._trt_high_precision_dtype, + ) + else: + # Real quantize + if not self._dequantize: + outputs = self._real_quantize(inputs) + self._dequantize = True + else: + # De-quantize + if isinstance(inputs, QTensorWrapper): + inputs = inputs.get_qtensor() + assert isinstance( + inputs, BaseQuantizedTensor + ), "Expected input as real quantized tensors." + return self.dequantize(inputs) + return outputs + + def _check_onnx_readiness(self, inputs): + """Check if quantizer is ready for ONNX export.""" + assert hasattr(self, "_amax"), ( + "Quantizer has not been calibrated. ONNX export requires the quantizer to be" + " calibrated. Calibrate and load amax before exporting to ONNX." + ) + + if self._if_calib: + warnings.warn( + "Quantizer is in calibration mode. " + "Please complete calibration before exporting to ONNX for correct results." + ) + + amax = self._get_amax(inputs) + + # We only support scalar amax for E4M3 ONNX export + if isinstance(self.num_bits, tuple): + assert amax.numel() == 1, ( + "E4M3 supports ONNX export only for per-tensor quantization." + " Per-tensor quantization requires scalar amax. " + f"Received non-scalar amax of shape: {amax.shape}" + ) + + def _setup_for_blockquant(self, inputs: torch.Tensor): + # Get reshape sizes and paddings for block-quantization + def get_axis_quant_params(ax): + ax = ax if ax in self.block_sizes else ax - inputs.dim() + bsize = self.block_sizes.get(ax, None) + padding, ax_slice = None, None + if bsize is not None and inputs.shape[ax] % bsize != 0: + padding = (bsize - (inputs.shape[ax] % bsize), 0) + ax_slice = slice(inputs.shape[ax]) + return bsize, padding, ax_slice + + def set_quant_params(axis, block_reshape_size, padding, slices, amax_shape=None): + self._axis = tuple(axis) + if hasattr(self, "_calibrator"): + self._calibrator._axis = self._axis + self._original_shape = inputs.shape + self._block_reshape_size = torch.Size(block_reshape_size) + if padding is not None: + self._padding = tuple(padding) + self._original_shape = F.pad(inputs, self._padding, "constant", 0).shape + if slices is not None: + self._slices = slices + if amax_shape: + self._amax_shape_for_export = amax_shape + + # Reshape size have already been set + if hasattr(self, "_block_reshape_size"): + return + + reshape_size, quantize_axis, paddings, slices = [], [], [], [] + + # special handling for block-quantization along the last axis: + # flatten the input for faster execution + if (self.block_sizes.get(inputs.dim() - 1, None) or self.block_sizes.get(-1, None)) and len( + QuantizerAttributeConfig._get_block_quant_axes_and_sizes(self.block_sizes) + ) == 1: + bsize, padding, ax_slice = get_axis_quant_params(inputs.dim() - 1) + slices = None if ax_slice is None else (*(slice(None),) * (inputs.dim() - 1), ax_slice) + padding = padding if not padding else tuple(reversed(padding)) + amax_shape_for_export = (*(inputs.shape[:-1]), -1) + set_quant_params((0,), (-1, bsize), padding, slices, amax_shape_for_export) + return + + for ax in range(inputs.dim()): + bsize, padding, ax_slice = get_axis_quant_params(ax) + paddings.append(padding) + slices.append(ax_slice) + if bsize is not None: + reshape_size.extend([math.ceil(inputs.shape[ax] / bsize), bsize]) + quantize_axis.extend([True, False]) + else: + reshape_size.append(inputs.shape[ax]) + quantize_axis.append(True) + + quant_axis = [i for i in range(len(quantize_axis)) if quantize_axis[i]] + + if all(s is None for s in slices): + slices = None + else: + slices = [s if s else slice(None) for s in slices] + + if all(p is None for p in paddings): + paddings = None + else: + new_paddings = [] + for padding in paddings: + if not (new_paddings or padding): + continue + new_paddings.extend(padding if padding else (0, 0)) + paddings = tuple(reversed(new_paddings)) + + set_quant_params(quant_axis, reshape_size, paddings, slices) + + def _process_for_blockquant(self, inputs: torch.Tensor): + if hasattr(self, "_padding"): + inputs = F.pad(inputs, self._padding, "constant", 0) + assert inputs.shape == self._original_shape, ( + f"Input shape has changed from {self._original_shape} to {inputs.shape}." + " Block-quantization requires a fixed input shape." + ) + inputs = inputs.reshape(self._block_reshape_size) + return inputs + + def _reset_to_original_shape(self, outputs: torch.Tensor): + outputs = outputs.reshape(self._original_shape) + if hasattr(self, "_slices"): + outputs = outputs[self._slices] + return outputs + + def export_amax(self) -> Optional[torch.Tensor]: + """Export correctly formatted/shaped amax.""" + if self.block_sizes is not None and self.block_sizes.get("type", None) == "dynamic": + return self.amax + + if self.amax is None: + return None + + if not hasattr(self, "_amax_shape_for_export"): + amax = self.amax + else: + amax = self.amax.reshape(self._amax_shape_for_export) + amax[amax == 0] = self.maxbound + clamp_min, clamp_max = torch.finfo(amax.dtype).tiny, torch.finfo(amax.dtype).max + amax = amax.clamp(min=clamp_min, max=clamp_max) + + self._validate_amax(amax) + + if self.block_sizes is None: + # tensorrt_llm assumes the scaling_factor dim >= 1 for per-tensor. + if self.axis is None: + amax = amax.unsqueeze(0) + + # If single-axis quantization, squeeze amax + elif isinstance(self.axis, int) or ( + isinstance(self.axis, (list, tuple)) and len(self.axis) == 1 + ): + amax = amax.squeeze() + + return amax + + def forward(self, inputs): + """Apply tensor_quant function to inputs. + + Args: + inputs: A Tensor of type float32/float16/bfloat16. + + Returns: + outputs: A Tensor of type output_dtype + """ + if isinstance(inputs, BaseQuantizedTensor): + assert self._dequantize, "No dequantization stats in the tensor quantizer." + return self._quant_forward(inputs) + + # Early return if nothing is collected during the forward (e.g. MoE) + # len(inputs) will break the dynamic shape for torch_export + if not is_torch_export_mode() and len(inputs) == 0: + return inputs + + # Activation scaling for smoothquant + if self.pre_quant_scale is not None: + inputs = inputs * self.pre_quant_scale + + if self._disabled: + # if quantizer is disabled, we still need to track the input dtype for saving the model + # TODO: This is a temporary solution and needs to be removed once megatron supports + # non-homogeneous layers + self._input_dtype = inputs.dtype if hasattr(inputs, "dtype") else None + return inputs + + # GLOBALS could break TorchDynamo for some Pytorch versions (i.e., 2.3.0) + if not is_torch_export_mode(): + if GLOBALS.in_onnx_export: + self._check_onnx_readiness(inputs) + + if ( + self.block_sizes is not None + and not self.block_sizes.get("type", None) == "dynamic" + and self._fake_quant + ): + # Dynamic block quantization is handled seperately by the quantization kernels + self._setup_for_blockquant(inputs) + inputs = self._process_for_blockquant(inputs) + + outputs = inputs + + if self._if_calib and not self._dynamic: + if self._calibrator is None: + raise RuntimeError("Calibrator was not created.") + # Shape is only known when it sees the first tensor + self._calibrator.collect(inputs) + + if self._if_clip: + if not self._learn_amax: + raise RuntimeError("Clip without learning amax is not implemented.") + outputs = self.clip(inputs) + + if self._if_quant: + outputs = self._quant_forward(inputs) + + if ( + self.block_sizes is not None + and not self.block_sizes.get("type", None) == "dynamic" + and self._fake_quant + ): + outputs = self._reset_to_original_shape(outputs) + + return outputs + + def _short_amax(self, fmt=".4f"): + """Short description of amax. + + Returns: + 'dynamic': if _amax is not registered + 'amax': if _amax is per-tensor + '[min, max](size)': if _amax is per-channel + """ + if self.is_mx_format: + return "None" + if not hasattr(self, "_amax"): + return "dynamic" + if self._amax is None: + return "None" + if self._amax.is_meta: + return "meta" + if self._amax.numel() == 1: + return f"{self._amax.item():{fmt}}" + return ( + f"[{self._amax.min().item():{fmt}}," + f" {self._amax.max().item():{fmt}}]({self._amax.numel()})" + ) + + def extra_repr(self): + """Set the extra information about this module.""" + if self._disabled: + return "disabled" + s = f"{'unsigned ' if self._unsigned else ''}{self._num_bits} bit" + s += " narrow" if (self._narrow_range) else "" + s += " fake" if (self._fake_quant) else "" + if self.block_sizes is not None: + s += f" block_sizes={self._block_sizes}," + else: + s += f" axis={self._axis}" if self._axis is not None else " per-tensor" + s += f" amax={self._short_amax()}" + s += " pre_quant_scale" if self.pre_quant_scale is not None else "" + s += " learned" if (self._learn_amax) else "" + s += ( + f" calibrator={self._calibrator.__class__.__name__}" + if (self._calibrator is not None) + else "" + ) + s += " quant" if (self._if_quant) else "" + s += " clip" if (self._if_clip) else "" + s += " calib" if (self._if_calib) else "" + return s + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Overloaded module function. + + Adds warnings during state_dict loading. + A workaround is implemented for loading amax from checkpoint and only supports CUDA. + + Args: + state_dict: A dict containing the state of the top level module + prefix: A string that prefixes all of this modules state in state_dict, e.g. 'model.conv1.' + """ + dst_has_amax = "_amax" in self._buffers + src_has_amax = prefix + "_amax" in state_dict + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + if not src_has_amax and dst_has_amax: + warnings.warn(f"{prefix[:-1]}: No amax in state_dict.") + elif src_has_amax and not dst_has_amax: + warnings.warn( + f"{prefix[:-1]}: No '_amax' buffer to load amax into." + " '_amax` will be created as WAR for now. " + "This behavior will change in future." + ) + self.register_buffer("_amax", state_dict[prefix + "_amax"].clone().detach().to(device)) + + dst_has_pre_quant_scale = "_pre_quant_scale" in self._buffers + src_has_pre_quant_scale = prefix + "_pre_quant_scale" in state_dict + + if not src_has_pre_quant_scale and dst_has_pre_quant_scale: + warnings.warn(f"{prefix[:-1]}: No pre_quant_scale in state_dict.") + elif src_has_pre_quant_scale and not dst_has_pre_quant_scale: + warnings.warn( + f"{prefix[:-1]}: No '_pre_quant_scale' buffer to load pre_quant_scale into." + " '_pre_quant_scale` will be created as WAR for now. " + "This behavior will change in future." + ) + self.register_buffer( + "_pre_quant_scale", + state_dict[prefix + "_pre_quant_scale"].clone().detach().to(device), + ) + + super(TensorQuantizer, self)._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def _get_skip_properties_for_modelopt_state(self): + return {"clip", "_calibrator", "_original_shape", "_block_reshape_size", "_padding"} + + def _get_properties_for_modelopt_state(self): + return ( + self.__dict__.keys() + - nn.Module().__dict__.keys() + - self._get_skip_properties_for_modelopt_state() + ) + + def get_modelopt_state(self, properties_only: bool = False) -> dict[str, Any]: + """Get meta state to be saved in checkpoint. + + If `properties_only` is True, only the quantizer properties such as `num_bits`, `axis` etc are included. + For restoring the quantizer fully, use `properties_only=False`. + """ + modelopt_state = {} + for k in self._get_properties_for_modelopt_state(): + modelopt_state[k] = getattr(self, k) + + if properties_only: + return modelopt_state + + if hasattr(self, "_amax"): + modelopt_state["_has_amax"] = True + + if hasattr(self, "_pre_quant_scale"): + modelopt_state["_has_pre_quant_scale"] = True + + if hasattr(self, "clip"): + modelopt_state["_init_clip"] = True + + return modelopt_state + + def set_from_modelopt_state(self, modelopt_state, prefix=""): + """Set meta state from checkpoint.""" + # Set all properties except the skip properties; this is done for backward compatibility + for key in modelopt_state.keys() - self._get_skip_properties_for_modelopt_state(): + setattr(self, key, modelopt_state[key]) + + # Set the calibrator properties + # TODO: This might not be sufficient for the custom calibrators - however there is no use-case for it yet + for key in ["_num_bits", "_axis", "_unsigned"]: + setattr(self._calibrator, key, getattr(self, key)) + + if "_init_clip" in modelopt_state: + # clip min and max parameters will be loaded from checkpoint + self.clip = Clip(-1.0, 1.0, learn_min=True, learn_max=True) + + # Create a temporary variable to indicate if the quantizer had amax in the checkpoint + if "_has_amax" in modelopt_state or "_amax" in modelopt_state: + self._has_amax = modelopt_state.get("_has_amax", "_amax" in modelopt_state) + + # Create a temporary variable to indicate if the quantizer had pre_quant_scale in the checkpoint + if "_has_pre_quant_scale" in modelopt_state or "_pre_quant_scale" in modelopt_state: + self._has_pre_quant_scale = modelopt_state.get( + "_has_pre_quant_scale", "_pre_quant_scale" in modelopt_state + ) + + def clean_up_after_set_from_modelopt_state(self, prefix=""): + """Clean up temporary variables created during set_from_modelopt_state.""" + warning_msg = ( + f"Could not initialize the quantizer states for {prefix}. The quantizer" + " states after `load_state_dict` could be in the wrong device. Please move" + " the modules to the correct device after loading the state dict." + ) + + if hasattr(self, "_has_amax"): + if self._has_amax and self.amax is None: + warnings.warn(warning_msg, UserWarning) + delattr(self, "_has_amax") + + if hasattr(self, "_has_pre_quant_scale"): + if self._has_pre_quant_scale and self.pre_quant_scale is None: + warnings.warn(warning_msg, UserWarning) + delattr(self, "_has_pre_quant_scale") + + def sync_amax_across_distributed_group(self, parallel_group: DistributedProcessGroup): + """Synchronize the amax across all ranks in the given group.""" + if parallel_group.is_initialized() and self.amax is not None: + try: + dist.all_reduce(self.amax, op=dist.ReduceOp.MAX, group=parallel_group.group) + except RuntimeError as e: + # This error happens if the distributed backend is using GPU and + # the tensor is not on GPU (or vice versa). + warnings.warn( + ( + f"Failed to synchronize amax: {e}, probably because the tensor is on a device which is not" + "supported by the current distributed backend. This warning can be ignored" + "if happening during modelopt restore." + ) + ) + + +class SequentialQuantizer(nn.Sequential): + """A sequential container for :class:`TensorQuantizer` modules. + + This modules is used to quantize a tensor in multiple formats sequentially. It takes as input + :class:`TensorQuantizer` modules and containerize them similar to :class:`torch.nn.Sequential`. + + Args: + quantizers (TensorQuantizer): :class:`TensorQuantizer` modules to be added to the container. + + """ + + def __init__(self, *quantizers: TensorQuantizer): # noqa: N803 + """Initialize SequentialQuantizer module.""" + assert not any( + not isinstance(q, TensorQuantizer) for q in quantizers + ), "All quantizers must be a TensorQuantizer." + super().__init__(*quantizers) + + def get_modelopt_state(self) -> dict[str, Any]: + """Get meta state to be saved in checkpoint.""" + return {"num_quantizers": len(self), "is_sequential_quantizer": True} + + def disable(self): + """Disable the quantizer modules.""" + for quantizer in self: + quantizer.disable() + + def reset_amax(self): + """Reset amax of the quantizers.""" + for quantizer in self: + quantizer.reset_amax() + + def set_from_attribute_config( + self, + attributes: Union[ + list[Union[dict[str, Any], QuantizerAttributeConfig]], + Union[dict[str, Any], QuantizerAttributeConfig], + ], + ): + """Set the attributes of contained quantizers from a list of attribute_dicts.""" + if not isinstance(attributes, (list, tuple)): + assert isinstance( + attributes, (dict, QuantizerAttributeConfig) + ), "attributes must be a list or a dict." + attributes = [attributes] * len(self) + + for attribute, quantizer in zip(attributes, self): + quantizer.set_from_attribute_config(attribute) + + @staticmethod + @contextlib.contextmanager + def replace_sequential_quantizer_with_single_quantizer(model, indx: int = 0): + """Replace instances of :class:`SequentialQuantizer` in the model with single quantizers. + + The quantizer indexed by ``indx`` from the sequential quantizer is used to replace it. + This method is useful for individually calibrating the quantizers in a sequential quantizer. + """ + original_sequential_quantizers: dict[nn.Module, list] = {} + for name, module in list(model.named_modules()): + if isinstance(module, SequentialQuantizer): + assert len(module) > indx + parent_module = model.get_submodule(name.rpartition(".")[0]) + if parent_module not in original_sequential_quantizers: + original_sequential_quantizers[parent_module] = [] + original_sequential_quantizers[parent_module].append( + (name.rpartition(".")[-1], module) + ) + setattr(parent_module, name.rpartition(".")[-1], module[indx]) + + yield + + for parent_module, sequential_quantizers_list in original_sequential_quantizers.items(): + for name, sequential_quantizer in sequential_quantizers_list: + setattr(parent_module, name, sequential_quantizer) + + @staticmethod + def tensor_quantizer_iterator(quantizers): + """Iterator for the quantizers in the container (but yield itself if its a TensorQuantizer).""" + if quantizers is None: + return + if isinstance(quantizers, TensorQuantizer): + yield quantizers + elif isinstance(quantizers, SequentialQuantizer): + for quantizer in quantizers: + yield quantizer + else: + raise ValueError("Invalid quantizer type.") diff --git a/modelopt/torch/quantization/optim.py b/modelopt/torch/quantization/optim.py new file mode 100644 index 0000000000..36b4ca0d4d --- /dev/null +++ b/modelopt/torch/quantization/optim.py @@ -0,0 +1,38 @@ +# 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. + +"""Deprecated. Placeholder module for throwing deprecated error.""" + +from modelopt.torch.utils import DeprecatedError + + +def match_parameters(*args, **kwargs): + """Deprecated. Placeholder function for throwing deprecated error.""" + raise DeprecatedError("This method is deprecated. ") + + +def group_parameters(*args, **kwargs): + """Deprecated. Placeholder function for throwing deprecated error.""" + raise DeprecatedError("This method is deprecated. ") + + +def freeze_parameters(*args, **kwargs): + """Deprecated. Placeholder function for throwing deprecated error.""" + raise DeprecatedError("This module is deprecated. ") + + +def quant_weight_inplace(*args, **kwargs): + """Deprecated. Placeholder function for throwing deprecated error.""" + raise DeprecatedError("This method is deprecated. ") diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py new file mode 100644 index 0000000000..a57ce44c0d --- /dev/null +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -0,0 +1,51 @@ +# 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. + +"""Handles quantization plugins to correctly quantize third-party modules. + +Please check out the source code of this module for examples of how plugins work and how you can +write your own one. Currently, we support plugins for + +- :meth:`apex` +- :meth:`diffusers` +- :meth:`huggingface` +- :meth:`megatron` +- :meth:`nemo` +- :meth:`fairscale` +- :meth:`peft` +""" + +from modelopt.torch.utils import import_plugin + +with import_plugin("apex"): + from .apex import * + +with import_plugin("diffusers"): + from .diffusers import * + +with import_plugin("huggingface"): + from .huggingface import * + +with import_plugin("megatron"): + from .megatron import * + +with import_plugin("nemo"): + from .nemo import * + +with import_plugin("fairscale"): + from .fairscale import * + +with import_plugin("peft"): + from .peft import * diff --git a/modelopt/torch/quantization/plugins/apex.py b/modelopt/torch/quantization/plugins/apex.py new file mode 100644 index 0000000000..974fa3722e --- /dev/null +++ b/modelopt/torch/quantization/plugins/apex.py @@ -0,0 +1,54 @@ +# 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. + +"""Support quantization for apex linear layers.""" + +from functools import partial + +import apex.transformer.tensor_parallel.layers as apex_parallel +from apex.transformer.parallel_state import get_data_parallel_group, get_tensor_model_parallel_group + +from modelopt.torch.utils.distributed import ParallelState + +from ..nn import QuantModuleRegistry +from ..nn.modules.quant_linear import _QuantLinear +from .custom import _ParallelLinear + + +class _ApexParallelLinear(_ParallelLinear): + def initialize_parallel_state(self): + self._parallel_state = ParallelState( + get_data_parallel_group(), get_tensor_model_parallel_group() + ) + + def _setup(self): + quantized_linear_fn = partial( + _QuantLinear.quantized_linear_fn, + apex_parallel, + "linear_with_grad_accumulation_and_async_allreduce", + self, + ) + self._forward_impl = quantized_linear_fn + super()._setup() + + +@QuantModuleRegistry.register({apex_parallel.ColumnParallelLinear: "apex_ColumnParallelLinear"}) +class _ApexColumnParallelLinear(_ApexParallelLinear): + _is_column_parallel = True + + +@QuantModuleRegistry.register({apex_parallel.RowParallelLinear: "apex_RowParallelLinear"}) +class _ApexRowParallelLinear(_ApexParallelLinear): + _is_row_parallel = True diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py new file mode 100644 index 0000000000..d4eae3c017 --- /dev/null +++ b/modelopt/torch/quantization/plugins/attention.py @@ -0,0 +1,229 @@ +# Inspired by https://github.com/ELS-RD/transformer-deploy/blob/6b88e24ade6ce199e825adc0477b28a07f51f17d/src/transformer_deploy/QDQModels/ast_operator_patch.py + +# Apache License +# Copyright 2022, Lefebvre Dalloz Services + +# 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. + +# 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. + +"""Support quantization for KV Cache in attention layers.""" + +import ast +import inspect +import tempfile +import types + +from ..conversion import register +from ..nn import TensorQuantizer + +__all__ = ["register_attention_for_kv_quant"] + + +def register_attention_for_kv_quant(attention_cls: type) -> bool: + """Register attention layer for quantization of KV Cache. + + Generate a quantized version of the attention class on the fly, + and register it with the original class for quantization. + """ + source_code = inspect.getsource(attention_cls) + model_module = inspect.getmodule(attention_cls) + head = ast.parse(source_code) + + bmm_ops = ("matmul", "bmm", "baddbmm") + sdpa_ops = ("scaled_dot_product_attention",) + + def is_bmm(node): + return ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + # and isinstance(node.func.value, ast.Name) + # and node.func.value.id == "torch" + and node.func.attr in bmm_ops + ) + + def is_sdpa(node): + return ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in sdpa_ops + ) + + def is_bin_matmul(node): + return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult) + + def patch(node, quantizer_names, transpose=False): + for index, quantizer_name in enumerate(quantizer_names): + if quantizer_name is None: + continue + arg = node.args[index] + + if not transpose: + node.args[index] = ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_name, + ctx=ast.Load(), + ), + args=[arg], + keywords=[], + ) + else: + node.args[index] = ast.Call( + func=ast.Attribute( + value=ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_name, + ctx=ast.Load(), + ), + args=[ + ast.Call( + func=ast.Attribute(value=arg, attr="transpose", ctx=ast.Load()), + args=[ast.Constant(value=-1), ast.Constant(value=-2)], + keywords=[], + ) + ], + keywords=[], + ), + attr="transpose", + ctx=ast.Load(), + ), + args=[ast.Constant(value=-1), ast.Constant(value=-2)], + keywords=[], + ) + + def patch_binop(node, quantizer_names): + assert len(quantizer_names) == 2 + if quantizer_names[0] is not None: + node.left = ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_names[0], + ctx=ast.Load(), + ), + args=[node.left], + keywords=[], + ) + if quantizer_names[1] is not None: + node.right = ast.Call( + func=ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr=quantizer_names[1], + ctx=ast.Load(), + ), + args=[node.right], + keywords=[], + ) + + nodes = list(ast.walk(head)) + org_class_name = nodes[1].name # type: ignore[attr-defined] + new_class_name = nodes[1].name = "_Quant" + nodes[1].name # type: ignore[attr-defined] + + bmm_nodes = [] + sdpa_nodes = [] + bin_matmul_nodes = [] + for node in ast.walk(head): + if is_bmm(node): + bmm_nodes.append(node) + if is_sdpa(node): + sdpa_nodes.append(node) + if is_bin_matmul(node): + bin_matmul_nodes.append(node) + + if len(bmm_nodes) != 2 and len(sdpa_nodes) != 1 and len(bin_matmul_nodes) != 2: + print(f"Expect 2 bmm/matmul op in the {org_class_name}, found {len(bmm_nodes)}") + print(f"Or expect 1 sdpa op in the {org_class_name}, found {len(sdpa_nodes)}") + print(f"Or expect 2 @ op in the {org_class_name}, found {len(bin_matmul_nodes)}") + print("Auto quantization of KV Cache fails") + return False + + if len(bmm_nodes) == 2: + # transpose k cache here to enable per-token quantization + # without transpose, the quantization will be per-channel, i.e., + # self.k_bmm_quantizer(key_states.transpose(-1, -2)) + # after transpose, the quantization will be per-token, i.e., + # self.k_bmm_quantizer(key_states.transpose(-1, -2).transpose(-1, -2)).transpose(-1, -2) + # removing the additional tranpose is doable but not trivial + patch(bmm_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) + patch(bmm_nodes[1], quantizer_names=(None, "k_bmm_quantizer"), transpose=True) + print("Patching 2 BMM/Matmul operators with quantizers") + if len(bin_matmul_nodes) == 2: + patch_binop(bin_matmul_nodes[0], quantizer_names=(None, "k_bmm_quantizer")) + patch_binop(bin_matmul_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) + print("Patching 2 @ operators with quantizers") + + if len(sdpa_nodes) == 1: + patch(sdpa_nodes[0], quantizer_names=(None, "k_bmm_quantizer", "v_bmm_quantizer")) + print("Patching 1 scaled_dot_product_attention operator with quantizers") + + head = ast.fix_missing_locations(head) + org_class = model_module.__dict__[org_class_name] + + module_code_str = ast.unparse(head) + with tempfile.NamedTemporaryFile(prefix="modelopt_", suffix=".py", delete=False) as temp_file: + temp_file.write(module_code_str.encode()) + print(f"Definition of {new_class_name} saved to {temp_file.name}") + + # Exec with python runtime and extract the new class + # This could lead to side effects if the class code is not properly isolated + # Therefore, it is recommended to run this function only when necessary + # exec( + # new_class_code, + # globals=model_module.__dict__, + # locals=model_module.__dict__ + # ) # bandit throws error here + # quant_class = model_module.__dict__[new_class_name] + + # Extract the bytecode and create a new class on the fly + # This is more tricky but doesn't require runtime execution + module_code = compile(head, filename="modelopt_generated", mode="exec") + class_code = module_code.co_consts[0] + assert class_code.co_name == new_class_name + method_codes = [const for const in class_code.co_consts if isinstance(const, types.CodeType)] + + new_methods = {} + for method_code in method_codes: + method_name = method_code.co_name + original_method = getattr(org_class, method_name, None) + if not isinstance(original_method, types.FunctionType): + continue + # Create a new class method from bytecode + new_methods[method_name] = types.FunctionType( + method_code, globals=original_method.__globals__, closure=original_method.__closure__ + ) + + def setup_method(self): + self.k_bmm_quantizer = TensorQuantizer() + self.v_bmm_quantizer = TensorQuantizer() + + assert "_setup" not in new_methods, "Method _setup already exists" + new_methods["_setup"] = setup_method + + # Create a new subclass on the fly + quant_class = type(new_class_name, (org_class,), new_methods) + + register(original_cls=org_class, quantized_cls=quant_class) + print(f"Successfully registered {org_class_name} for quantization") + return True diff --git a/modelopt/torch/quantization/plugins/custom.py b/modelopt/torch/quantization/plugins/custom.py new file mode 100644 index 0000000000..9f1ef5955d --- /dev/null +++ b/modelopt/torch/quantization/plugins/custom.py @@ -0,0 +1,110 @@ +# 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. + +"""Custom plugin base modules and utilities for quantization.""" + +from functools import partial +from types import ModuleType +from typing import Callable, Iterator + +from modelopt.torch.opt.dynamic import DynamicModule +from modelopt.torch.utils.distributed import ParallelState + +from ..nn import TensorQuantizer +from ..nn.modules.quant_linear import _QuantLinear +from ..utils import multi_context, replace_function + +try: + from .huggingface import register_dbrx_moe_on_the_fly, register_falcon_linears_on_the_fly +except ImportError: + + def _dummy_register(model): + pass + + register_falcon_linears_on_the_fly = _dummy_register + register_dbrx_moe_on_the_fly = _dummy_register + + +# TODO: This is a temporary solution +# In future implement a decorator to register methods updating QUANT_MODULE on the fly +def register_custom_model_plugins_on_the_fly(model): + """Registers custom modules as QUANT_MODULE on the fly.""" + register_falcon_linears_on_the_fly(model) + register_dbrx_moe_on_the_fly(model) + + +class _QuantFunctionalMixin(DynamicModule): + """Mixin class for quantized functionals. + + Often we need to replace a functional with a quantized version. This class provides a way to do that. + """ + + # List of functionals to replace with quantized versions, e.g. [(package, func_name, quantized_func), ...] + _functionals_to_replace: list[tuple[ModuleType, str, Callable]] = [] + + @property + def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]: + return ( + (package, func_name, quantized_func) + for package, func_name, quantized_func in self._functionals_to_replace + if hasattr(package, func_name) + ) + + def forward(self, *args, **kwargs): + with multi_context( + *( + replace_function(package, func_name, quantized_func) + for package, func_name, quantized_func in self.functionals_to_replace + ) + ): + return super().forward(*args, **kwargs) + + +class _ParallelLinear(_QuantFunctionalMixin): + """Quantized base class for ParallelLinear type classes. + + For this type of modules, we need to quantize the inputs and weights just before calling the actual linear + functional. This is accomplished by replacing the linear functional with a custom one that quantizes the inputs + and weights before calling the original functional. + """ + + # List of functionals to replace [(package, func_name), ...] + _functionals_to_replace: list[tuple[ModuleType, str]] = [] + + _parallel_state: ParallelState + _is_column_parallel = False + _is_row_parallel = False + + @property + def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]: + for package, func_name in self._functionals_to_replace: + if not hasattr(package, func_name): + continue + quantized_func = partial( + _QuantLinear.quantized_linear_fn, package, "_" + func_name, self + ) + if hasattr(getattr(package, func_name), "__dict__"): + quantized_func.__dict__.update(getattr(package, func_name).__dict__) + yield package, func_name, quantized_func + + def initialize_parallel_state(self): + self._parallel_state = ParallelState() + + def _setup(self): + self.initialize_parallel_state() + self.input_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_input) + self.weight_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_weight) + self.output_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_output) + self.output_quantizer.disable() diff --git a/modelopt/torch/quantization/plugins/diffusers.py b/modelopt/torch/quantization/plugins/diffusers.py new file mode 100644 index 0000000000..4b75330df9 --- /dev/null +++ b/modelopt/torch/quantization/plugins/diffusers.py @@ -0,0 +1,231 @@ +# 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. + +"""Support quantization of diffusers layers.""" + +import functools +from functools import partial +from types import ModuleType +from typing import Callable, Iterator, Optional + +import onnx +import torch +from diffusers.models.attention_processor import Attention +from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear +from torch.autograd import Function +from torch.nn import functional as F +from torch.onnx import symbolic_helper +from torch.onnx._internal import jit_utils, registration + +from ..export_onnx import export_fp8_mha +from ..nn import ( + QuantConv2d, + QuantInputBase, + QuantLinear, + QuantLinearConvBase, + QuantModuleRegistry, + TensorQuantizer, +) +from .custom import _QuantFunctionalMixin + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=18) + +onnx_dtype_map = { + "BFloat16": onnx.TensorProto.BFLOAT16, + "Float": onnx.TensorProto.FLOAT, + "Float8": onnx.TensorProto.FLOAT8E4M3FN, + "Half": onnx.TensorProto.FLOAT16, + "INT8": onnx.TensorProto.INT8, + "UINT8": onnx.TensorProto.UINT8, +} +mha_valid_precisions = {"Half", "BFloat16"} + + +class _QuantLoRACompatibleLinearConvBase(QuantLinearConvBase): + def _setup(self): + assert self.lora_layer is None, ( + f"To quantize {self}, lora_layer should be None. Please fuse the LoRA layer before" + " quantization." + ) + return super()._setup() + + +@QuantModuleRegistry.register({LoRACompatibleConv: "LoRACompatibleConv"}) +class _QuantLoRACompatibleConv(_QuantLoRACompatibleLinearConvBase): + default_quant_desc_weight = QuantConv2d.default_quant_desc_weight + + +@QuantModuleRegistry.register({LoRACompatibleLinear: "LoRACompatibleLinear"}) +class _QuantLoRACompatibleLinear(_QuantLoRACompatibleLinearConvBase): + default_quant_desc_weight = QuantLinear.default_quant_desc_weight + + +def _quantized_bmm(self, input, mat2, *args, **kwargs): + attn, v = input, mat2 + return self.bmm2_output_quantizer( + torch._bmm(self.softmax_quantizer(attn), self.v_bmm_quantizer(v), *args, **kwargs) + ) + + +def _quantized_baddbmm(self, input, batch1, batch2, *args, **kwargs): + q, k = batch1, batch2 + return torch._baddbmm(input, self.q_bmm_quantizer(q), self.k_bmm_quantizer(k), *args, **kwargs) + + +def _quantized_sdpa(self, *args, **kwargs): + fp8_sdpa = FP8SDPA.apply + parameters = [ + "query", + "key", + "value", + "attn_mask", + "dropout_p", + "is_causal", + "scale", + "q_quantized_scale", + "k_quantized_scale", + "v_quantized_scale", + "high_precision_flag", + ] + default_values = [None, None, None, None, 0.0, False, None, None, None, None, "Half"] + param_dict = dict(zip(parameters, default_values)) + for i, arg in enumerate(args): + param_dict[parameters[i]] = arg + param_dict.update(kwargs) + fp8_sdpa_args = [param_dict[param] for param in parameters] + while fp8_sdpa_args and fp8_sdpa_args[-1] is None: + fp8_sdpa_args.pop() + query, key, value = fp8_sdpa_args[:3] + q_quantized_scale = self.q_bmm_quantizer._get_amax(query) + k_quantized_scale = self.k_bmm_quantizer._get_amax(key) + v_quantized_scale = self.v_bmm_quantizer._get_amax(value) + + # We don't need to calibrate the output of softmax + return self.bmm2_output_quantizer( + fp8_sdpa( + query, + key, + value, + *fp8_sdpa_args[3:7], + q_quantized_scale, + k_quantized_scale, + v_quantized_scale, + self.q_bmm_quantizer.trt_high_precision_dtype + if hasattr(self.q_bmm_quantizer, "trt_high_precision_dtype") + else "Half", + self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True, + ) + ) + + +class _QuantAttention(_QuantFunctionalMixin): + """FP8 processor for performing attention-related computations.""" + + _functionals_to_replace = [ + (torch, "bmm", _quantized_bmm), + (torch, "baddbmm", _quantized_baddbmm), + (F, "scaled_dot_product_attention", _quantized_sdpa), + ] + + @property + def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]: + for package, func_name, quantized_func in self._functionals_to_replace: + if not hasattr(package, func_name): + continue + quantized_func = partial(quantized_func, self) + yield package, func_name, quantized_func + + def _setup(self): + self.q_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input) + self.k_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input) + self.v_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input) + self.softmax_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input) + self.bmm2_output_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input) + + +QuantModuleRegistry.register({Attention: "Attention"})(_QuantAttention) + + +original_scaled_dot_product_attention = F.scaled_dot_product_attention + + +class FP8SDPA(Function): + """A customized FP8 SDPA op for the onnx export.""" + + @staticmethod + def forward( + ctx, + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None, + q_quantized_scale=None, + k_quantized_scale=None, + v_quantized_scale=None, + high_precision_flag=None, + disable_fp8_mha=True, + ): + """Forward method.""" + ctx.save_for_backward(query, key, value, attn_mask) + ctx.q_quantized_scale = q_quantized_scale + ctx.k_quantized_scale = k_quantized_scale + ctx.v_quantized_scale = v_quantized_scale + # During runtime, ignore x or use it as needed + return original_scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + ) + + @staticmethod + @symbolic_helper.parse_args("v", "v", "v", "v", "f", "b", "v", "t", "t", "t", "s", "b") + def symbolic( + g: jit_utils.GraphContext, + query: torch._C.Value, + key: torch._C.Value, + value: torch._C.Value, + attn_mask: Optional[torch._C.Value] = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: Optional[torch._C.Value] = None, + q_quantized_scale: float = 1.0, + k_quantized_scale: float = 1.0, + v_quantized_scale: float = 1.0, + high_precision_flag: str = "Half", + disable_fp8_mha: bool = True, + ): + """Symbolic method.""" + return export_fp8_mha( + g, + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + q_quantized_scale, + k_quantized_scale, + v_quantized_scale, + high_precision_flag, + disable_fp8_mha, + ) diff --git a/modelopt/torch/quantization/plugins/fairscale.py b/modelopt/torch/quantization/plugins/fairscale.py new file mode 100644 index 0000000000..9544edab32 --- /dev/null +++ b/modelopt/torch/quantization/plugins/fairscale.py @@ -0,0 +1,44 @@ +# 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. + +"""Support quantization for megatron linear layers.""" + +import torch.nn.functional as F +from fairscale.nn.model_parallel.initialize import get_data_parallel_group, get_model_parallel_group +from fairscale.nn.model_parallel.layers import ColumnParallelLinear, RowParallelLinear + +from modelopt.torch.utils.distributed import ParallelState + +from ..nn import QuantModuleRegistry +from .custom import _ParallelLinear + +__all__ = [] + + +class _FairscaleParallelLinear(_ParallelLinear): + _functionals_to_replace = [(F, "linear")] + + def initialize_parallel_state(self): + self._parallel_state = ParallelState(get_data_parallel_group(), get_model_parallel_group()) + + +@QuantModuleRegistry.register({ColumnParallelLinear: "fairscale_ColumnParallelLinear"}) +class _FairscaleColumnParallelLinear(_FairscaleParallelLinear): + _is_column_parallel = True + + +@QuantModuleRegistry.register({RowParallelLinear: "fairscale_RowParallelLinear"}) +class _FairscaleRowParallelLinear(_FairscaleParallelLinear): + _is_row_parallel = True diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py new file mode 100644 index 0000000000..2735c51d71 --- /dev/null +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -0,0 +1,251 @@ +# 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. + +"""Support quantization for huggingface layers.""" + +import warnings +from contextlib import contextmanager + +import torch +import torch.nn as nn +import transformers + +from modelopt.core.torch.quantization.algorithms import AutoQuantizeSearcher +from modelopt.torch.opt.dynamic import DynamicModule + +from ..nn import QuantModuleRegistry +from ..nn.modules.quant_linear import _QuantLinear +from .attention import register_attention_for_kv_quant + +__all__ = ["register_hf_attentions_on_the_fly"] + +if transformers.modeling_utils.Conv1D not in QuantModuleRegistry: + # transformers.modeling_utils.Conv1D used in HF-GPT2 is not a real Conv1D + # It is actually a Linear layer where weight is transposed and torch.addmm is used + @QuantModuleRegistry.register({transformers.modeling_utils.Conv1D: "Conv1D"}) + class _QuantConv1D(_QuantLinear): + @classmethod + @torch.no_grad() + def convert(cls, module: nn.Module) -> "_QuantConv1D": + module.weight = nn.Parameter(module.weight.T.contiguous()) + module.out_features, module.in_features = module.weight.shape + # We want the forward method of nn.Linear to be called instead of the forward method of Conv1D + dyn_cls: DynamicModule = QuantModuleRegistry.get(nn.Linear) + return dyn_cls.convert(module) + + +if hasattr(transformers.models, "dbrx"): + # For more information on DbrxExpert, see https://github.com/huggingface/transformers/blame/dcdda5324bcc7a750b5e40e11dd795442204ff27/src/transformers/models/dbrx/modeling_dbrx.py#L756 + class _QuantDbrxExperts(DynamicModule): + def _setup(self): + """Modify the DbrxExpert.""" + # No setup is needed for DbrxExpert, we only need to update DbrxExpertGLU + pass + + # forward method copied from the original dbrx repo - https://github.com/databricks/dbrx/blob/a3200393e678387a6f30f3e903108c650625eb21/model/modeling_dbrx.py#L795 + def forward( + self, + x: torch.Tensor, + weights: torch.Tensor, + top_weights: torch.Tensor, + top_experts: torch.LongTensor, + ) -> torch.Tensor: + bsz, q_len, hidden_size = x.shape + x = x.view(-1, hidden_size) + out = torch.zeros_like(x) + + expert_mask = nn.functional.one_hot( + top_experts, num_classes=self.moe_num_experts + ).permute(2, 1, 0) + for expert_idx in range(0, self.moe_num_experts): + topk_idx, token_idx = torch.where(expert_mask[expert_idx]) + if token_idx.shape[0] == 0: + continue + + token_list = token_idx.tolist() + topk_list = topk_idx.tolist() + + expert_tokens = x[None, token_list].reshape(-1, hidden_size) + expert_out = ( + self.mlp(expert_tokens, expert_idx) * top_weights[token_list, topk_list, None] + ) + + out.index_add_(0, token_idx, expert_out) + + out = out.reshape(bsz, q_len, hidden_size) + return out + + class _QuantDbrxExpertGLU(DynamicModule): + def _setup(self): + """Modify the DbrxExpertGLU by using nn.Linear layers.""" + dtype, device = self.w1.dtype, self.w1.device + + def _copy_weights(modules, weights): + modules.to(dtype=dtype, device=device) + for expert_idx, module in enumerate(modules): + with torch.no_grad(): + module.weight.copy_(weights[expert_idx].detach()) + + self.w1_linear = nn.ModuleList( + [ + nn.Linear(self.hidden_size, self.ffn_hidden_size, bias=False) + for _ in range(self.moe_num_experts) + ] + ) + _copy_weights( + self.w1_linear, + self.w1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size), + ) + delattr(self, "w1") + + self.v1_linear = nn.ModuleList( + [ + nn.Linear(self.hidden_size, self.ffn_hidden_size, bias=False) + for _ in range(self.moe_num_experts) + ] + ) + _copy_weights( + self.v1_linear, + self.v1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size), + ) + delattr(self, "v1") + + self.w2_linear = nn.ModuleList( + [ + nn.Linear(self.ffn_hidden_size, self.hidden_size, bias=False) + for _ in range(self.moe_num_experts) + ] + ) + _copy_weights( + self.w2_linear, + self.w2.view( + self.moe_num_experts, self.ffn_hidden_size, self.hidden_size + ).transpose(1, 2), + ) + delattr(self, "w2") + + def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor: + x1 = self.w1_linear[expert_idx](x) + x2 = self.v1_linear[expert_idx](x) + x1 = self.activation_fn(x1) + x1 = x1 * x2 + return self.w2_linear[expert_idx](x1) + + if transformers.models.dbrx.modeling_dbrx.DbrxExperts not in QuantModuleRegistry: + QuantModuleRegistry.register( + {transformers.models.dbrx.modeling_dbrx.DbrxExperts: "hf.DbrxExperts"} + )(_QuantDbrxExperts) + + if transformers.models.dbrx.modeling_dbrx.DbrxExpertGLU not in QuantModuleRegistry: + QuantModuleRegistry.register( + {transformers.models.dbrx.modeling_dbrx.DbrxExpertGLU: "hf.DbrxExpertGLU"} + )(_QuantDbrxExpertGLU) + + +if hasattr(transformers.models, "falcon") and hasattr( + transformers.models.falcon.modeling_falcon, "FalconLinear" +): + QuantModuleRegistry.register( + {transformers.models.falcon.modeling_falcon.FalconLinear: "FalconLinear"} + )(_QuantLinear) + + +def register_dbrx_moe_on_the_fly(model): + """Register DBRX MoE modules as QUANT_MODULE. + + The MoE class in DBRX is `transformers_modules.modeling_dbrx.DbrxExpertGLU`, which loads dynamically. + """ + if type(model).__name__ in ["DbrxForCausalLM"]: + moe_type = type(model.transformer.blocks[0].ffn.experts.mlp) + # Create a QuantDbrxExpertGLU class on the fly + if QuantModuleRegistry.get(moe_type) is None: + QuantModuleRegistry.register({moe_type: moe_type.__name__})(_QuantDbrxExpertGLU) + + +def register_falcon_linears_on_the_fly(model): + """Register Falcon linear modules as a QUANT_MODULE. + + Certain falcon models (for example, falcon 40b) use remote code, which are loaded dynamically, to build their model. + Therefore, we need to register the linear on the fly before quantization. + """ + if type(model).__name__ in ["RWForCausalLM", "FalconForCausalLM"]: + linear_type = type(model.transformer.h[0].self_attention.dense) + # Create a QuantFalconLinear class on the fly + if QuantModuleRegistry.get(linear_type) is None: + QuantModuleRegistry.register({linear_type: linear_type.__name__})(_QuantLinear) + + +def register_hf_attentions_on_the_fly(model): + """Find HF Attention modules in the model and register them for KV Cache quantization. + + This function attempts to find child modules ending with "Attention" in the name. + If such child modules are not found, or the corresponding class does not contain + identifiable attention patterns, the function will not register any new modules. + """ + attention_cls = {} + for name, module in model.named_modules(): + if type(module).__name__.endswith("Attention"): + attention_cls[type(module)] = type(module).__name__ + + success = any([register_attention_for_kv_quant(cls) for cls in attention_cls]) + if not success: + raise RuntimeError("No Attention class found for KV Cache quantization.") + + +def _is_supported_hf_model(model): + """Check if the model a valid model for transformers quantization specific support.""" + return isinstance(model, transformers.PreTrainedModel) + + +@contextmanager +def setup_model_for_gradient_checkpointing(model: nn.Module): + use_cache = None + if hasattr(model, "config") and hasattr(model.config, "use_cache"): + # Disable use_cache explicitly before forward is called + use_cache = model.config.use_cache + model.config.use_cache = False + + if not hasattr(model, "gradient_checkpointing_enable") or not ( + hasattr(model, "supports_gradient_checkpointing") and model.supports_gradient_checkpointing + ): + warnings.warn( + "AutoQuantize: Huggingface model without gradient checkpointing support detected. " + "AutoQuantize will consume more memory." + ) + else: + try: + warnings.warn( + "AutoQuantize: Huggingface model detected - Enabling gradient checkpointing. " + "Disable gradient checkpointing after AutoQuantize if this is not desired!" + ) + model.gradient_checkpointing_enable({"use_reentrant": True}) + model.train() # Model needs to be in training mode to enable gradient checkpointing + # Set all dropout layers to eval mode for deterministic auto-quantize scores + for name, module in model.named_modules(): + if isinstance(model, torch.nn.Dropout): + module.eval() + except Exception as e: + warnings.warn( + f"AutoQuantize: Error enabling gradient checkpointing for huggingface model due to: {e}, " + "AutoQuantize will consume more memory." + ) + yield + if use_cache is not None: + model.config.use_cache = use_cache + + +AutoQuantizeSearcher.register_gradient_checkpointing_enable_context( + _is_supported_hf_model, setup_model_for_gradient_checkpointing +) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py new file mode 100644 index 0000000000..100b1c8b85 --- /dev/null +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -0,0 +1,298 @@ +# 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. + +"""Support quantization for megatron linear layers.""" + +import math +import warnings +from contextlib import contextmanager + +import megatron.core.tensor_parallel.layers as megatron_parallel +import torch +import torch.nn.functional as F +from megatron.core.parallel_state import get_data_parallel_group, get_tensor_model_parallel_group +from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint +from packaging.version import Version + +from modelopt import __version__ +from modelopt.torch.utils.distributed import ParallelState + +from ..nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer +from .custom import _ParallelLinear + +__all__ = [] + + +class _MegatronParallelLinear(_ParallelLinear): + _functionals_to_replace = [ + (megatron_parallel, "linear_with_grad_accumulation_and_async_allreduce"), + (megatron_parallel, "linear_with_frozen_weight"), + ] + + _QUANT_MIN_BLOCK_SIZE = 128 + + def initialize_parallel_state(self): + self._parallel_state = ParallelState( + get_data_parallel_group(), get_tensor_model_parallel_group() + ) + + def _process_weight_quantizer_amax(self, k, v, quantizer_state_dict): + quantizer_state_dict[k] = v.view(self.weight.shape[0], -1) if v.numel() > 1 else v.view(-1) + + def _process_activation_quantizer_amax(self, k, v, quantizer_state_dict): + # if quantizer is not enabled, we need to honor the input dtype for consistency + # TODO: This is a temporary solution and needs to be removed once megatron supports + # non-homogeneous layers + def _sanitize_amax_dtype_for_disabled_quantizer(k, v): + quantizer_name = k.split(".")[0] + if hasattr(self, quantizer_name) and not getattr(self, quantizer_name).is_enabled: + input_dtype = getattr(self, k.split(".")[0])._input_dtype + v = v.to(input_dtype) + return v + + quantizer_state_dict[k] = _sanitize_amax_dtype_for_disabled_quantizer(k, v).view(-1) + + def _process_activation_quantizer_pre_quant_scale(self, k, v, quantizer_state_dict): + quantizer_state_dict[k] = v + + def _get_shard_axis_dict(self): + raise NotImplementedError + + @contextmanager + def _get_quantizer_states_for_sharded_state_dict(self): + """Get the quantizer states for sharded state_dict. + + This is a workaround to support MCore distributed checkpointing. MCore distributed checkpointing + requires all the transformer blocks to be uniform in terms the parameters and buffers. + + However with partial quantization/mixed precision quantization, the quantizers of Linear layers are not uniform + across the transformer blocks. For example, linear layer from a particular block might have INT4 quantization + with `weight_quantizer.amax` of shape (Cout x Cin//block_size) while the linear layer from another block might + be disabled and `weight_quantizer.amax` could be None. + + MCore uses `shared_state_dict` to get the state_dict of the Linear layers. + So lets call `sharded_state_dict` from this context. We will temporarily insert dummy states for the quantizers + with the maximum size. We will store the original quantizer state if it exists in the dummy states. + After the `sharded_state_dict` is called, we will restore the original quantizer states. + This will ensure that the quantizer states `shared_state_dict` are uniform across the entire model. + """ + if self.adapt_for_old_ckpt: + yield + return + + original_quantizer_states = {} + + _is_original_weight_quantizer_sequential = isinstance( + self.weight_quantizer, SequentialQuantizer + ) + if not isinstance(self.weight_quantizer, SequentialQuantizer): + self.weight_quantizer = SequentialQuantizer(self.weight_quantizer, TensorQuantizer()) + else: + assert ( + len(self.weight_quantizer) == 2 + ), f"Invalid weight_quantizer: {self.weight_quantizer}" + + for name, module in self.named_modules(): + # lm_head or output_layer has no weight (reusing the embeddings) + if self.weight is None: + continue + if not isinstance(module, TensorQuantizer): + continue + + original_quantizer_states[module] = { + "_amax": getattr(module, "_amax", None), + "_pre_quant_scale": getattr(module, "_pre_quant_scale", None), + } + if "weight_quantizer" in name: + weight_shape = self.weight.shape + dummy_tensor_shape = ( + weight_shape[0], + # if dim % block_size != 0, then we will pad the dim, so we need ceil here + math.ceil(weight_shape[1] / self._QUANT_MIN_BLOCK_SIZE), + ) + if hasattr(module, "_amax"): + assert module._amax.ndim in [0, 2], "Invalid amax" + if module._amax.ndim == 2: + dummy_tensor = module._amax.view(weight_shape[0], -1) + if dummy_tensor.shape[-1] == 1: + # Per-Channel quantization + dummy_tensor = module._amax.repeat(1, dummy_tensor_shape[-1]) + else: + # Per-block quantization + cur_amax_dim = dummy_tensor.shape[-1] + # if block_size > 128, padding the amax to max possible shape + if cur_amax_dim < dummy_tensor_shape[-1]: + dummy_tensor = F.pad( + dummy_tensor, (0, dummy_tensor_shape[1] - cur_amax_dim) + ) + elif cur_amax_dim > dummy_tensor_shape[-1]: + raise ValueError( + "Expecting the block_size >= 128 for static block-wise quantization!" + ) + elif module._amax.ndim == 0: + # Per-tensor quantization + dummy_tensor = ( + torch.ones( + dummy_tensor_shape, + device=module._amax.device, + dtype=module._amax.dtype, + ) + * module._amax + ) + else: + raise ValueError("Invalid amax") + delattr(module, "_amax") + else: + dummy_tensor = torch.zeros( + dummy_tensor_shape, device=self.weight.device, dtype=self.weight.dtype + ) + else: + if hasattr(module, "_amax"): + assert module._amax.numel() == 1, "Invalid amax" + dummy_tensor = module._amax.view(-1) + delattr(module, "_amax") + else: + dummy_tensor = torch.zeros( + 1, device=self.weight.device, dtype=self.weight.dtype + ) + module.register_buffer("_amax", dummy_tensor) + if "input_quantizer" in name: + if hasattr(module, "_pre_quant_scale"): + dummy_tensor = module._pre_quant_scale.view(-1) + delattr(module, "_pre_quant_scale") + else: + dummy_tensor = torch.zeros( + self.weight.shape[1], device=self.weight.device, dtype=self.weight.dtype + ) + module.register_buffer("_pre_quant_scale", dummy_tensor) + + yield + + for module, original_state in original_quantizer_states.items(): + for k, v in original_state.items(): + if hasattr(module, k): + delattr(module, k) + if v is not None: + module.register_buffer(k, v) + + if not _is_original_weight_quantizer_sequential: + self.weight_quantizer = self.weight_quantizer[0] + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + with self._get_quantizer_states_for_sharded_state_dict(): + sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets) + # TODO: Clean this up; We dont need all the custom handling here + quantizer_state_dict, sharded_axis_dict = {}, self._get_shard_axis_dict() + for k, v in self.state_dict(prefix="", keep_vars=True).items(): + if "weight_quantizer" in k and "_amax" in k: + self._process_weight_quantizer_amax(k, v, quantizer_state_dict) + elif ("input_quantizer" in k or "output_quantizer" in k) and k.endswith("._amax"): + self._process_activation_quantizer_amax(k, v, quantizer_state_dict) + elif k == "input_quantizer._pre_quant_scale": + self._process_activation_quantizer_pre_quant_scale(k, v, quantizer_state_dict) + + sharded_state_dict.update( + **make_sharded_tensors_for_checkpoint( + quantizer_state_dict, prefix, sharded_axis_dict, sharded_offsets + ) + ) + return sharded_state_dict + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + if ( + not isinstance(self.weight_quantizer, SequentialQuantizer) + and (prefix + "weight_quantizer.0._amax") in state_dict + ): + # state_dict generated from sharded_state_dict which temporarily SequentialQuantizer + state_dict[prefix + "weight_quantizer._amax"] = state_dict.pop( + prefix + "weight_quantizer.0._amax" + ) + state_dict.pop(prefix + "weight_quantizer.1._amax") + + for k in list(state_dict.keys()): + if not any( + k.startswith(prefix + name) + for name in ["weight_quantizer", "input_quantizer", "output_quantizer"] + ): + continue + + name = k.split(prefix)[-1] + if name not in self.state_dict(): + state_dict.pop(k) + continue + + if "weight_quantizer" in name and self.state_dict()[name].ndim == 2: + num_cols_to_crop = self.state_dict()[name].view(self.weight.shape[0], -1).shape[-1] + state_dict[k] = state_dict[k][..., :num_cols_to_crop].reshape(-1, 1) + elif "_amax" in name: + state_dict[k] = state_dict[k].view(-1)[0].view_as(self.state_dict()[name]) + + super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + def _setup(self): + super()._setup() + self._modelopt_load_version = __version__ # Will get overwritten by `replace_quant_module` + + @property + def adapt_for_old_ckpt(self) -> bool: + if ( + self._modelopt_load_version + and Version(self._modelopt_load_version) < Version("0.20") + and Version(self._modelopt_load_version) != Version("0.0.0") + ): # version in NeMo container is 0.0.0 if installed from source without history + warnings.warn("Old checkpoint detected. Please re-save model to avoid this warning.") + return True + return False + + +@QuantModuleRegistry.register( + {megatron_parallel.ColumnParallelLinear: "megatron_ColumnParallelLinear"} +) +class _MegatronColumnParallelLinear(_MegatronParallelLinear): + _is_column_parallel = True + + def _get_shard_axis_dict(self): + shard_axis_dict = {} + for k, v in self.state_dict().items(): + if "weight_quantizer" in k and "_amax" in k: + shard_axis_dict[k] = 0 + return shard_axis_dict + + +@QuantModuleRegistry.register({megatron_parallel.RowParallelLinear: "megatron_RowParallelLinear"}) +class _MegatronRowParallelLinear(_MegatronParallelLinear): + _is_row_parallel = True + + def _get_shard_axis_dict(self): + if self.adapt_for_old_ckpt: + shard_weight_for_block_wise_quant_only = True + else: + shard_weight_for_block_wise_quant_only = False + + shard_axis_dict = {} + for k, v in self.state_dict().items(): + if "weight_quantizer" in k and "_amax" in k: + if shard_weight_for_block_wise_quant_only: + quantizer = self.get_submodule(k.split("._amax")[0]) + if ( + quantizer.block_sizes + and quantizer.block_sizes.get("type", None) != "dynamic" + ): + shard_axis_dict[k] = 1 + else: + shard_axis_dict[k] = 1 + if k == "input_quantizer._pre_quant_scale": + shard_axis_dict[k] = 0 + return shard_axis_dict diff --git a/modelopt/torch/quantization/plugins/nemo.py b/modelopt/torch/quantization/plugins/nemo.py new file mode 100644 index 0000000000..6b2f9636e6 --- /dev/null +++ b/modelopt/torch/quantization/plugins/nemo.py @@ -0,0 +1,67 @@ +# 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. + +"""Customization for Nemo Megatron GPT.""" + +from functools import partial +from types import ModuleType +from typing import Callable, Iterator + +import torch + +# New nemo version should depend on megatron.core +from nemo.collections.nlp.modules.common.megatron.attention import CoreAttention + +from ..nn import QuantInputBase, QuantModuleRegistry, TensorQuantizer +from .custom import _QuantFunctionalMixin + +__all__ = [] + + +def _quantized_bmm(self, input, mat2, *args, **kwargs): + """Quantized version of BMM2 in nemo CoreAttention.""" + attn, v = input, mat2 + return torch._bmm(attn, self.v_bmm_quantizer(v), *args, **kwargs) + + +def _quantized_baddbmm(self, input, batch1, batch2, *args, **kwargs): + """Quantized version of BMM1 in nemo CoreAttention.""" + q, k = batch1, batch2 + return torch._baddbmm(input, self.q_bmm_quantizer(q), self.k_bmm_quantizer(k), *args, **kwargs) + + +class _QuantCoreAttention(_QuantFunctionalMixin): + """Quantized base class for CoreAttention.""" + + _functionals_to_replace = [ + (torch, "bmm", _quantized_bmm), + (torch, "baddbmm", _quantized_baddbmm), + ] + + @property + def functionals_to_replace(self) -> Iterator[tuple[ModuleType, str, Callable]]: + for package, func_name, quantized_func in self._functionals_to_replace: + if not hasattr(package, func_name): + continue + quantized_func = partial(quantized_func, self) + yield package, func_name, quantized_func + + def _setup(self): + self.q_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input) + self.k_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input) + self.v_bmm_quantizer = TensorQuantizer(QuantInputBase.default_quant_desc_input) + + +QuantModuleRegistry.register({CoreAttention: "nemo_core_attention"})(_QuantCoreAttention) diff --git a/modelopt/torch/quantization/plugins/peft.py b/modelopt/torch/quantization/plugins/peft.py new file mode 100644 index 0000000000..b7ef5b4ee3 --- /dev/null +++ b/modelopt/torch/quantization/plugins/peft.py @@ -0,0 +1,56 @@ +# 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. + +"""Support quantization for peft LoRA linear layers.""" + +import torch.nn.functional as F +from peft.tuners.lora.layer import Linear as LoraLinear + +from modelopt.torch.opt.dynamic import DynamicModule + +from ..nn import QuantModuleRegistry, TensorQuantizer + +__all__ = [] + + +@QuantModuleRegistry.register({LoraLinear: "LoraLinear"}) +class _QuantLoraLinear(DynamicModule): + def _setup(self): + self.input_quantizer = TensorQuantizer() + self.weight_quantizer = TensorQuantizer() + self.output_quantizer = TensorQuantizer() + + def forward(self, x, *args, **kwargs): + adapter_names = kwargs.pop("adapter_names", None) + if self.disable_adapters or adapter_names is not None or self.merged: + return super().forward(x, args, kwargs) + + weight = self.base_layer.weight + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + lora_a = self.lora_A[active_adapter] + lora_b = self.lora_B[active_adapter] + scaling = self.scaling[active_adapter] + + if not self.use_dora[active_adapter]: + weight = weight + scaling * lora_b.weight @ lora_a.weight + else: + raise NotImplementedError("dora not implemented") + + x = self.input_quantizer(x) + weight = self.weight_quantizer(weight) + output = self.output_quantizer(F.linear(x, weight, self.base_layer.bias)) + return output diff --git a/modelopt/torch/quantization/qtensor/__init__.py b/modelopt/torch/quantization/qtensor/__init__.py new file mode 100644 index 0000000000..5e14143762 --- /dev/null +++ b/modelopt/torch/quantization/qtensor/__init__.py @@ -0,0 +1,22 @@ +# 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. + +"""Tensor Class for Real Quantization.""" + +from modelopt.core.torch.quantization.qtensor.nvfp4_tensor import * + +from .base_qtensor import * +from .int4_tensor import * +from .nf4_tensor import * diff --git a/modelopt/torch/quantization/qtensor/base_qtensor.py b/modelopt/torch/quantization/qtensor/base_qtensor.py new file mode 100644 index 0000000000..63947d0d24 --- /dev/null +++ b/modelopt/torch/quantization/qtensor/base_qtensor.py @@ -0,0 +1,162 @@ +# 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. + +"""Base Class for Real Quantized Tensor.""" + +import enum + +import torch + + +class QTensorType(enum.Enum): + """Enumeration for defining types of quantization.""" + + INT4 = 1 + INT8 = 2 + FP8 = 3 + NF4 = 4 + + +__all__ = ["BaseQuantizedTensor", "QTensorWrapper", "pack_real_quantize_weight"] + + +class BaseQuantizedTensor: + """Base class for quantized tensors, providing methods for quantization and dequantization. + + This class should be subclassed to implement specific types of quantized tensors. It handles the + storage of quantized data along with the necessary configurations and original attributes. + + Attributes: + original_meta_tensor (torch.Tensor): Original meta to keep attributes of original tensors. + quantized_data (torch.Tensor): Storage for the quantized tensor data. Quantized_data dtype is + customized per QuantizedTensor implementation. + """ + + _quantized_data: torch.Tensor + + def __init__( + self, + original_shape: torch.Size, + original_dtype: torch.dtype, + quantized_data: torch.Tensor, + ): + """Initialize data attributes.""" + self.metadata = { + "shape": original_shape, + "dtype": original_dtype, + } + self._quantized_data = quantized_data + + @classmethod + def quantize(cls, input: torch.Tensor, block_size: int): + """Pack a fake torch.Tensor into a real quantized tensor. + + Args: + fake_quant_tensor (torch.Tensor): The fake quantized tensor. + + Returns: + A real quantized tensor, scales. + """ + raise NotImplementedError("This method must be implemented by subclasses.") + + def dequantize(self, dtype=None, **kwarg): + """Converts the quantized tensor back to a standard torch.Tensor. + + Returns: + torch.Tensor: The dequantized tensor. + """ + raise NotImplementedError("This method must be implemented by subclasses.") + + +class QTensorWrapper(torch.nn.Parameter): + """A wrapper class for quantized tensors to make them compatible with torch.nn.Parameter. + + Args: + qtensor (BaseQuantizedTensor): The quantized tensor to be wrapped. + """ + + def __new__(cls, qtensor: BaseQuantizedTensor): + """Create a new QTensorWrapper instance.""" + quantized_tensor = qtensor._quantized_data.view(torch.uint8) + instance = super().__new__(cls, quantized_tensor, requires_grad=False) + instance.metadata = qtensor.metadata + instance.metadata["qtensor_class"] = qtensor.__class__ + return instance + + def dim(self): + """Return the number of dimensions of the meta_tensor.""" + return len(self.metadata["shape"]) + + def to(self, *args, **kwargs): + """Override the `to` method to move real quantized tensors to the specified device.""" + changing_device, changing_dtype, *_ = torch._C._nn._parse_to(*args, **kwargs) + if changing_device: + self.data = self.data.to(device=changing_device) + dtype = changing_dtype if changing_dtype else self.metadata["dtype"] + return QTensorWrapper( + self.metadata["qtensor_class"](self.metadata["shape"], dtype, self.data) + ) + + def get_qtensor(self): + """Get the quantized tensor class from QTensorWrapper.""" + return self.metadata["qtensor_class"]( + self.metadata["shape"], self.metadata["dtype"], self.data + ) + + +def pack_real_quantize_weight(module, force_quantize: bool = False): + """Pack real quantized tensors to a compressed format and set proper load_state_dict function.""" + + # Function to dynamically override load_state_dict + def dynamically_update_state_methods(module): + # Original method + original_load_from_state_dict = module._load_from_state_dict + + def custom_load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Override _load_from_state_dict to handle custom parameters dynamically.""" + args_list = list(args) + + deleted_stat_dict = {} + for name, param in self.named_parameters(): + if isinstance(param, QTensorWrapper) and prefix + name in state_dict: + param.copy_(state_dict[prefix + name]) + deleted_stat_dict[prefix + name] = state_dict[prefix + name] + del state_dict[prefix + name] + # Set strict=False because weight keys are removed + kwargs = {} + if len(args_list) > 3: + args_list[1] = False + else: + kwargs = {"strict": False} + original_load_from_state_dict(state_dict, prefix, *args_list, **kwargs) + state_dict.update(**deleted_stat_dict) + + module._load_from_state_dict = custom_load_from_state_dict.__get__(module, type(module)) + + with torch.no_grad(): + for _, m in module.named_modules(): + if ( + hasattr(m, "weight_quantizer") + and m.weight_quantizer.is_enabled + and not m.weight_quantizer._fake_quant + ): + if force_quantize: + m.weight_quantizer._dequantize = False + assert ( + not m.weight.is_meta + ), "Real quantization does not support tensors on meta device." + real_quant_tensor = m.weight_quantizer(m.weight) + m.weight = QTensorWrapper(real_quant_tensor) + dynamically_update_state_methods(m) diff --git a/modelopt/torch/quantization/qtensor/int4_tensor.py b/modelopt/torch/quantization/qtensor/int4_tensor.py new file mode 100644 index 0000000000..6ebd7728e1 --- /dev/null +++ b/modelopt/torch/quantization/qtensor/int4_tensor.py @@ -0,0 +1,116 @@ +# 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. + +"""Implements INT4 quantization for efficient tensor storage and computation.""" + +import torch + +from ..extensions import get_cuda_ext +from ..qtensor.base_qtensor import BaseQuantizedTensor +from ..utils import reduce_amax + +__all__ = ["INT4QTensor"] + + +class INT4QTensor(BaseQuantizedTensor): + """Implements the INT4 quantization on tensors for more efficient storage or computation. + + Attributes: + quantized_data (torch.Tensor): The quantized data stored as a packed uint8 tensor. + """ + + @staticmethod + def _get_quant_maxbound(num_bits): + return 2 ** (num_bits - 1) - 1 + + @classmethod + def quantize(cls, input: torch.Tensor, block_size: int) -> torch.Tensor: + """Converting a tensor to a quantized format based on INT4 (AWQ) quantization. + + Args: + input (torch.Tensor): The input tensor to be quantized. + block_size (int): The size of each block for quantization. + + Returns: + tuple: Contains quantized data, input quantization config, and scale quantization config. + """ + cuda_ext = get_cuda_ext() + + scale_quant_maxbound = cls._get_quant_maxbound(num_bits=4) + + # get scales for each block + assert ( + input.numel() % block_size == 0 + ), "Number of input elements is not divisible by the block size." + block_input = input.view(-1, block_size) + scales = scale_quant_maxbound / reduce_amax(block_input, -1) + # expand scalers to match shape of input + scales = scales.view(block_input.shape[0], -1) # shape: (block_input.shape[0], 1) + + if cuda_ext and input.is_cuda: + # use a custom cuda kernel if available + packed_output_uint8 = cuda_ext.INT4_quantize(input, scales, block_size) + else: + scaled_blocks = block_input * scales + flattened = scaled_blocks.flatten() + # uint4: 0 - 15 + flattened = flattened.round().clamp( + -(scale_quant_maxbound + 1), scale_quant_maxbound + ) + (scale_quant_maxbound + 1) + flattened = flattened.to(torch.uint8) + + packed_output_uint8 = torch.empty( + input.numel() // 2, dtype=torch.uint8, device=input.device + ) + # pack the int4 weights into a uint8 tensor + # packing format: w0, w1, w2, w3, w4, w5, ... + # | byte | byte | byte | + packed_output_uint8 = flattened[::2] << 4 | flattened[1::2] + + return cls(input.shape, input.dtype, packed_output_uint8), scales + + def dequantize(self, dtype: torch.dtype = None, **kwarg): + """Dequantze INT4 packed tensor to a target dtype.""" + if dtype is None: + dtype = self.metadata["dtype"] + cuda_ext = get_cuda_ext() + + # get kwargs + scales = kwarg["scale"] + block_sizes = kwarg["block_sizes"] + scale_quant_maxbound = self._get_quant_maxbound(num_bits=4) + + if cuda_ext and self._quantized_data.is_cuda: + # use a custom cuda kernel if available + output = cuda_ext.INT4_dequantize(self._quantized_data, scales, block_sizes[-1]) + return output.view(self.metadata["shape"]).to(dtype) + else: + # indexing in torch required long dtype, we may need to optimize this with customized kernels + # convert (0, 15) -> (-8, 7) + first_half = (self._quantized_data >> 4).to(torch.long) - (scale_quant_maxbound + 1) + second_half = (self._quantized_data & 0x0F).to(torch.long) - (scale_quant_maxbound + 1) + + # de-quantize tensor + first_half = first_half.view(-1, block_sizes[-1] // 2) / scales.view(-1, 1) + second_half = second_half.view(-1, block_sizes[-1] // 2) / scales.view(-1, 1) + + # merge the interleaving elements + first_half = first_half.flatten().unsqueeze(-1).transpose(0, 1) + second_half = second_half.flatten().unsqueeze(-1).transpose(0, 1) + return ( + torch.stack([first_half, second_half], dim=-1) + .reshape(self.metadata["shape"]) + .to(dtype) + ) diff --git a/modelopt/torch/quantization/qtensor/nf4_tensor.py b/modelopt/torch/quantization/qtensor/nf4_tensor.py new file mode 100644 index 0000000000..c1da7e6642 --- /dev/null +++ b/modelopt/torch/quantization/qtensor/nf4_tensor.py @@ -0,0 +1,185 @@ +# 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. + +"""Implements NF4 quantization for efficient tensor storage and computation.""" + +import torch + +from ..extensions import get_cuda_ext +from ..qtensor.base_qtensor import BaseQuantizedTensor +from ..utils import reduce_amax + +nf4_table = torch.tensor( + [ + -1.0000, + -0.6962, + -0.5251, + -0.3949, + -0.2844, + -0.1848, + -0.0911, + 0.0000, + 0.0796, + 0.1609, + 0.2461, + 0.3379, + 0.4407, + 0.5626, + 0.7230, + 1.0000, + ], + dtype=torch.bfloat16, +) + +__all__ = ["NF4QTensor"] + + +def _dequantize_scalers(scales, double_scale, scale_zeros, dtype): + return (scales / double_scale.unsqueeze(-1)).to(dtype) + scale_zeros + + +def _quantize_to_nearest_lut(flatten_tensor: torch.Tensor, lut: torch.Tensor) -> torch.Tensor: + """Quantize a float16 tensor to nearest value and return the indices of the look-up table.""" + assert ( + flatten_tensor.dim() == 1 + ), f"Expect flatten tensor but got input with {flatten_tensor.dim()} dimensions." + diff = (flatten_tensor[:, None] - lut).abs() # Shape: (numel, 16) + indices = diff.argmin(dim=-1) + return indices + + +def _nf4_lookup(quantized_idx: torch.Tensor) -> torch.Tensor: + """Dequantize uint4 index to nf4 value.""" + return nf4_table.to(quantized_idx.device)[quantized_idx] + + +class NF4QTensor(BaseQuantizedTensor): + """Implements the NF4 quantization on tensors for more efficient storage or computation. + + Attributes: + quantized_data (torch.Tensor): The quantized data stored as a packed uint8 tensor. + """ + + @classmethod + def quantize(cls, input: torch.Tensor, block_size: int) -> torch.Tensor: + """Converting a tensor to a quantized format based on NF4 double quantization. + + Args: + input (torch.Tensor): The input tensor to be quantized. + block_size (int): The size of each block for quantization. + scale_block_size (int): The block size for scaling during quantization. + + Returns: + tuple: Contains quantized data, input quantization config, and scale quantization config. + """ + cuda_ext = get_cuda_ext() + + # get scales for each block + assert ( + input.numel() % block_size == 0 + ), "Number of input elements is not divisible by the block size." + block_input = input.view(-1, block_size) + scales = reduce_amax(block_input, -1) + + if cuda_ext and input.is_cuda: + packed_output_uint8 = cuda_ext.NF4_quantize(input, scales, block_size) + else: + # expand scalers to match shape of input + scales = scales.view(block_input.shape[0], -1) # shape: (block_input.shape[0], 1) + scaled_blocks = block_input / scales + + quantized_output = torch.empty( + block_input.numel(), dtype=torch.uint8, device=block_input.device + ) + flattened = scaled_blocks.flatten() + quantized_output = _quantize_to_nearest_lut( + flattened, nf4_table.to(device=input.device, dtype=input.dtype) + ) + quantized_output_uint8 = quantized_output.to(torch.uint8) + + # pack the int4 weights into a uint8 tensor + # packing format: w0, w1, w2, w3, w4, w5, ... + # | byte | byte | byte | + packed_output_uint8 = quantized_output_uint8[::2] << 4 | quantized_output_uint8[1::2] + + return cls(input.shape, input.dtype, packed_output_uint8), scales + + @classmethod + def double_quantization(cls, scales: torch.Tensor, scale_block_size: int, num_scale_bits: int): + """Perform double quantization on the scales. + + Unlike the `quantize` method quantizing input data, this function quantizes float scales into + int8 to further reduce memory usage of scales. + """ + # Double quantization for the scales, int8 per-block quantization + assert ( + scales.numel() % scale_block_size == 0 + ), "Number of scales elements is not divisible by the scale block size." + scale_quant_maxbound = 2 ** (num_scale_bits - 1) - 1 + block_scales = scales.view(-1, scale_block_size) + num_scale_blocks = block_scales.shape[0] + scalers_zero_point = block_scales.mean() + block_scales = block_scales - scalers_zero_point + + double_quant_scales = scale_quant_maxbound / reduce_amax(block_scales, -1) + quantized_scales = block_scales * double_quant_scales.expand( + num_scale_blocks, scale_block_size + ) + quantized_scales = ( + quantized_scales.round() + .clamp(-scale_quant_maxbound, scale_quant_maxbound) + .to(torch.int8) + ) + + return quantized_scales, double_quant_scales.flatten(), scalers_zero_point + + def dequantize(self, dtype: torch.dtype = None, **kwarg): + """Dequantze NF4 packed tensor to a target dtype.""" + if dtype is None: + dtype = self.metadata["dtype"] + cuda_ext = get_cuda_ext() + + # get kwargs + scales = kwarg["scale"] + block_sizes = kwarg["block_sizes"] + double_scale = kwarg["double_scale"] + scale_zeros = kwarg["scale_zeros"] + + if cuda_ext and self._quantized_data.is_cuda: + # with a custom cuda kernel + scales = _dequantize_scalers(scales, double_scale, scale_zeros, dtype).flatten() + output = cuda_ext.NF4_dequantize(self._quantized_data, scales, block_sizes[-1]) + return output.reshape(self.metadata["shape"]).to(dtype) + else: + # de-qauntize scales + scales = _dequantize_scalers(scales, double_scale, scale_zeros, dtype).flatten() + # indexing in torch required long dtype, we may need to optimize this with customized kernels + first_half_idx = (self._quantized_data >> 4).to(torch.long) + second_half_idx = (self._quantized_data & 0x0F).to(torch.long) + scaled_first_half = _nf4_lookup(first_half_idx) # w0, w2, w4... + scaled_second_half = _nf4_lookup(second_half_idx) # w1, w3, w5... + + # de-quantize tensor + first_half = scaled_first_half.view(-1, block_sizes[-1] // 2) * scales.view(-1, 1) + second_half = scaled_second_half.view(-1, block_sizes[-1] // 2) * scales.view(-1, 1) + + # merge the interleaving elements + first_half = first_half.flatten().unsqueeze(-1).transpose(0, 1) + second_half = second_half.flatten().unsqueeze(-1).transpose(0, 1) + return ( + torch.stack([first_half, second_half], dim=-1) + .reshape(self.metadata["shape"]) + .to(dtype) + ) diff --git a/modelopt/torch/quantization/quant_modules.py b/modelopt/torch/quantization/quant_modules.py new file mode 100644 index 0000000000..7c686a32d9 --- /dev/null +++ b/modelopt/torch/quantization/quant_modules.py @@ -0,0 +1,50 @@ +# 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. + +"""Deprecated. Placeholder module for throwing deprecated error.""" + +from contextlib import contextmanager + +from modelopt.torch.utils import DeprecatedError + + +def initialize(*args, **kwargs): + """Deprecated. This API is no longer supported. + + Use :meth:`mtq.quantize() ` + instead to quantize the model. + """ + raise DeprecatedError( + "This API is no longer supported. Use `modelopt.torch.quantization.model_quant.quantize()`" + " API instead to quantize the model." + ) + + +def deactivate(): + """Deprecated. This API is no longer supported.""" + raise DeprecatedError( + "This API is no longer supported. Use " + "`modelopt.torch.quantization.model_quant.disable_quantizer()` API instead " + "to disable quantization." + ) + + +@contextmanager +def enable_onnx_export(): + """Deprecated. You no longer need to use this context manager while exporting to ONNX.""" + raise DeprecatedError( + "You no longer need to use this context manager while exporting to ONNX! please call" + " `torch.onnx.export` directly." + ) diff --git a/modelopt/torch/quantization/src/tensor_quant.cpp b/modelopt/torch/quantization/src/tensor_quant.cpp new file mode 100644 index 0000000000..73ba63c77b --- /dev/null +++ b/modelopt/torch/quantization/src/tensor_quant.cpp @@ -0,0 +1,77 @@ +/* + * 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. + */ + +#include "tensor_quant.h" + +at::Tensor NF4_dequantize(torch::Tensor quantized_data, torch::Tensor scales, int block_size) { + TORCH_CHECK(quantized_data.is_cuda()); + return NF4_dequantize_cuda(quantized_data, scales, block_size); +} +at::Tensor NF4_quantize(torch::Tensor quantized_data, torch::Tensor scales, int block_size) { + TORCH_CHECK(quantized_data.is_cuda()); + return NF4_quantize_cuda(quantized_data, scales, block_size); +} +at::Tensor INT4_dequantize(torch::Tensor quantized_data, torch::Tensor scales, int block_size) { + TORCH_CHECK(quantized_data.is_cuda()); + return INT4_dequantize_cuda(quantized_data, scales, block_size); +} +at::Tensor INT4_quantize(torch::Tensor quantized_data, torch::Tensor scales, int block_size) { + TORCH_CHECK(quantized_data.is_cuda()); + return INT4_quantize_cuda(quantized_data, scales, block_size); +} + +void fake_tensor_quant_(at::Tensor inputs, at::Tensor amax, int num_bits = 8, + bool is_unsigned = false, bool narrow_range = true) { + TORCH_CHECK(inputs.is_cuda()); + TORCH_CHECK(inputs.is_contiguous()) // in-place on non-contiguous tensor is + // more difficult + TORCH_CHECK(amax.numel(), 1); + fake_tensor_quant_cuda_inplace(inputs, amax, num_bits, is_unsigned, narrow_range); +} +// TODO: Can we add support for CPU tensors here? +at::Tensor fake_tensor_quant(at::Tensor inputs, at::Tensor amax, int num_bits = 8, + bool is_unsigned = false, bool narrow_range = true) { + TORCH_CHECK(inputs.is_cuda()); + TORCH_CHECK(amax.numel(), 1); + return fake_tensor_quant_cuda(inputs.contiguous(), amax.contiguous(), num_bits, is_unsigned, + narrow_range); +} + +at::Tensor fake_tensor_quant_with_axis(at::Tensor inputs, at::Tensor amax, int axis, + int num_bits = 8, bool is_unsigned = false, + bool narrow_range = true) { + TORCH_CHECK(inputs.is_cuda()); + TORCH_CHECK(amax.numel(), inputs.size(axis)); + return fake_tensor_quant_with_axis_cuda(inputs.contiguous(), amax.contiguous(), axis, num_bits, + is_unsigned, narrow_range); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fake_tensor_quant_", &fake_tensor_quant_, "Fake Tensor Quant Inplace", py::arg("inputs"), + py::arg("amax"), py::arg("num_bits") = 8, py::arg("unsigned") = false, + py::arg("narrow_range") = true); + m.def("fake_tensor_quant", &fake_tensor_quant, "Fake Tensor Quant", py::arg("inputs"), + py::arg("amax"), py::arg("num_bits") = 8, py::arg("unsigned") = false, + py::arg("narrow_range") = true); + m.def("fake_tensor_quant_with_axis", &fake_tensor_quant_with_axis, "Fake Tensor Quant with axis", + py::arg("inputs"), py::arg("amax"), py::arg("axis"), py::arg("num_bits") = 8, + py::arg("unsigned") = false, py::arg("narrow_range") = true); + m.def("NF4_dequantize", &NF4_dequantize, "Dequantize NF4 weights."); + m.def("NF4_quantize", &NF4_quantize, "Quantize NF4 weights."); + m.def("INT4_dequantize", &INT4_dequantize, "Dequantize INT4 weights."); + m.def("INT4_quantize", &INT4_quantize, "Quantize INT4 weights."); +} diff --git a/modelopt/torch/quantization/src/tensor_quant.h b/modelopt/torch/quantization/src/tensor_quant.h new file mode 100644 index 0000000000..de0c9a9cd0 --- /dev/null +++ b/modelopt/torch/quantization/src/tensor_quant.h @@ -0,0 +1,69 @@ +/* + * 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. + */ + +#include + +// TODO: add descriptions for functions +void fake_tensor_quant_cuda_inplace(at::Tensor, at::Tensor, int, bool, bool); +at::Tensor fake_tensor_quant_cuda(at::Tensor, at::Tensor, int, bool, bool); +at::Tensor fake_tensor_quant_with_axis_cuda(at::Tensor, at::Tensor, int, int, bool, bool); +float bits_to_bound(int, int); +at::Tensor fake_e4m3fy_cuda(at::Tensor inputs); + +// Dequantizes data using NF4 quantization scheme and per-block scaling factors. +// +// Args: +// quantized_data (torch::Tensor): Quantized data packed into uint8. +// scales (torch::Tensor): Scaling factors for each data block. +// block_size (int): Number of elements per data block. +// +// Returns: +// at::Tensor: Dequantized data as bf16 tensor. +at::Tensor NF4_dequantize_cuda(torch::Tensor, torch::Tensor, int); + +// Quantizes data using NF4 quantization scheme and per-block scaling factors. +// +// Args: +// quantized_data (torch::Tensor): High precision data in fp32/fp16/bf16. +// scales (torch::Tensor): Scaling factors for each data block. +// block_size (int): Number of elements per data block. +// +// Returns: +// at::Tensor: Quantized data packed in uint8. +at::Tensor NF4_quantize_cuda(torch::Tensor, torch::Tensor, int); + +// Dequantizes data using INT4 quantization scheme and per-block scaling factors. +// +// Args: +// quantized_data (torch::Tensor): Quantized data packed into uint8. +// scales (torch::Tensor): Scaling factors for each data block. +// block_size (int): Number of elements per data block. +// +// Returns: +// at::Tensor: Dequantized data as the same dtype of scales. +at::Tensor INT4_dequantize_cuda(torch::Tensor, torch::Tensor, int); + +// Quantizes data using INT4 quantization scheme and per-block scaling factors. +// +// Args: +// quantized_data (torch::Tensor): High precision data in fp32/fp16/bf16. +// scales (torch::Tensor): Scaling factors for each data block. +// block_size (int): Number of elements per data block. +// +// Returns: +// at::Tensor: Quantized data packed in uint8. +at::Tensor INT4_quantize_cuda(torch::Tensor, torch::Tensor, int); diff --git a/modelopt/torch/quantization/src/tensor_quant_fp8.cpp b/modelopt/torch/quantization/src/tensor_quant_fp8.cpp new file mode 100644 index 0000000000..2bf9669ac7 --- /dev/null +++ b/modelopt/torch/quantization/src/tensor_quant_fp8.cpp @@ -0,0 +1,48 @@ +/* + * 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. + */ + +#include +#include +#include + +at::Tensor fake_e4m3fy_cuda(at::Tensor inputs); +at::Tensor fused_fake_e4m3fy_cuda(at::Tensor inputs, at::Tensor amax, const float zero_threshold); + +at::Tensor fake_e4m3fy(at::Tensor inputs) { + if (inputs.is_cuda()) { + return fake_e4m3fy_cuda(inputs.contiguous()); + } else { + TORCH_CHECK(inputs.dtype() == at::ScalarType::Float); + TORCH_CHECK(inputs.is_contiguous()); + auto out = at::zeros_like(inputs); + for (int i = 0; i < inputs.numel(); ++i) { + out.data_ptr()[i] = + static_cast(static_cast<__nv_fp8_e4m3>(inputs.data_ptr()[i])); + } + return out; + } +} + +at::Tensor fused_fake_e4m3fy(at::Tensor inputs, at::Tensor amax, const float zero_threshold) { + return fused_fake_e4m3fy_cuda(inputs.contiguous(), amax, zero_threshold); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fake_e4m3fy", &fake_e4m3fy, "Reduce precision to E4M3", py::arg("inputs")); + m.def("fused_fake_e4m3fy", &fused_fake_e4m3fy, "Reduce precision to E4M3 (fused)", + py::arg("inputs"), py::arg("amax"), py::arg("zero_threshold")); +} diff --git a/modelopt/torch/quantization/src/tensor_quant_gpu.cu b/modelopt/torch/quantization/src/tensor_quant_gpu.cu new file mode 100644 index 0000000000..2982ea589f --- /dev/null +++ b/modelopt/torch/quantization/src/tensor_quant_gpu.cu @@ -0,0 +1,361 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include + +#define BLOCK_SIZE 128 +#define EPSILON (1. / (1 << 24)) // Minimum representable of fp16 + +#define AT_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define AT_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +__host__ __device__ float bits_to_bound(int num_bits, int is_unsigned) { + float bound = (1 << (num_bits - 1 + int(is_unsigned))) - 1; + return bound; +} + +__device__ float fake_tensor_quant_device(float input, float amax, int min_bound, int max_bound) { + CUDA_KERNEL_ASSERT(amax >= 0); + + if (amax < EPSILON) { + return 0.f; + } + + float scale = max_bound / amax; + float output = rint(input * scale); + output = output > max_bound ? max_bound : output; + output = output < min_bound ? min_bound : output; + + return output / scale; +} + +template +__global__ void fake_tensor_quant_kernel(const T *inputs, size_t n, T *outputs, const float *amax, + int num_bits = 8, bool is_unsigned = false, + bool narrow_range = true) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + if (tid < n) { + if (is_unsigned) { + CUDA_KERNEL_ASSERT(inputs[tid] >= 0); + } + float bound = bits_to_bound(num_bits, is_unsigned); + float max_bound = bound; + float min_bound = -(bound + !narrow_range); + outputs[tid] = fake_tensor_quant_device((float)inputs[tid], amax[0], min_bound, max_bound); + } +} + +void fake_tensor_quant_cuda_inplace(at::Tensor inputs, at::Tensor amax, int num_bits = 8, + bool is_unsigned = false, bool narrow_range = true) { + size_t numel = inputs.numel(); + AT_DISPATCH_FLOATING_TYPES(inputs.type().scalarType(), "fake_tensor_quant_cuda_inplace", [&] { + fake_tensor_quant_kernel<<>>( + inputs.data_ptr(), numel, inputs.data_ptr(), + amax.to(at::ScalarType::Float).data_ptr(), num_bits, is_unsigned, narrow_range); + }); +} + +at::Tensor fake_tensor_quant_cuda(at::Tensor inputs, at::Tensor amax, int num_bits = 8, + bool is_unsigned = false, bool narrow_range = true) { + size_t numel = inputs.numel(); + auto outputs = torch::empty_like(inputs); + AT_DISPATCH_FLOATING_TYPES(inputs.type().scalarType(), "fake_tensor_quant_cuda", [&] { + fake_tensor_quant_kernel<<>>( + inputs.data_ptr(), numel, outputs.data_ptr(), + amax.to(at::ScalarType::Float).data_ptr(), num_bits, is_unsigned, narrow_range); + }); + + return outputs; +} + +template +__global__ void +fake_tensor_quant_with_axis_cuda_kernel(const T *inputs, size_t n, T *outputs, const float *amax, + int axis_size, int outer_size, int num_bits = 8, + bool is_unsigned = false, bool narrow_range = true) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + float bound = bits_to_bound(num_bits, is_unsigned); + float max_bound = bound; + float min_bound = -(bound + !narrow_range); + + for (int idx = 4 * tid; idx < 4 * (tid + 1) && idx < n; ++idx) { + if (is_unsigned) { + CUDA_KERNEL_ASSERT(inputs[idx] >= 0); + } + int axis_idx = (idx / outer_size) % axis_size; + + outputs[idx] = + fake_tensor_quant_device((float)inputs[idx], amax[axis_idx], min_bound, max_bound); + } +} + +at::Tensor fake_tensor_quant_with_axis_cuda(at::Tensor inputs, at::Tensor amax, int axis, + int num_bits = 8, bool is_unsigned = false, + bool narrow_range = true) { + auto outputs = torch::empty_like(inputs); + size_t numel = inputs.numel(); + int axis_size = inputs.size(axis); + + int outer_size = inputs.stride(axis); + + AT_DISPATCH_FLOATING_TYPES(inputs.type().scalarType(), "fake_tensor_quant_cuda_with_axis", [&] { + fake_tensor_quant_with_axis_cuda_kernel<<>>( + inputs.data_ptr(), numel, outputs.data_ptr(), + amax.to(at::ScalarType::Float).data_ptr(), axis_size, outer_size, num_bits, + is_unsigned, narrow_range); + }); + return outputs; +} + +__constant__ float NF4_LUT[16] = {-1.0000f, -0.6962f, -0.5251f, -0.3949f, -0.2844f, -0.1848f, + -0.0911f, 0.0000f, 0.0796f, 0.1609f, 0.2461f, 0.3379f, + 0.4407f, 0.5626f, 0.7230f, 1.0000f}; + +template +__global__ void NF4_dequantize_kernel(const uint8_t *__restrict__ quantized_data, + const T *__restrict__ scales, + at::BFloat16 *__restrict__ output, int block_size) { + const int byte_idx = blockIdx.x * blockDim.x + threadIdx.x; + + // load scale + const auto scale = at::BFloat16(scales[blockIdx.x]); + + // load the int8 bytes + const int packed_value = quantized_data[byte_idx]; + const int first_half_idx = (packed_value >> 4) & 0x0F; + const int second_half_idx = packed_value & 0x0F; + + // NF4 look up + const at::BFloat16 first_half_value = NF4_LUT[first_half_idx]; + const at::BFloat16 second_half_value = NF4_LUT[second_half_idx]; + + // de-quantize + output[2 * byte_idx] = first_half_value * scale; + output[2 * byte_idx + 1] = second_half_value * scale; +} + +at::Tensor NF4_dequantize_cuda(torch::Tensor quantized_data, torch::Tensor scales, int block_size) { + auto device = quantized_data.device(); + const int num_bytes = quantized_data.numel(); + const int num_scales = scales.numel(); + + const auto options = quantized_data.options().dtype(at::kBFloat16).device(device); + auto output = torch::empty({2 * num_bytes}, options); + + const int blocks = num_bytes / (block_size / 2); + + // checks before launching the kernel + CUDA_KERNEL_ASSERT(quantized_data.is_cuda() && scales.is_cuda()); + CUDA_KERNEL_ASSERT(device.index() == scales.device().index() && + device.index() == output.device().index()); + CUDA_KERNEL_ASSERT(quantized_data.dtype() == torch::kUInt8); + CUDA_KERNEL_ASSERT(num_bytes * 2 == num_scales * block_size); + + cudaSetDevice(device.index()); + // each thread de-quantize a byte of quantized_data + AT_DISPATCH_FLOATING_TYPES(scales.type().scalarType(), "NF4_dequantize", [&] { + NF4_dequantize_kernel<<>>(quantized_data.data_ptr(), + scales.data_ptr(), + output.data_ptr(), block_size); + }); + + return output; +} + +__device__ int find_closest_index(float *LUT, float value, int lut_size) { + float min_diff = fabs(LUT[0] - value); + int min_index = 0; + for (int i = 1; i < lut_size; ++i) { + float diff = fabs(LUT[i] - value); + if (diff < min_diff) { + min_diff = diff; + min_index = i; + } + } + return min_index; +} + +template +__global__ void NF4_quantize_kernel(const T *__restrict__ input, const T *__restrict__ scales, + uint8_t *__restrict__ output, int block_size) { + const int byte_idx = blockIdx.x * blockDim.x + threadIdx.x; + + // Load input and scale + const auto first_input = input[byte_idx * 2]; + const auto second_input = input[byte_idx * 2 + 1]; + const auto scale = scales[blockIdx.x]; + + // Scale input + const auto scaled_first_input = first_input / scale; + const auto scaled_second_input = second_input / scale; + + // Find the index + const int first_index = find_closest_index(NF4_LUT, scaled_first_input, 16); + const int second_index = find_closest_index(NF4_LUT, scaled_second_input, 16); + + // Pack the value + const uint8_t packed_value = ((first_index & 0xF) << 4) | ((second_index & 0xF)); + + // Store packed value in output tensor + output[byte_idx] = packed_value; +} + +torch::Tensor NF4_quantize_cuda(torch::Tensor input, torch::Tensor scales, int block_size) { + auto device = input.device(); + const int numel = input.numel(); + const int num_scales = scales.numel(); + const int output_byte = numel / 2; + + auto options = input.options().dtype(torch::kUInt8).device(device); + auto output = torch::empty({output_byte}, options); + + // checks before launching the kernel + CUDA_KERNEL_ASSERT(input.is_cuda() && scales.is_cuda() && output.is_cuda()); + CUDA_KERNEL_ASSERT(device.index() == scales.device().index() && + device.index() == output.device().index()); + CUDA_KERNEL_ASSERT(block_size % 2 == 0); + CUDA_KERNEL_ASSERT(numel % (2 * block_size) == 0); + + const int blocks = numel / block_size; + cudaSetDevice(device.index()); + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "NF4_quantize", [&] { + NF4_quantize_kernel<<>>(input.data_ptr(), + scales.data_ptr(), + output.data_ptr(), block_size); + }); + return output; +} + +template +__global__ void INT4_dequantize_kernel(const uint8_t *__restrict__ quantized_data, + const T *__restrict__ scales, T *__restrict__ output, + int block_size) { + const int byte_idx = blockIdx.x * blockDim.x + threadIdx.x; + const T scale_quant_maxbound = 7; + + // load scale + const auto scale = scales[blockIdx.x]; + + // load the scaled values and convert them to (-(scale_quant_maxbound+1), scale_quant_maxbound) + const int packed_value = quantized_data[byte_idx]; + const T first_half_value = ((packed_value >> 4) & 0x0F) - (scale_quant_maxbound + 1); + const T second_half_value = (packed_value & 0x0F) - (scale_quant_maxbound + 1); + + // de-quantize + output[2 * byte_idx] = first_half_value / scale; + output[2 * byte_idx + 1] = second_half_value / scale; +} + +at::Tensor INT4_dequantize_cuda(torch::Tensor quantized_data, torch::Tensor scales, + int block_size) { + auto device = quantized_data.device(); + const int num_bytes = quantized_data.numel(); + const int num_scales = scales.numel(); + + const auto options = scales.options().device(device); + auto output = torch::empty({2 * num_bytes}, options); + + const int blocks = num_bytes / (block_size / 2); + + // checks before launching the kernel + CUDA_KERNEL_ASSERT(quantized_data.is_cuda() && scales.is_cuda() && output.is_cuda()); + CUDA_KERNEL_ASSERT(device.index() == scales.device().index() && + device.index() == output.device().index()); + CUDA_KERNEL_ASSERT(quantized_data.dtype() == torch::kUInt8); + CUDA_KERNEL_ASSERT(num_bytes * 2 == num_scales * block_size); + + // each thread de-quantize a byte of quantized_data + cudaSetDevice(device.index()); + AT_DISPATCH_FLOATING_TYPES(scales.type().scalarType(), "INT4_dequantize", [&] { + INT4_dequantize_kernel<<>>(quantized_data.data_ptr(), + scales.data_ptr(), + output.data_ptr(), block_size); + }); + + return output; +} + +template +__global__ void INT4_quantize_kernel(const T *__restrict__ input, const T *__restrict__ scales, + uint8_t *__restrict__ output, int block_size) { + const int byte_idx = blockIdx.x * blockDim.x + threadIdx.x; + const T scale_quant_maxbound = 7; + + // Load input and scale + const auto first_input = input[byte_idx * 2]; + const auto second_input = input[byte_idx * 2 + 1]; + const auto scale = scales[blockIdx.x]; + + // Scale input and clamp to (-(scale_quant_maxbound+1), scale_quant_maxbound) + auto scaled_first_input = first_input * scale; + scaled_first_input = + max(-(scale_quant_maxbound + 1), min(scale_quant_maxbound, scaled_first_input)); + auto scaled_second_input = second_input * scale; + scaled_second_input = + max(-(scale_quant_maxbound + 1), min(scale_quant_maxbound, scaled_second_input)); + + // int4 to uint4: 0 - 15 + const int uint4_first_input = + static_cast(roundf(scaled_first_input + (scale_quant_maxbound + 1))); + const int uint4_second_input = + static_cast(roundf(scaled_second_input + (scale_quant_maxbound + 1))); + + // Pack the value + const uint8_t packed_value = ((uint4_first_input & 0xF) << 4) | ((uint4_second_input & 0xF)); + + // Store packed value in output tensor + output[byte_idx] = packed_value; +} + +torch::Tensor INT4_quantize_cuda(torch::Tensor input, torch::Tensor scales, int block_size) { + auto device = input.device(); + const int numel = input.numel(); + const int num_scales = scales.numel(); + const int output_byte = numel / 2; + + auto options = input.options().dtype(torch::kUInt8); + auto output = torch::empty({output_byte}, options); + + // checks before launching the kernel + CUDA_KERNEL_ASSERT(input.is_cuda() && scales.is_cuda() && output.is_cuda()); + CUDA_KERNEL_ASSERT(device.index() == scales.device().index() && + device.index() == output.device().index()); + CUDA_KERNEL_ASSERT(block_size % 2 == 0); + CUDA_KERNEL_ASSERT(numel % (2 * block_size) == 0); + + const int blocks = numel / block_size; + cudaSetDevice(device.index()); + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "INT4_quantize", [&] { + INT4_quantize_kernel<<>>(input.data_ptr(), + scales.data_ptr(), + output.data_ptr(), block_size); + }); + return output; +} diff --git a/modelopt/torch/quantization/src/tensor_quant_gpu_fp8.cu b/modelopt/torch/quantization/src/tensor_quant_gpu_fp8.cu new file mode 100644 index 0000000000..9081d0e143 --- /dev/null +++ b/modelopt/torch/quantization/src/tensor_quant_gpu_fp8.cu @@ -0,0 +1,122 @@ +/* + * 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. + */ + +#include +#include +#include +#include + +#define BLOCK_SIZE 128 + +#define AT_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define AT_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +template __global__ void fake_e4m3fy_kernel(const T *inputs, size_t n, T *outputs) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + for (int idx = 4 * tid; idx < 4 * (tid + 1) && idx < n; ++idx) { + outputs[idx] = static_cast( + static_cast(static_cast<__nv_fp8_e4m3>(static_cast(inputs[idx])))); + } +} + +template +__global__ void fused_fake_e4m3fy_kernel(const T *inputs, size_t n, float *amax, + bool per_block_scaling_factor, size_t blocksize, + float zero_threshold, T *outputs) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + for (int idx = 4 * tid; idx < 4 * (tid + 1) && idx < n; ++idx) { + float x = static_cast(inputs[idx]); + + // generate mask for zeroing tiny values + float x_abs = fabsf(x); + bool zero_mask = x_abs < zero_threshold; + + // grab the global scaling factor + size_t amax_idx = (per_block_scaling_factor) ? (idx / blocksize) : 0; + + // compute scale and inverse-scales + float scale = 448.f / (amax[amax_idx]); + float inv_scale = 1.f / scale; + + // compute the output + float output = static_cast(static_cast<__nv_fp8_e4m3>(scale * x)) * inv_scale; + + // zero out small values + if (zero_mask) { + output = 0.f; + } + + outputs[idx] = output; + } +} + +at::Tensor fused_fake_e4m3fy_cuda(at::Tensor inputs, at::Tensor amax, const float zero_threshold) { + size_t numel = inputs.numel(); + auto outputs = torch::empty_like(inputs); + + bool per_block_scaling_factor = false; + size_t blocksize = numel; + + int amax_ndim = amax.dim(); + int input_ndim = inputs.dim(); + + // 3 options: + // 1. + // inputs[numel], amax[1] -> per-tensor scaling + // 2. + // inputs[numel], amax[numel/num_cols] -> per-row / per-channel scaling + // 3. + // inputs[numel/bs, bs], amax[numel/bs, 1] -> blockwise scaling + if (amax.numel() == 1) { + // case 1. + per_block_scaling_factor = false; + } else if (amax.numel() > 1 && (amax_ndim > 1 && (amax.size(-1) == amax.numel()))) { + // case 2. + per_block_scaling_factor = true; + blocksize = numel / amax.numel(); + } else { + throw std::runtime_error("invalid combination of inputs and amax shapes/sizes"); + } + + auto stream = c10::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES(inputs.type().scalarType(), "fused_fake_e4m3fy_cuda", [&] { + fused_fake_e4m3fy_kernel<<>>( + inputs.data_ptr(), numel, amax.data_ptr(), per_block_scaling_factor, + blocksize, zero_threshold, outputs.data_ptr()); + }); + return outputs; +} + +at::Tensor fake_e4m3fy_cuda(at::Tensor inputs) { + size_t numel = inputs.numel(); + auto outputs = torch::empty_like(inputs); + auto stream = c10::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES(inputs.type().scalarType(), "fake_e4m3fy_cuda", [&] { + fake_e4m3fy_kernel<<>>( + inputs.data_ptr(), numel, outputs.data_ptr()); + }); + return outputs; +} diff --git a/modelopt/torch/quantization/src/tensor_quant_mx.cu b/modelopt/torch/quantization/src/tensor_quant_mx.cu new file mode 100644 index 0000000000..b8278dd7aa --- /dev/null +++ b/modelopt/torch/quantization/src/tensor_quant_mx.cu @@ -0,0 +1,401 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +#include "tensor_quant_mx.h" + +namespace { + +static __device__ __host__ __inline__ float quantize(const float x, const float scale, + const float unscale, const Types format) { + // Get the sign. + int sign; + if (x < 0.0) + sign = -1; + if (x > 0.0) + sign = 1; + // Absolute value of x. + float xabs = fabsf(x); + // Scale. + xabs = xabs * scale; + // Quantize. + xabs = convert_to_types(xabs, format); + // Unscale. + xabs = xabs * unscale; + + // Apply sign. + return sign * xabs; +} + +/* + * Get the maximum for a format for quantization purposes. + * + * Note: we split out e8mY formats as a) they don't need scaling for + * dynamic range, and if we _do_ scale, we can easily get numerical + * weirdnesses - 1/scale has a habit of flushing -> 0 and causing + * NaN in answers. For this reason, effectively disable scaling + * for these formats (scale=unscale=1) + */ + +// Algorithms for E8M0 scaling factor computation: OCP and NV +// Both implementations require or assume the following: +// * amax is the maximum magnitude input value, +// * dmax is the maximum magnitude representable by config, +// * the range of config is symmetric. + +// OCP algorithm: +// scale = 2 ** clamp(floor(log2(amax)) - floor(log2(dmax)), -127, 127), +__device__ __host__ cuda::std::tuple compute_scale_e8m0_OCP(float amax, + Types format) { + // lg2_dmax = log2(maximum power-of-two representable by config) + int32_t lg2_dmax = get_format_lg2_max(format); + + // lg2_amax = floor(log2(amax)) + // The following code assumes amax is a positive, finite FP32 value. + union { + float f32; + uint32_t u32; + } amax_bits = {.f32 = amax}; + uint32_t amax_exponent_field = (amax_bits.u32 >> 23) & 0xff; + uint32_t amax_significand_field = amax_bits.u32 & 0x7fffff; + int32_t lg2_amax = amax_exponent_field == 0 + ? static_cast(cuda::std::bit_width(amax_significand_field)) - 150 + : static_cast(amax_exponent_field) - 127; + + int unscale_exponent = std::clamp(lg2_amax - lg2_dmax, -127, 127); + + float scale = ldexpf(1.f, -unscale_exponent); + float unscale = ldexpf(1.f, unscale_exponent); + + return cuda::std::make_tuple(scale, unscale); +} + +// NVIDIA algorithm: +// unscale = 2 ** clamp (ceil(log2(amax/dmax)), -127, 127) +__device__ __host__ cuda::std::tuple compute_scale_e8m0_NV(float amax, Types format) { + // TODO: currently this FP32 division uses RNE. When the exact result + // is just above a binade boundary, it will be rounded down, in which + // case this routine may return the floor(), not the ceil()! + // + // We can fix this by performing the fdiv using RU. On __device__, + // this can be accomplished by __fdiv_ru intrinsic. On __host__, + // we may have to use . + + float ratio = amax / get_format_max(format); + + // scale_exponent = clamp(ceil(log2(ratio)), -127, 127) + // The following code assumes ratio is a positive, finite FP32 value. + union { + float f32; + uint32_t u32; + } ratio_bits = {.f32 = ratio}; + uint32_t ratio_exponent_field = (ratio_bits.u32 >> 23) & 0xff; + uint32_t ratio_significand_field = ratio_bits.u32 & 0x7fffff; + int32_t unscale_exponent = + ratio_significand_field > 0 && ratio_exponent_field != 0xfe && + !(ratio_exponent_field == 0 && ratio_significand_field <= 0x400000) + ? static_cast(ratio_exponent_field) - 126 + : static_cast(ratio_exponent_field) - 127; + + float scale = ldexpf(1.f, -unscale_exponent); + float unscale = ldexpf(1.f, unscale_exponent); + + return cuda::std::make_tuple(scale, unscale); +} + +// Compute scaling factor. +__device__ __host__ __inline__ cuda::std::tuple +compute_scale(const float amax, const Types format, const Types scale_format) { + if (amax == 0 || isinf(amax) || isnan(amax)) { + return cuda::std::make_tuple(1.f, 1.f); + } else { + // intercept e8m0 specifically + if (scale_format == Types::E8M0) { + return compute_scale_e8m0_NV(amax, format); + } + // otherwise take the general path + const float emax = get_format_max(format); + double scale = emax / amax; + double inv_scale = convert_to_types(1. / scale, scale_format); + scale = 1. / inv_scale; + return cuda::std::make_tuple(scale, inv_scale); + } +} + +__device__ __host__ __inline__ cuda::std::tuple +compute_scale_with_global(const float global_amax, const float local_amax, const Types format, + const Types scale_format) { + bool local_amax_invalid = local_amax == 0 || isinf(local_amax) || isnan(local_amax); + bool global_amax_invalid = global_amax == 0 || isinf(global_amax) || isnan(global_amax); + if (local_amax_invalid || global_amax_invalid) { + return cuda::std::make_tuple(1.f, 1.f); + } else { + const float elem_format_max = get_format_max(format); + const float scale_format_max = get_format_max(scale_format); + + const float local_unscale = local_amax / elem_format_max; + // Note: if elem_format_max / global_amax > 1 and scale_format_max + // is close to FLT_MAX, we can overflow here + // WAR is to just do these few calculations in double. + const double two_level_scale = + static_cast(scale_format_max) * (elem_format_max / global_amax); + + const double local_unscale_q = + convert_to_types(static_cast(local_unscale * two_level_scale), scale_format) / + two_level_scale; + const float scale = 1. / local_unscale_q; + + return cuda::std::make_tuple(scale, static_cast(local_unscale_q)); + } +} + +}; // anonymous namespace + +namespace { + +template +__device__ float compute_max_block(const float x, const int block_size, + const float clamp = FLT_MAX) { + float xabs = fabsf(x); + xabs = (xabs > clamp) ? clamp : xabs; + + typedef cub::BlockReduce BlockReduce; + + __shared__ typename BlockReduce::TempStorage temp_storage[THREAD_Y]; + + float amax = BlockReduce(temp_storage[threadIdx.y]).Reduce(xabs, cub::Max()); + + __shared__ float amax_smem[THREAD_Y]; + if (threadIdx.x == 0) + amax_smem[threadIdx.y] = amax; + + __syncthreads(); + + return amax_smem[threadIdx.y]; +} + +template +__device__ float compute_max_warp(const float x, const int block_size, + const float clamp = FLT_MAX) { + float xabs = fabsf(x); + xabs = (xabs > clamp) ? clamp : xabs; + + typedef cub::WarpReduce WarpReduce; + __shared__ typename WarpReduce::TempStorage temp_storage[THREAD_Y]; + + float amax = WarpReduce(temp_storage[threadIdx.y]).Reduce(xabs, cub::Max()); + amax = __shfl_sync(0xffffffff, amax, 0); + return amax; +} + +template +__device__ float compute_max(const float x, const int block_size, const float clamp = FLT_MAX) { + if (THREAD_X == 32) { + return compute_max_warp(x, block_size, clamp); + } else { + return compute_max_block(x, block_size, clamp); + } +} + +}; // namespace + +template +__global__ void amax_and_convert_kernel(const T *x, T_y *y, const int64_t n, const int blocksize, + const int actual_cols, const int padded_cols, Types format, + Types scale_format, const float *global_amax = nullptr, + const bool per_channel_scaling = false, + const int ncols = -1) { + // map thread -> element + const uint64_t block_idx = blockIdx.x * blockDim.y + threadIdx.y; + const int thread_id = threadIdx.x; + const int warp_id = threadIdx.y; + + // Unique id in the padded matrix + const int64_t idx = block_idx * blocksize + thread_id; + + const int blocks_per_row = padded_cols / blocksize; + + if (idx >= n) + return; + + // bool output_scaled = (block_unscale_ptr != nullptr); + + float scale = 1.f, unscale = 1.f; + + // row and col are in terms of the "padded" size that we launched based off + const int row = block_idx / blocks_per_row; + const int block_in_row = block_idx % blocks_per_row; + const int col = block_in_row * blocksize + thread_id; + // Differentiate between real and padded elements + const bool in_padding = col >= actual_cols; + + // map threads to the actual memory locations + // Example: x = [2, 17], bs=16 + // thread @ 0, 0 = 0 + // thread @ 0, 15 = 15 + // thread @ 1, 0 = 17 + const size_t real_idx = row * actual_cols + col; + + // read either real or padded (zero) value + float thread_val = (in_padding) ? 0.f : static_cast(x[real_idx]); + float block_amax = compute_max(thread_val, blocksize); + + if (global_amax) { + int global_idx = (per_channel_scaling) ? idx / ncols : 0; + float global_amax_val = global_amax[global_idx]; + cuda::std::tie(scale, unscale) = + compute_scale_with_global(global_amax_val, block_amax, format, scale_format); + } else { + cuda::std::tie(scale, unscale) = compute_scale(block_amax, format, scale_format); + } + + // if thread isn't handling a real value, nothing to output. + if (in_padding) + return; + + y[real_idx] = quantize(thread_val, scale, unscale, format); +} + +std::tuple get_global_scaling(at::Tensor x, std::optional global_amax) { + float *global_amax_ptr = + global_amax.has_value() ? global_amax.value().data_ptr() : nullptr; + + // if we need per-channel scaling, following must be true: + // 1) 2d input tensor + // 2) 1d blocksize + // 3) # global_amax elements == # rows of collapsed matrix + bool per_channel_scaling = false; + int cols = x.size(-1); + if (global_amax_ptr) { + per_channel_scaling = (x.dim() == 2) && (global_amax.value().numel() == x.size(0)); + } + + return std::make_tuple(global_amax_ptr, per_channel_scaling); +} + +using output_t = std::tuple>; + +#define DISPATCH_FUSED_AMAX_KERNEL(T, TX, TY) \ + amax_and_convert_kernel \ + <<>>( \ + x.data_ptr(), y.data_ptr(), n, blocksize, actual_cols, padded_cols, format, \ + scale_format, global_amax_ptr, per_channel_scaling, x.size(-1)) + +#define DISPATCH_SINGLE_TYPE(T) \ + if (blocksize == 8) { \ + DISPATCH_FUSED_AMAX_KERNEL(T, 8, blocks_per_cta); \ + } else if (blocksize == 16) { \ + DISPATCH_FUSED_AMAX_KERNEL(T, 16, blocks_per_cta); \ + } else if (blocksize == 32) { \ + DISPATCH_FUSED_AMAX_KERNEL(T, 32, blocks_per_cta); \ + } else { \ + throw std::invalid_argument("Blocksize for fused call must be one of {8, 16, 32}"); \ + } + +#define DISPATCH_FLOAT_AND_HALF_AND_BFLOAT(TYPE, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Float: { \ + using scalar_t = float; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::Half: { \ + using scalar_t = at::Half; \ + __VA_ARGS__; \ + break; \ + } \ + case at::ScalarType::BFloat16: { \ + using scalar_t = at::BFloat16; \ + __VA_ARGS__; \ + break; \ + } \ + default: \ + AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } + +int pad_dim_to_blocksize(const int dim, const int bs) { return dim + (bs - (dim % bs)); } + +at::Tensor fused_amax_convert(at::Tensor x, const int blocksize, Types format, Types scale_format, + std::optional global_amax) { + std::optional global_amax_float = std::nullopt; + + if (global_amax.has_value()) { + const at::Tensor &_global_amax = *global_amax; + global_amax_float = at::_cast_Float(_global_amax); + } + // const bool native_mode = is_native(output_mode); + auto y_options = x.options(); + + at::Tensor y = at::empty_like(x, y_options); + + size_t n = x.numel(); + int ndim = x.dim(); + // Handle partial blocks by fake-padding to full blocks + const int actual_cols = x.size(-1); + const int padded_cols = pad_dim_to_blocksize(actual_cols, blocksize); + // pad the total size, replace the final dimension with padded one + n = n / actual_cols * padded_cols; + + const int blocks_per_cta = 4; + + const dim3 block(blocksize, blocks_per_cta); + const int blocks = ceil_div(n, blocksize * blocks_per_cta); + + auto [global_amax_ptr, per_channel_scaling] = get_global_scaling(x, global_amax_float); + + DISPATCH_FLOAT_AND_HALF_AND_BFLOAT(x.scalar_type(), "fused_amax_convert_kernel", + { DISPATCH_SINGLE_TYPE(scalar_t); }); + + return y; +} + +#undef DISPATCH_FUSED_AMAX_KERNEL +#undef DISPATCH_SINGLE_TYPE +#undef DISPATCH_FLOAT_AND_HALF_AND_BFLOAT + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fused_amax_convert", &fused_amax_convert, + "Compute block amax and quantize the tensor to mx formats", py::arg("inputs"), + py::arg("block_size"), py::arg("format"), py::arg("scale_format"), + py::arg("global_amax") = py::none()); + m.def("convert_to_exmy", &convert_to_types, "Convert value to exmy format.", py::arg("x"), + py::arg("format")); + py::enum_(m, "Types") + .value("E4M3", Types::E4M3) + .value("E5M2", Types::E5M2) + .value("INT8", Types::INT8) + .value("E0M3", Types::E0M3) + .value("E1M2", Types::E1M2) + .value("E3M0", Types::E3M0) + .value("E2M1", Types::E2M1) + .value("E3M2", Types::E3M2) + .value("E2M3", Types::E2M3) + .value("E8M0", Types::E8M0); +} diff --git a/modelopt/torch/quantization/src/tensor_quant_mx.h b/modelopt/torch/quantization/src/tensor_quant_mx.h new file mode 100644 index 0000000000..d1714330f4 --- /dev/null +++ b/modelopt/torch/quantization/src/tensor_quant_mx.h @@ -0,0 +1,247 @@ +/* + * 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. + */ + +#pragma once + +#ifndef CUDA_HOST_DEVICE +#ifdef __CUDA_ARCH__ +#define CUDA_HOST_DEVICE __host__ __device__ __forceinline__ +#define MODELOPT_FORMAT_FAIL(str) __trap(); +#else +#define CUDA_HOST_DEVICE inline +#define MODELOPT_FORMAT_FAIL(str) throw std::runtime_error(str); +#endif +#endif + +#ifdef __CUDA_ARCH__ +#define CONSTANT __constant__ +#else +#define CONSTANT static +#endif + +#include +#include + +enum class Types { E4M3, E5M2, INT8, E0M3, E1M2, E3M0, E2M1, E3M2, E2M3, E8M0 }; + +// FP4 +CONSTANT float e2m1_table_values[8] = {0, 0.5, 1, 1.5, 2, 3, 4, 6}; +CONSTANT float e2m1_table_bounds[7] = {0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5}; + +CONSTANT float e1m2_table_values[8] = {0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5}; +CONSTANT float e1m2_table_bounds[7] = {0.25, 0.75, 1.25, 1.75, 2.25, 2.75, 3.25}; + +CONSTANT float e0m3_table_values[8] = {0, 1, 2, 3, 4, 5, 6, 7}; +CONSTANT float e0m3_table_bounds[7] = {0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5}; + +CONSTANT float e3m0_table_values[8] = {0, 0.25, 0.5, 1, 2, 4, 8, 16}; +CONSTANT float e3m0_table_bounds[7] = {0.125, 0.375, 0.75, 1.5, 3, 6, 12}; + +// FP6 +CONSTANT float e3m2_table_values[32] = {0, 0.0625, 0.125, 0.1875, 0.25, 0.3125, 0.375, 0.4375, + 0.5, 0.625, 0.75, 0.875, 1, 1.25, 1.5, 1.75, + 2, 2.5, 3, 3.5, 4, 5, 6, 7, + 8, 10, 12, 14, 16, 20, 24, 28}; +CONSTANT float e3m2_table_bounds[31] = { + 0.03125, 0.09375, 0.15625, 0.21875, 0.28125, 0.34375, 0.40625, 0.46875, 0.5625, 0.6875, 0.8125, + 0.9375, 1.125, 1.375, 1.625, 1.875, 2.25, 2.75, 3.25, 3.75, 4.5, 5.5, + 6.5, 7.5, 9, 11, 13, 15, 18, 22, 26}; + +CONSTANT float e2m3_table_values[32] = { + 0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1, 1.125, 1.25, 1.375, 1.5, 1.625, 1.75, 1.875, + 2, 2.25, 2.5, 2.75, 3, 3.25, 3.5, 3.75, 4, 4.5, 5, 5.5, 6, 6.5, 7, 7.5}; +CONSTANT float e2m3_table_bounds[31] = { + 0.0625, 0.1875, 0.3125, 0.4375, 0.5625, 0.6875, 0.8125, 0.9375, 1.0625, 1.1875, 1.3125, + 1.4375, 1.5625, 1.6875, 1.8125, 1.9375, 2.125, 2.375, 2.625, 2.875, 3.125, 3.375, + 3.625, 3.875, 4.25, 4.75, 5.25, 5.75, 6.25, 6.75, 7.25}; + +CUDA_HOST_DEVICE float convert_to_fp8_e4m3(const float x) { + return static_cast(static_cast<__nv_fp8_e4m3>(static_cast(x))); +} + +CUDA_HOST_DEVICE float convert_to_fp8_e5m2(const float x) { + return static_cast(static_cast<__nv_fp8_e5m2>(static_cast(x))); +} + +CUDA_HOST_DEVICE float convert_int8_saturating(const float x) { + // Regular conversion w/rounding + int converted = std::rint(x); + + // saturate high & low + converted = std::clamp(converted, -INT8_MAX, INT8_MAX); + + // do the actual conversion. + return static_cast(converted); +} + +// // Rounding function. Round to nearest, ties away from zero. +CUDA_HOST_DEVICE float round_away_from_zero(float x, const float *values, const float *bounds, + int n) { + // Within current bounds. + for (int i = 0; i < n - 1; i++) { + if (x < bounds[i]) + return values[i]; + } + // Max bounds. + return values[n - 1]; +} + +// Rounding function. Round to nearest, ties to even +CUDA_HOST_DEVICE float round_to_even(float x, const float *values, const float *bounds, int n) { + // Within current bounds. + for (int i = 0; i < n - 1; i++) { + // If less than tie, round down + if (x < bounds[i]) + return values[i]; + // If equals to tie, round to representable value that has 0 in lsb of the mantissa + if (x == bounds[i]) { + // Since lsb alternates between 0 and 1, we can round down + // if the index is even, and round up if the index is odd + if (i % 2 == 0) + return values[i]; + else + return values[i + 1]; + } + } + // Max bounds. + return values[n - 1]; +} + +CUDA_HOST_DEVICE float convert_to_table_types(float x, const int n, const float *values, + const float *bounds) { + float sign; + if (x < 0.f) { + sign = -1.f; + x = -x; + } else { + sign = 1.f; + } + // Saturate infs and nans. + if (isinf(x) || isnan(x)) + return sign * values[n - 1]; + + // Round. + float value = round_to_even(x, values, bounds, n); + + return value * sign; +} + +CUDA_HOST_DEVICE float convert_to_e3m0(float x, const int n, const float *values, + const float *bounds) { + float sign; + if (x < 0.f) { + sign = -1.f; + x = -x; + } else { + sign = 1.f; + } + // Encoding table for 4-bit formats. + // Saturate infs and nans. + if (isinf(x) || isnan(x)) + return sign * values[n - 1]; + + // Round. + float value = round_away_from_zero(x, values, bounds, n); + + return value * sign; +} + +CUDA_HOST_DEVICE float convert_to_types(const float x, const Types format) { + switch (format) { + case Types::E4M3: + return convert_to_fp8_e4m3(x); + case Types::E5M2: + return convert_to_fp8_e5m2(x); + case Types::INT8: + return convert_int8_saturating(x); + case Types::E2M1: + return convert_to_table_types(x, 8, e2m1_table_values, e2m1_table_bounds); + case Types::E0M3: + return convert_to_table_types(x, 8, e0m3_table_values, e0m3_table_bounds); + case Types::E1M2: + return convert_to_table_types(x, 8, e1m2_table_values, e1m2_table_bounds); + case Types::E3M0: + return convert_to_e3m0(x, 8, e3m0_table_values, e3m0_table_bounds); + case Types::E3M2: + return convert_to_table_types(x, 32, e3m2_table_values, e3m2_table_bounds); + case Types::E2M3: + return convert_to_table_types(x, 32, e2m3_table_values, e2m3_table_bounds); + default: + return 0.f; + } +} + +CUDA_HOST_DEVICE +float get_format_max(const Types format) { + switch (format) { + case Types::E4M3: + return 448.f; + case Types::E5M2: + return 57344.f; + case Types::INT8: + return 127.f; + case Types::E2M1: + return 6.f; + case Types::E0M3: + return 7.f; + case Types::E3M0: + return 16.f; + case Types::E1M2: + return 3.5f; + case Types::E3M2: + return 28.f; + case Types::E2M3: + return 7.5f; + default: + // error + return 0.f; + } +} + +// Returns the exponent of the maximum power-of-two representable in the given format. +__host__ __device__ __inline__ int get_format_lg2_max(Types format) { + switch (format) { + case Types::E5M2: + return 15; + case Types::E4M3: + return 8; + case Types::E3M2: + case Types::E3M0: + return 4; + case Types::E2M3: + case Types::E2M1: + return 2; + case Types::E1M2: + return 1; + case Types::E0M3: + return 2; + case Types::INT8: + // This follows the psx-encodings definition of "INT8", + // which are the 8-bit two's complement integers, symmetrized + // by removing -128. + // + // In the context of OCP MX-formats, "INT8" refers to a + // symmetrized Q1.6 type. In that case, this function + // should return 0. + return 6; + default: + MODELOPT_FORMAT_FAIL("Unsupported format for E8M0 scaling"); + } + return -1; +} + +inline int ceil_div(const int a, const int b) { return (a + b - 1) / b; } diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py new file mode 100644 index 0000000000..b4923fb661 --- /dev/null +++ b/modelopt/torch/quantization/tensor_quant.py @@ -0,0 +1,653 @@ +# 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. + +"""Basic tensor quantization functions.""" + +import torch +from packaging.version import Version +from torch.autograd import Function +from torch.onnx import symbolic_helper + +from .config import QuantizerAttributeConfig +from .export_onnx import export_fp4, export_fp8, export_int8 +from .extensions import get_cuda_ext, get_cuda_ext_fp8, get_cuda_ext_mx +from .utils import is_torch_library_supported + +mx_format_map = { + (4, 3): "E4M3", + (5, 2): "E5M2", + (3, 2): "E3M2", + (2, 3): "E2M3", + 8: "INT8", + (8, 0): "E8M0", + (2, 1): "E2M1", + (1, 2): "E1M2", + (0, 3): "E0M3", + (3, 0): "E3M0", +} + + +def scaled_e4m3_impl( + inputs: torch.Tensor, # TODO: check support for multiple inputs + amax: torch.Tensor, + disable_fused_kernel=True, +) -> torch.Tensor: + """Implementation of fake quantizing input to FP8. + + Args: + inputs: Torch tensor. + amax: Absolute max range of the input tensor. + + Returns: + Input tensors faked quantized to FP8. + """ + cuda_ext_fp8 = get_cuda_ext_fp8(raise_if_failed=True) + + def is_fusable(): + # ignore no scaling and shape([]) cases + if amax is None or len(amax.shape) == 0: + return False + else: + # can't have amax.shape = [1, 1, 4, 1] and the like + amax_last_dim_only = amax.numel() == amax.shape[-1] + # must be cuda + all_cuda = inputs.is_cuda and amax.is_cuda + + # also check explicit disable. + return amax_last_dim_only and all_cuda and (not disable_fused_kernel) + + with torch.cuda.device( + None if inputs.device.index == torch.cuda.current_device() else inputs.device.index + ): + # differentiate between fused & unfused cases + if is_fusable(): + zero_threshold = 1.0 / (1 << 24) + outputs = cuda_ext_fp8.fused_fake_e4m3fy(inputs, amax.float(), zero_threshold) + else: + zero_mask = inputs.abs() < 1.0 / (1 << 24) + + if amax is None: + outputs = cuda_ext_fp8.fake_e4m3fy(inputs) + else: + scale = 448.0 / amax + outputs = cuda_ext_fp8.fake_e4m3fy(inputs * scale) / scale + + # Zero out values that are tiny. + # Tiny values could lead to tiny amax and then large scale which cause overflow/saturation + # and won't go back to normal value after dividing by scale. The right behavior is to mark them + # as zero which also get rid of inf/nan + outputs[zero_mask] = 0.0 + + return outputs + + +def fake_quant_impl( + inputs: torch.Tensor, + amax: torch.Tensor, + num_bits=8, + unsigned=False, + narrow_range=True, +): + """Implementation of fake quantizing input according to number of bits.""" + cuda_ext = get_cuda_ext() + + with torch.cuda.device( + None if inputs.device.index == torch.cuda.current_device() else inputs.device.index + ): + if amax.numel() == 1: + outputs = cuda_ext.fake_tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) + else: + axis = amax.shape.index(amax.numel()) + outputs = cuda_ext.fake_tensor_quant_with_axis( + inputs, amax.squeeze(), axis, num_bits, unsigned, narrow_range + ) + return outputs + + +def _quantize_impl( + inputs: torch.Tensor, + amax: torch.Tensor, + num_bits: int = 8, + exponent_bits: int = 0, + unsigned: bool = False, + narrow_range: bool = True, +): + if num_bits == 8 and exponent_bits == 4: + return scaled_e4m3_impl(inputs=inputs, amax=amax) + elif isinstance(num_bits, int): + return fake_quant_impl( + inputs=inputs, + amax=amax, + num_bits=num_bits, + unsigned=unsigned, + narrow_range=narrow_range, + ) + else: + raise ValueError( + f"Invalid combination of (num_bits, exponent_bits): ({num_bits}, {exponent_bits})." + ) + + +def _quantize_impl_abstract( + input: torch.Tensor, + amax: torch.Tensor, + num_bits: int = 8, + exponent_bits: int = 0, + unsigned: bool = False, + narrow_range: bool = True, +) -> torch.Tensor: + """Register an abstract implementation for quantizing tensor. + + This abstract function returns an empty tensor with the same shape and dtype. + """ + output = torch.empty_like(input) + + return output + + +def _dynamic_block_quantize_impl( + inputs: torch.Tensor, + block_size: int, + amax: torch.Tensor, + num_bits: int, + exponent_bits: int, + scale_num_bits: int, + scale_exponent_bits: int, +): + scale_bits = (scale_exponent_bits, scale_num_bits - scale_exponent_bits - 1) + num_bits = (exponent_bits, num_bits - exponent_bits - 1) + if num_bits in mx_format_map: + assert scale_bits in mx_format_map, f"Scale bits should be in {mx_format_map.keys()}" + if scale_bits != (8, 0): + assert amax.is_cuda, "amax must be a CUDA tensor for dynamic block quantization." + if amax.numel() != 1: + amax = amax.amax() + cuda_ext_mx = get_cuda_ext_mx(raise_if_failed=True) + with torch.cuda.device( + None if inputs.device.index == torch.cuda.current_device() else inputs.device.index + ): + return cuda_ext_mx.fused_amax_convert( + inputs, + block_size, + getattr(cuda_ext_mx.Types, mx_format_map[num_bits]), + getattr(cuda_ext_mx.Types, mx_format_map[scale_bits]), + amax, + ) + else: + raise NotImplementedError("Dynamic block quantization only support E2M1 for now.") + + +def _dynamic_block_quantize_impl_abstract( + inputs: torch.Tensor, + block_size: int, + amax: torch.Tensor, + num_bits: int, + exponent_bits: int, + scale_num_bits: int, + scale_exponent_bits: int, +): + """Register an abstract implementation for dynamic block quantization. + + This abstract function returns an empty tensor with the same shape and dtype. + """ + output = torch.empty_like(inputs) + + return output + + +quantize_op = _quantize_impl +dynamic_block_quantize_op = _dynamic_block_quantize_impl +# Define custom operators via torch.library if supported: +# 1. quantize_op: Applies static quantization to the input tensor using a specified amax. +# 2. dynamic_block_quantize_op: Performs blockwise double quantization with dynamically +# determined scales, governed by the given quantization format (scale_num_bits, scale_exponent_bits). +if is_torch_library_supported(): + try: + torch.library.define( + "tensorrt::quantize_op", + "(Tensor input, Tensor amax, int num_bits, int exponent_bits, " + "bool unsigned, bool narrow_range) -> Tensor", + ) + torch.library.define( + "tensorrt::dynamic_block_quantize_op", + "(Tensor input, int block_size, Tensor amax, int num_bits, int exponent_bits, " + "int scale_num_bits, int scale_exponent_bits) -> Tensor", + ) + torch.library.impl("tensorrt::quantize_op", ["cpu", "cuda"])(_quantize_impl) + torch.library.impl("tensorrt::dynamic_block_quantize_op", ["cpu", "cuda"])( + _dynamic_block_quantize_impl + ) + if Version(torch.__version__) < Version("2.4.0"): + torch.library.impl_abstract("tensorrt::quantize_op")(_quantize_impl_abstract) + torch.library.impl_abstract("tensorrt::dynamic_block_quantize_op")( + _dynamic_block_quantize_impl_abstract + ) + else: + torch.library.register_fake("tensorrt::quantize_op")(_quantize_impl_abstract) + torch.library.register_fake("tensorrt::dynamic_block_quantize_op")( + _dynamic_block_quantize_impl_abstract + ) + quantize_op = torch.ops.tensorrt.quantize_op + dynamic_block_quantize_op = torch.ops.tensorrt.dynamic_block_quantize_op + except (AttributeError, RuntimeError): + # torch.library is an experiemental feature, the function signatures may change overtime. + print( + "Unable to register operators with torch.library. Exporting quantized models with" + " torch.export will not be supported." + ) + +# Predefined descriptors +QUANT_DESC_8BIT_PER_TENSOR = QuantizerAttributeConfig(num_bits=8) +QUANT_DESC_UNSIGNED_8BIT_PER_TENSOR = QuantizerAttributeConfig(num_bits=8, unsigned=True) +QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0)) +QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0)) +QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0)) +QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW = QuantizerAttributeConfig(num_bits=8, axis=(0)) +QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0)) +QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0)) +QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL = QuantizerAttributeConfig(num_bits=8, axis=(0)) + + +@torch.jit.script +def _fake_tensor_quant_backward(inputs, amax, grad_outputs): + zero = grad_outputs.new_zeros(1) + grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero) + return grad_inputs + + +class FakeTensorQuantFunction(Function): + """Fake version of TensorQuantFunction use CUDA extension.""" + + @staticmethod + @symbolic_helper.parse_args("v", "t", "i", "b", "b", "s") + def symbolic( + g, + inputs, + amax, + num_bits=8, + unsigned=False, + narrow_range=True, + trt_high_precision_dtype="Float", + ): + """ONNX symbolic function.""" + return export_int8( + g, inputs, amax, num_bits, unsigned, narrow_range, trt_high_precision_dtype + ) + + @staticmethod + def forward( + ctx, + inputs, + amax, + num_bits=8, + unsigned=False, + narrow_range=True, + trt_high_precision_dtype="Float", + ): + """Forward method.""" + ctx.save_for_backward(inputs, amax) + + def legacy_quant_func(): + # The LegacyFakeTensorQuantFunction support cpu and amax with any shape that can be broadcasted to inputs. + outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) + return outputs / scale.to(inputs.dtype) + + if not inputs.is_cuda: + outputs = legacy_quant_func() + else: + try: + outputs = quantize_op( + inputs, + amax, + num_bits=num_bits, + exponent_bits=0, + unsigned=unsigned, + narrow_range=narrow_range, + ) + except (AttributeError, ValueError): + # AttributeError: cuda_ext is not imported, possibly due to CPU only installation + # ValueError: cuda_ext is installed, but trying to perform multidimensional quantization (amax dim > 1) + outputs = legacy_quant_func() + + return outputs + + @staticmethod + def backward(ctx, grad_outputs): + """Implements straight through estimation with clipping.""" + inputs, amax = ctx.saved_tensors + return _fake_tensor_quant_backward(inputs, amax, grad_outputs), None, None, None, None, None + + +class ScaledE4M3Function(Function): + """E4M3fy input with scale.""" + + @staticmethod + @symbolic_helper.parse_args("v", "t", "i", "b", "b", "s") + def symbolic(g, inputs, amax=None, E=4, M=3, trt_high_precision_dtype="Float"): # noqa: N803 + """ONNX symbolic function.""" + return export_fp8(g, inputs, amax, trt_high_precision_dtype) + + @staticmethod + # Default values could cause errors from TorchDynamo during torch.export + def forward(ctx, inputs, amax, E, M, trt_high_precision_dtype="Float"): # noqa: N803 + """Forward method.""" + if E != 4 or M != 3: + raise NotImplementedError("Only support E=4 & M=3 for now.") + + ctx.save_for_backward(inputs) + ctx.amax = amax + outputs = quantize_op( + inputs, amax, num_bits=8, exponent_bits=4, unsigned=False, narrow_range=False + ) + + return outputs + + @staticmethod + def backward(ctx, grad_outputs): + """Implements straight through estimation with clipping.""" + (inputs,) = ctx.saved_tensors + amax = torch.tensor( + ctx.amax if ctx.amax is not None else 448.0, dtype=torch.float32, device=inputs.device + ) + grad_inputs = _fake_tensor_quant_backward(inputs, amax, grad_outputs) + return grad_inputs, None, None, None, None + + +class DynamicBlockQuantizationFunction(Function): + """Dynamic block quantization functional.""" + + @staticmethod + @symbolic_helper.parse_args("v", "i", "v", "is", "is", "s", "s") + def symbolic( + g, + inputs, + block_size, + amax, + num_bits, + scale_bits, + trt_high_precision_dtype="Half", + onnx_quantizer_type="dynamic", + ): # noqa: N803 + """ONNX symbolic function.""" + return export_fp4( + g, + inputs, + block_size, + amax, + num_bits, + trt_high_precision_dtype, + onnx_quantizer_type, + ) + + @staticmethod + def forward( + ctx, + inputs, + block_size, + amax, + num_bits, + scale_bits, + trt_high_precision_dtype="Half", + onnx_quantizer_type="dynamic", + ): + """Forward method.""" + ctx.save_for_backward(inputs, amax) + + assert isinstance(num_bits, tuple) and len(num_bits) == 2 + exponent_bits = num_bits[0] + num_bits = num_bits[0] + num_bits[1] + 1 + assert isinstance(scale_bits, tuple) and len(scale_bits) == 2 + scale_exponent_bits = scale_bits[0] + scale_num_bits = scale_bits[0] + scale_bits[1] + 1 + outputs = dynamic_block_quantize_op( + inputs, block_size, amax, num_bits, exponent_bits, scale_num_bits, scale_exponent_bits + ) + return outputs + + @staticmethod + def backward(ctx, grad_outputs): + """Implements straight through estimation with clipping.""" + # Dynamic block quantization does not clip the tensors (i.e in effect) in the forward loop + # as the block quantization amax is always the block amax. So we dont need to clip the gradients for STE + inputs, amax = ctx.saved_tensors + return ( + _fake_tensor_quant_backward(inputs, amax, grad_outputs), + None, + None, + None, + None, + None, + None, + ) + + +class TensorQuantFunction(Function): + """A universal tensor quantization function. + + Take an input tensor, output an quantized tensor. The granularity of scale can be interpreted from the + shape of amax. + output_dtype indicates whether the quantized value will be stored in integer or float. The reason we want to store + it in float is the pytorch function takes the quantized value may not accept integer input, e.g. Conv2D. + + It uses 2^num_bits -1 values instead of 2^num_bits. e.g., for num_bits=8, it uses [-127, 127] instead of [-128, 127] + """ + + @staticmethod + @symbolic_helper.parse_args("v", "t", "i", "b", "b", "s") + def symbolic( + g, + inputs, + amax, + num_bits=8, + unsigned=False, + narrow_range=True, + trt_high_precision_dtype="Float", + ): + """ONNX symbolic function.""" + return export_int8( + g, inputs, amax, num_bits, unsigned, narrow_range, trt_high_precision_dtype + ) + + @staticmethod + def forward( + ctx, + inputs, + amax, + num_bits=8, + unsigned=False, + narrow_range=True, + trt_high_precision_dtype="Float", + ): + """Forward method. + + Follow tensorflow convention, max value is passed in and used to decide scale, instead of inputing scale + directly. Though inputing scale directly may be more natural to use. + + Args: + ctx: A Context object to store tensors for backward. + inputs: A Tensor of type float32. + amax: A Tensor of type float32. Inputs will be quantized within range [-amax, amax] + amax will be broadcasted to inputs tensor. + num_bits: A integer used to calculate scaling factor, scale = (2^(num_bits-1) - 1) / max + Effectively, it indicates how many integer bits is used to represent the value. Default 8. + output_dtype: A type of Tensor. torch.int32 or torch.float32. + unsigned: A boolean. Use unsigned integer range. E.g. [0, 255] for num_bits=8. Default False. + narrow_range: A boolean. Use symmetric integer range for signed quantization + E.g. [-127,127] instead of [-128,127] for num_bits=8. Default True. + + Returns: + outputs: A Tensor of type output_dtype. + scale: A Tensor of type float32. outputs / scale will dequantize outputs tensor. + + Raises: + ValueError: + """ + ctx.save_for_backward(inputs, amax) + outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) + # Check if scale overflows FP16 + if outputs.dtype == torch.half and scale.max() > 65504: + raise ValueError(f"scale is too large for FP16 with amax={amax}") + return outputs, scale.to(inputs.dtype) + + @staticmethod + def backward(ctx, grad_outputs, grad_scale): + """Implements straight through estimation with clipping. + + For -amax <= input <= amax the gradient passes straight through, otherwise the gradient is zero. + + Args: + ctx: A Context object with saved tensors from forward. + grad_outputs: A tensor of gradient of outputs. + grad_scale: A tensor of gradient of scale. + + Returns: + grad_inputs: A tensor of gradient. + """ + inputs, amax = ctx.saved_tensors + zero = grad_outputs.new_zeros(1) # create a zero tensor with the same type and device + grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero) + return grad_inputs, None, None, None, None, None + + +class LegacyFakeTensorQuantFunction(Function): + """Fake version of TensorQuantFunction. + + See comments of TensorQuantFunction, arguments are the same. + """ + + @staticmethod + def forward(ctx, inputs, amax, num_bits=8, unsigned=False, narrow_range=True): + """Forward method.""" + ctx.save_for_backward(inputs, amax) + outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) + return outputs / scale.to(inputs.dtype) + + @staticmethod + def backward(ctx, grad_outputs): + """Implements straight through estimation.""" + inputs, amax = ctx.saved_tensors + zero = grad_outputs.new_zeros(1) + grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero) + return grad_inputs, None, None, None, None + + +def _tensor_quant(inputs, amax, num_bits=8, unsigned=False, narrow_range=True): + """Shared function body between TensorQuantFunction and FakeTensorQuantFunction.""" + # Fine scale, per channel scale will be handled by broadcasting, which could be tricky. Pop a warning. + if unsigned: + if inputs.min() < 0.0: + raise TypeError("Negative values encountered in unsigned quantization.") + + # Computation can be done in FP32 to prevent potential over flow. + input_dtype = inputs.dtype + if inputs.dtype == torch.half: + inputs = inputs.float() + if amax.dtype == torch.half: + amax = amax.float() + + min_amax = amax.min() + if min_amax < 0: + raise ValueError("Negative values in amax") + + max_bound = torch.tensor((2.0 ** (num_bits - 1 + int(unsigned))) - 1.0, device=amax.device) + if unsigned: + min_bound = 0 + elif narrow_range: + min_bound = -max_bound + else: + min_bound = -max_bound - 1 + scale = max_bound / amax + + epsilon = 1.0 / (1 << 24) + if min_amax <= epsilon: # Treat amax smaller than minimum representable of fp16 0 + zero_amax_mask = amax <= epsilon + scale[zero_amax_mask] = 0 # Value quantized with amax=0 should all be 0 + + outputs = torch.clamp((inputs * scale).round_(), min_bound, max_bound) + + if min_amax <= epsilon: + scale[zero_amax_mask] = ( + 1.0 # Return 1 makes more sense for values quantized to 0 with amax=0 + ) + + if input_dtype == torch.half: + outputs = outputs.half() + + return outputs, scale + + +class FakeAffineTensorQuantFunction(Function): + """Fake version of affine quantization. + + gemmlowp style scale+shift quantization. See more details in + https://github.com/google/gemmlowp/blob/master/doc/quantization.md. + + We DO NOT recommend affine quantization on weights for performance reason. There might be value to affine quantize + activation as it can be cancelled by bias and comes with no performance penalty. This functionality is only added + for experimental purpose. + """ + + @staticmethod + def forward(ctx, inputs, min_range, max_range, num_bits=8): + """As it will be only applied on activation with per tensor granularity, broadcast is not needed. + + Args: + ctx: Pytorch convention. + inputs: A Tensor of type float32. + min_range: A float. + max_range: A float. + num_bits: An integer + + Returns: + outputs: A Tensor of type output_dtype + """ + ctx.save_for_backward(inputs, min_range, max_range) + + step_size = (max_range - min_range) / (2.0**num_bits - 1) + + min_bound = -(2.0 ** (num_bits - 1)) + max_bound = 2.0 ** (num_bits - 1) - 1 + + quant_zero = torch.round(min_range / step_size) - min_bound + quantized = torch.round(inputs / step_size) - quant_zero + quantized = torch.clamp(quantized, min_bound, max_bound) + + outputs = (quantized + quant_zero) * step_size + + return outputs + + @staticmethod + def backward(ctx, grad_outputs): + """Implements straight through estimation with clipping. + + Args: + ctx: Pytorch convention. + grad_output: A tensor of gradient of outputs. + + Returns: + grad_inputs: A tensor of gradient + """ + inputs, min_range, max_range = ctx.saved_tensors + zero = grad_outputs.new_zeros(1) + grad_inputs = torch.where((inputs <= max_range) * (inputs >= min_range), grad_outputs, zero) + return grad_inputs, None, None, None + + +tensor_quant = TensorQuantFunction.apply +legacy_fake_tensor_quant = LegacyFakeTensorQuantFunction.apply +fake_tensor_quant = FakeTensorQuantFunction.apply +fake_affine_tensor_quant = FakeAffineTensorQuantFunction.apply +scaled_e4m3 = ScaledE4M3Function.apply +dynamic_block_quant = DynamicBlockQuantizationFunction.apply diff --git a/modelopt/torch/quantization/utils.py b/modelopt/torch/quantization/utils.py new file mode 100644 index 0000000000..66c2b17088 --- /dev/null +++ b/modelopt/torch/quantization/utils.py @@ -0,0 +1,184 @@ +# 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. + +"""Quantization utilities.""" + +from contextlib import ExitStack, contextmanager + +import torch +from packaging.version import Version + +from modelopt.torch.utils.distributed import ParallelState + +__all__ = [ + "reduce_amax", + "is_quantized", + "is_quantized_layer_with_weight", + "is_quantized_column_parallel_linear", + "is_quantized_row_parallel_linear", + "replace_function", + "EXPORT_MODE", + "export_torch_mode", + "is_torch_library_supported", + "get_parallel_state", +] + + +def reduce_amax(input, axis=None, keepdims=True): + """Compute the absolute maximum value of a tensor. + + Reduces input_tensor along the dimensions given in axis. Unless keepdims is true, + the rank of the tensor is reduced by 1 for each entry in axis. If keepdims is true, + the reduced dimensions are retained with length 1. + + .. note:: + Gradient computation is disabled as this function is never meant learning reduces amax + + Args: + input: Input tensor + axis: The dimensions to reduce. None or int or tuple of ints. If None (the default), + reduces all dimensions. Must be in the range [-rank(input_tensor), rank(input_tensor)). + keepdims: A boolean. If true, retains reduced dimensions with length 1. Default True + granularity: DEPRECTED. specifies if the statistic has to be calculated at tensor or channel granularity + + Returns: + The reduced tensor. + + Raises: + ValueError: Any axis which doesn't make sense or is not supported + ValueError: If unknown granularity is passed in. + """ + with torch.no_grad(): + # A memory-efficient implementation that avoids copying input tensor + if axis is None: + max_val = torch.max(input) + min_val = torch.min(input) + output = torch.maximum(torch.abs(max_val), torch.abs(min_val)) + else: + if isinstance(axis, int): + axis = (axis,) + max_val = torch.amax(input, dim=axis, keepdim=keepdims) + min_val = torch.amin(input, dim=axis, keepdim=keepdims) + output = torch.maximum(torch.abs(max_val), torch.abs(min_val)) + if output.numel() == 1: + output.squeeze_() + return output + + +def is_quantized(module): + """Check if a module is quantized.""" + from .nn import TensorQuantizer + + for _module in module.modules(): + if isinstance(_module, TensorQuantizer): + return True + return False + + +def is_quantized_layer_with_weight(module): + """Check if a module is quantized with weights.""" + return is_quantized(module) and getattr(module, "weight", None) is not None + + +def is_quantized_linear(module): + """Check if a module is a quantized linear module.""" + return ( + hasattr(module, "input_quantizer") + and hasattr(module, "weight_quantizer") + and getattr(module, "weight", None) is not None + and module.weight.dim() == 2 + ) + + +def is_quantized_column_parallel_linear(module): + """Check if a module is a quantized column parallel linear module.""" + return is_quantized_linear(module) and getattr(module, "_is_column_parallel", False) + + +def is_quantized_row_parallel_linear(module): + """Check if a module is a quantized row parallel linear module.""" + return is_quantized_linear(module) and getattr(module, "_is_row_parallel", False) + + +@contextmanager +def replace_function(package, name, new_func): + """Replace a function with a new one within a context.""" + old_func = getattr(package, name) + setattr(package, name, new_func) + setattr(package, "_" + name, old_func) + yield + setattr(package, name, old_func) + delattr(package, "_" + name) + + +@contextmanager +def multi_context(*cms): + """Context manager enabling variable number of context managers.""" + with ExitStack() as stack: + yield [stack.enter_context(cls) for cls in cms] + + +EXPORT_MODE: bool = False + + +@contextmanager +def export_torch_mode(): + """Context manager enabling the export mode.""" + global EXPORT_MODE + original_value = EXPORT_MODE + EXPORT_MODE = True + try: + yield + finally: + EXPORT_MODE = original_value + + +def is_torch_export_mode(): + """Check whether in the context of exporting model to torch.""" + return EXPORT_MODE + + +def is_torch_library_supported(): + """Check if the installed PyTorch version meets or exceeds a specified version.""" + # Require torch version >= 2.2.0 + # Adding checks for `impl` and `impl_abstract` as they are experimental features + return ( + Version(torch.__version__) >= Version("2.2.0") + and hasattr(torch.library, "impl") + and hasattr(torch.library, "impl_abstract") + ) or ( + Version(torch.__version__) >= Version("2.4.0") + and hasattr(torch.library, "impl") + and hasattr(torch.library, "register_fake") + ) + + +def get_parallel_state(model, name=None) -> ParallelState: + """Get the parallel state. + + Args: + model: Pytorch model. + name: The name of the submodule of the model to get the parallel state from. If None, + the parallel state of the model is returned. + """ + if name is None: + return getattr(model, "_parallel_state", ParallelState()) + + # If the submodule does not have a parallel state, get the parallel state of the parent module + module = model.get_submodule(name) + if hasattr(module, "_parallel_state"): + return module._parallel_state + parent_module = model.get_submodule(name.rpartition(".")[0]) + return getattr(parent_module, "_parallel_state", ParallelState()) diff --git a/modelopt/torch/sparsity/__init__.py b/modelopt/torch/sparsity/__init__.py new file mode 100644 index 0000000000..2ac26b8de5 --- /dev/null +++ b/modelopt/torch/sparsity/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""API for sparsification algorithms.""" + +from . import mode, module, plugins +from .sparsification import * diff --git a/modelopt/torch/sparsity/config.py b/modelopt/torch/sparsity/config.py new file mode 100644 index 0000000000..b88533885a --- /dev/null +++ b/modelopt/torch/sparsity/config.py @@ -0,0 +1,51 @@ +# 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. + +"""Default configurations for sparsity modes.""" + +from pydantic import create_model + +from modelopt.torch.opt.config import ModeloptBaseConfig, get_kwargs_for_create_model_with_rules + +from .module import SpDMRegistry + +SparseMagnitudeConfig = create_model( + "SparseMagnitudeConfig", + **get_kwargs_for_create_model_with_rules( + registry=SpDMRegistry, + default_rules={ + "nn.Linear": {"*": {}, "*lm_head*": None}, + "nn.Conv2d": {"*": {}, "*lm_head*": None}, + }, + doc='Configuration for the ``"sparse_magnitude"`` mode.', + ), +) + + +SparseGPTConfig = create_model( + "SparseGPTConfig", + **get_kwargs_for_create_model_with_rules( + registry=SpDMRegistry, + default_rules={ + "nn.Linear": {"*": {}, "*lm_head*": None}, + "nn.Conv2d": {"*": {}, "*lm_head*": None}, + }, + doc='Configuration for the ``"sparse_gpt"`` mode.', + ), +) + + +class ExportSparseConfig(ModeloptBaseConfig): + """Configuration (empty!) for the ``"export_sparse"`` mode.""" diff --git a/modelopt/torch/sparsity/magnitude.py b/modelopt/torch/sparsity/magnitude.py new file mode 100644 index 0000000000..21fcc97bdb --- /dev/null +++ b/modelopt/torch/sparsity/magnitude.py @@ -0,0 +1,148 @@ +# 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. + +"""Magnitude-base sparsity inspired by NVIDIA ASP (Automatic SParsity).""" + +import re +import warnings +from itertools import permutations + +import torch +import torch.nn as nn + +from .module import SparseModule +from .searcher import BaseSparseSearcher + + +def get_nmprune_info(pattern: str) -> tuple[bool, int, int]: + """Gets the n:m sparsity pattern information from a given string.""" + nm_prune = re.search(r"(\d+):(\d+) sparsity", pattern) + if nm_prune is not None: + n, m = map(int, nm_prune.groups()) + return nm_prune is not None, n, m + return False, 0, 0 + + +def fill(x): + """Calculates the ratio of non-zero elements in a tensor.""" + return float(x.nonzero().size(0)) / torch.numel(x) + + +def reshape_1d(matrix, m): + """Reshapes a given matrix into m-dimensional vectors: (h,w) -> (hw/m, m).""" + if matrix.shape[1] % m > 0: + new_cols = matrix.shape[1] + (m - matrix.shape[1] % m) + mat = matrix.new_empty(matrix.shape[0], new_cols).fill_(0) + mat[:, : matrix.shape[1]] = matrix + + return mat.view(-1, m), mat.shape + else: + return matrix.view(-1, m), matrix.shape + + +def compute_valid_1d_patterns(m, n): + """Computes all possible m:n patterns in a 1D vector. + + The function generates a tensor of size m with n ones and (m-n) zeros. + It then generates all permutations of this tensor, removes duplicates, + and returns the unique patterns as a tensor. + """ + patterns = torch.zeros(m) + patterns[:n] = 1 + valid_patterns = torch.tensor(list(set(permutations(patterns.tolist())))) + return valid_patterns + + +def mn_1d_best(matrix, m, n): + """Finds the best m:n pattern in a given matrix. + + The function computes all possible m:n patterns and selects the one + that maximizes the sum of non-masked weights in the matrix. The selected + pattern is then used to create a mask for the matrix. + """ + patterns = compute_valid_1d_patterns(m, n).to(matrix.device) + + # Find the best m:n pattern (sum of non-masked weights). + mask = torch.IntTensor(matrix.shape).fill_(1).view(-1, m) + mat, _ = reshape_1d(matrix, m) + pmax = torch.argmax(torch.matmul(mat.abs(), patterns.t()), dim=1) + mask[:] = patterns[pmax[:]] + mask = mask.view(matrix.shape) + return mask + + +def m4n2_1d(mat): + """Finds the best 2:4 pattern in a given matrix.""" + return mn_1d_best(mat, 4, 2) + + +def create_asp_mask(tensor: nn.Parameter, pattern: str) -> torch.BoolTensor: + """Creates a mask for a given tensor based on a specified sparse pattern. + + The function reshapes the tensor and applies the specified pattern to create a sparse mask. + The default pattern is m4n2_1d, which finds the best 2:4 sparsity pattern in the tensor. + """ + pattern_method_lut = {BaseSparseSearcher._pattern_2_4: m4n2_1d} + if pattern not in pattern_method_lut: + raise NotImplementedError(f"Unsupported pattern {pattern} for ASP sparsity") + func = pattern_method_lut[pattern] + + shape = tensor.shape + tensor.type() + t = tensor.float().contiguous() + + # 1d-tensor + if len(shape) == 1: + t = t.view(1, shape[0]) + mask = func(t) + # 2d-tensor (K, C) + elif len(shape) == 2: + # linear + t = t.view(shape[0], shape[1]) + mask = func(t) + # 3d-tensor (K, C, R) + elif len(shape) == 3: + # 1d convs + t = t.permute(0, 2, 1).contiguous().view(shape[0] * shape[2], shape[1]) + mask = func(t) + mask = mask.view(shape[0], shape[2], shape[1]).permute(0, 2, 1).contiguous() + # 4d-tensor (K, C, R, S) + elif len(shape) == 4: + # 2d convs + t = t.permute(2, 3, 0, 1).contiguous().view(shape[2] * shape[3] * shape[0], shape[1]) + mask = func(t) + mask = mask.view(shape[2], shape[3], shape[0], shape[1]).permute(2, 3, 0, 1).contiguous() + + return mask.view(shape).to(dtype=torch.bool) + + +class MagnitudeSearcher(BaseSparseSearcher): + """Searcher for magnitude-based sparsity.""" + + def _check_weight_size(self, weight: torch.nn.Parameter, mod_name: str) -> bool: + """Check if the weight size is supported.""" + # rules from ASP + if weight.size(0) % 8 != 0 or weight.size(1) % 16 != 0: + warnings.warn( + f"Skipping pruning {mod_name} of size={str(weight.size())} and" + f" type={str(weight.dtype)} for sparsity" + ) + return False + + return True + + def _compute_mask(self, module: SparseModule) -> torch.BoolTensor: + """Compute the mask (and weight update) for the given module.""" + return create_asp_mask(module.weight, self.config["pattern"]) diff --git a/modelopt/torch/sparsity/mode.py b/modelopt/torch/sparsity/mode.py new file mode 100644 index 0000000000..dda147049e --- /dev/null +++ b/modelopt/torch/sparsity/mode.py @@ -0,0 +1,204 @@ +# 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. + +"""Sparsity mode descriptor.""" + +from typing import Optional + +from torch import nn + +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.conversion import ApplyModeError +from modelopt.torch.opt.dynamic import DynamicSpace +from modelopt.torch.opt.mode import ( + ConvertEntrypoint, + ConvertReturnType, + MetadataDict, + RestoreEntrypoint, + UpdateEntrypoint, + _ModeDescriptor, + _ModeRegistryCls, +) +from modelopt.torch.opt.searcher import BaseSearcher +from modelopt.torch.utils import compare_dict, unwrap_model + +from .config import ExportSparseConfig, SparseGPTConfig, SparseMagnitudeConfig +from .magnitude import MagnitudeSearcher +from .module import SpDMRegistry +from .sparsegpt import SparseGPTSearcher + +SparsityModeRegistry = _ModeRegistryCls() + + +def convert_sparse_model(model: nn.Module, config: ModeloptBaseConfig) -> ConvertReturnType: + """Function for converting a model to a sparsity meta-model.""" + # we use the search space utility here with a custom registry to convert the model + dynamic_space = DynamicSpace(model) + dynamic_space.convert_to_dynamic(config.model_dump(), SpDMRegistry) + + return dynamic_space.model, {"subnet_config": DynamicSpace(model).config()} + + +def restore_sparse_model( + model: nn.Module, config: ModeloptBaseConfig, metadata: MetadataDict +) -> nn.Module: + """Function for restoring a previously convert model to a sparsity meta-model.""" + assert "subnet_config" in metadata + model, _ = convert_sparse_model(model, config) + DynamicSpace(model).select(metadata["subnet_config"]) + + return model + + +def update_sparse_metadata( + model: nn.Module, config: ModeloptBaseConfig, metadata: MetadataDict +) -> None: + """Update subnet config to current subnet config of model.""" + metadata["subnet_config"] = DynamicSpace(model).config() + + +def export_sparse(model: nn.Module, config: ExportSparseConfig) -> ConvertReturnType: + """Export a sparse model to a regular model.""" + # sanity check to avoid DP/DDP here in the entrypoint + model = unwrap_model(model, raise_error=True) + + # store config from model if we can find it for a future convert/restore process + metadata = {"subnet_config": DynamicSpace(model).config()} + + # export model in-place + model = DynamicSpace(model).export(SpDMRegistry) + + return model, metadata + + +def restore_export_sparse( + model: nn.Module, config: ExportSparseConfig, metadata: MetadataDict +) -> nn.Module: + """Restore & export a sparse model to a regular model.""" + # select activated/deactivated sparse modules + DynamicSpace(model).select(metadata["subnet_config"]) + + # run export + model, metadata_new = export_sparse(model, config) + + # double check metadata + unmatched_keys = compare_dict(metadata, metadata_new) + if unmatched_keys: + raise ApplyModeError(f"Unmatched metadata={unmatched_keys}!") + + return model + + +@SparsityModeRegistry.register_mode +class SparseMagnitudeModeDescriptor(_ModeDescriptor): + """Class to define and describe magnitude-based sparsification.""" + + @property + def name(self) -> str: + """Returns the name of the mode.""" + return "sparse_magnitude" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return SparseMagnitudeConfig + + @property + def next_modes(self) -> Optional[set[str]]: + """Specifies the next modes for the mode.""" + return {"export_sparse", "kd_loss", "quantize"} + + @property + def export_mode(self) -> Optional[str]: + """The mode that corresponds to the export mode of this mode.""" + return "export_sparse" + + @property + def search_algorithm(self) -> type[BaseSearcher]: + """Specifies the search algorithm for the mode.""" + return MagnitudeSearcher + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return convert_sparse_model + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_sparse_model + + @property + def update_for_save(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models metadata.""" + return update_sparse_metadata + + @property + def update_for_new_mode(self) -> UpdateEntrypoint: + """The mode's entrypoint for updating the models metadata.""" + return update_sparse_metadata + + +@SparsityModeRegistry.register_mode +class SparseGPTModeDescriptor(SparseMagnitudeModeDescriptor): + """Class to define and describe sparsification based on SparseGPT.""" + + @property + def name(self) -> str: + """Returns the name of the mode.""" + return "sparsegpt" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return SparseGPTConfig + + @property + def search_algorithm(self) -> type[BaseSearcher]: + """Specifies the search algorithm for the mode.""" + return SparseGPTSearcher + + +@SparsityModeRegistry.register_mode +class ExportSparseModeDescriptor(_ModeDescriptor): + """Class to describe the ``"export_sparse"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "export_sparse" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return ExportSparseConfig + + @property + def is_export_mode(self) -> bool: + """Specifies if this mode is an export mode.""" + return True + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return export_sparse + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_export_sparse diff --git a/modelopt/torch/sparsity/module.py b/modelopt/torch/sparsity/module.py new file mode 100644 index 0000000000..ee380d126a --- /dev/null +++ b/modelopt/torch/sparsity/module.py @@ -0,0 +1,89 @@ +# 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. + +"""Dynamic class for all sparse modules.""" + +from typing import Optional + +import torch +from torch import nn + +from modelopt.torch.opt.dynamic import DynamicModule, _DMRegistryCls +from modelopt.torch.opt.hparam import Hparam + +__all__ = ["SparseModule", "SpDMRegistry"] + + +SpDMRegistry = _DMRegistryCls(prefix="Sparse") # global instance for the sparsity registry + + +@SpDMRegistry.register({nn.Linear: "nn.Linear", nn.Conv2d: "nn.Conv2d"}) +class SparseModule(DynamicModule): + """Base dynamic class for all sparse modules.""" + + @staticmethod + def _get_weight(mod: "SparseModule", weight: torch.Tensor) -> torch.Tensor: + if mod.is_sparse and mod._weight_mask is not None: + masked_weight = weight * mod._weight_mask + # Quick workaround for the custom attribute for Megatron. + # TODO: maybe we need a more general way for customized attributes + if hasattr(weight, "main_grad"): + masked_weight.main_grad = weight.main_grad + return masked_weight + return weight + + def _setup(self): + # define hparam to check if sparsity is activated + hp = Hparam([0, -1], original=0) + hp.active = 0 + self._register_hparam("is_sparse", hp) + + # define the sparse mask here (don't pre-allocate memory to maximize memory savings) + self._register_temp_attribute("_weight_mask", None, lambda m, n, v: m.register_buffer(n, v)) + + # register dynamic attributes of the class + self._register_dynamic_attribute("weight", self._get_weight) + + def modify(self, *args, **kwargs): + """Initialize the sparsity mask when this is called. + + Note that for any module that is not frozen via ``None`` in the rules, this function will be + called. Hence, we use this function to initialize the sparsity mask only when necessary. + """ + hp = self.get_hparam("is_sparse") + if -1 in hp.choices and self._weight_mask is None: + hp.active = -1 + self._weight_mask = torch.ones_like(self.weight, dtype=torch.bool) + + def set_mask(self, value: Optional[torch.BoolTensor]): + """Set the active sparse mask of the module weights.""" + if value is None: + self._weight_mask = None + return + + # sanity checks on the mask + w_shape = self.weight.shape + assert value is not None, "Mask cannot be None." + assert value.shape == w_shape, f"Mask must have shape {w_shape}, got {value.shape} instead." + assert value.dtype == torch.bool, f"Mask must be of type torch.bool, but got {value.dtype}." + + # assign mask + with torch.no_grad(): + if torch.all(value): + self._weight_mask = None + elif self._weight_mask is None: + self._weight_mask = value.detach().clone().to(self.weight.device) + else: + self._weight_mask.copy_(value.to(self._weight_mask.device)) diff --git a/modelopt/torch/sparsity/plugins/__init__.py b/modelopt/torch/sparsity/plugins/__init__.py new file mode 100644 index 0000000000..f1c80b7d13 --- /dev/null +++ b/modelopt/torch/sparsity/plugins/__init__.py @@ -0,0 +1,27 @@ +# 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. + +"""Handles sparsity plugins for third-party modules. + +Currently, we support plugins for + +- :meth:`megatron` + +""" + +from modelopt.torch.utils import import_plugin + +with import_plugin("megatron"): + from .megatron import * diff --git a/modelopt/torch/sparsity/plugins/megatron.py b/modelopt/torch/sparsity/plugins/megatron.py new file mode 100644 index 0000000000..188492f5d1 --- /dev/null +++ b/modelopt/torch/sparsity/plugins/megatron.py @@ -0,0 +1,78 @@ +# 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. + +"""Support megatron parallel linear.""" + +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint + +from ..config import SparseGPTConfig, SparseMagnitudeConfig +from ..module import SparseModule, SpDMRegistry + + +class _MegatronParallelLinear(SparseModule): + def _get_shard_axis_dict(self): + raise NotImplementedError + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets) + + sparse_state_dict, sharded_axis_dict = {}, self._get_shard_axis_dict() + for k, v in self.state_dict(prefix="", keep_vars=True).items(): + if k == "_weight_mask": + sparse_state_dict[k] = v + if sparse_state_dict: + sharded_state_dict.update( + **make_sharded_tensors_for_checkpoint( + sparse_state_dict, prefix, sharded_axis_dict, sharded_offsets + ) + ) + + return sharded_state_dict + + +@SpDMRegistry.register( + {ColumnParallelLinear: "megatron.core.tensor_parallel.layers.ColumnParallelLinear"} +) +class _MegatronColumnParallelLinear(_MegatronParallelLinear): + def _get_shard_axis_dict(self): + return {"_weight_mask": 0} + + +@SpDMRegistry.register( + {RowParallelLinear: "megatron.core.tensor_parallel.layers.RowParallelLinear"} +) +class _MegatronRowParallelLinear(_MegatronParallelLinear): + def _get_shard_axis_dict(self): + return {"_weight_mask": 1} + + +def _get_extra_rules(): + """Get the extra rules for megatron.""" + return { + "megatron.core.tensor_parallel.layers.ColumnParallelLinear": { + "*": {}, + "*output_layer*": None, + }, + "megatron.core.tensor_parallel.layers.RowParallelLinear": { + "*": {}, + "*output_layer*": None, + }, + } + + +# Update the default rules +SparseMagnitudeConfig.register_default(_get_extra_rules()) +SparseGPTConfig.register_default(_get_extra_rules()) diff --git a/modelopt/torch/sparsity/searcher.py b/modelopt/torch/sparsity/searcher.py new file mode 100644 index 0000000000..7d5e574aa7 --- /dev/null +++ b/modelopt/torch/sparsity/searcher.py @@ -0,0 +1,84 @@ +# 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. + +"""Searcher interface for sparsity algorithms.""" + +from abc import abstractmethod +from typing import Iterator, Optional + +import torch +import torch.nn as nn + +from modelopt.torch.opt.searcher import BaseSearcher, SearchConfig, SearchStateDict +from modelopt.torch.utils import print_rank_0 + +from . import magnitude as asp +from .module import SparseModule + + +class BaseSparseSearcher(BaseSearcher): + """A generic sparse mask searching algorithm.""" + + _pattern_2_4 = "2:4 sparsity" + + @property + def default_search_config(self) -> SearchConfig: + """Get the default config for the searcher.""" + return {**super().default_search_config, "pattern": self._pattern_2_4} + + @property + def default_state_dict(self) -> SearchStateDict: + """Return default state dict.""" + return {} + + def sanitize_search_config(self, config: Optional[SearchConfig]) -> SearchStateDict: + """Sanitize the search config dict.""" + config_sanitized = super().sanitize_search_config(config) + + # sanity check of sparsity format + is_nm_prune, n, m = asp.get_nmprune_info(config_sanitized["pattern"]) + assert ( + is_nm_prune and n == 2 and m == 4 + ), f"Unsupported pattern {self.config['pattern']} for sparsity" + + return config_sanitized + + @abstractmethod + def _check_weight_size(self, weight, mod_name) -> bool: + """Check if the weight size is supported by the algorithm.""" + raise NotImplementedError + + @abstractmethod + def _compute_mask(self, module: SparseModule) -> torch.BoolTensor: + """Compute the mask and update weight for a given module.""" + raise NotImplementedError + + def _named_sparsifiable_modules(self) -> Iterator[tuple[str, nn.Module]]: + """Get the named sparsifiable modules.""" + for name, module in self.model.named_modules(): + if ( + isinstance(module, SparseModule) + and module.is_sparse + and self._check_weight_size(module.weight, name) + ): + yield name, module + + def run_search(self): + """Search for sparse mask.""" + for name, module in self._named_sparsifiable_modules(): + # compute the mask (and potentially weight update inside compute_mask) + print_rank_0(f"Searching for sparse mask and weight update for module {name}.") + with torch.no_grad(): + module.set_mask(self._compute_mask(module)) diff --git a/modelopt/torch/sparsity/sparsegpt.py b/modelopt/torch/sparsity/sparsegpt.py new file mode 100644 index 0000000000..c08849e5ef --- /dev/null +++ b/modelopt/torch/sparsity/sparsegpt.py @@ -0,0 +1,276 @@ +# 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. + +"""Utility functions of SparseGPT.""" + +import math +import warnings + +import torch +import torch.nn as nn + +from modelopt.torch.opt.searcher import SearchConfig +from modelopt.torch.utils import print_rank_0 + +from .magnitude import get_nmprune_info +from .module import SparseModule +from .searcher import BaseSparseSearcher + + +def invert(hessian: torch.Tensor) -> torch.Tensor: + """Invert a Hessian matrix.""" + try: + hessian_inv = torch.linalg.cholesky(hessian) + hessian_inv = torch.cholesky_inverse(hessian_inv) + hessian_inv = torch.linalg.cholesky(hessian_inv, upper=True) + except RuntimeError: + cols = hessian.size(0) + eps = 1e-6 * torch.eye(cols).to(hessian.device) + hessian_inv = torch.cholesky_inverse(torch.linalg.cholesky(hessian + eps)) + + return hessian_inv + + +def prepare( + tensor: torch.Tensor, hessian: torch.Tensor, hessian_damp: float +) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare the inverse Hessian matrix.""" + weight = tensor.detach().clone() + # move the hessian matrix from CPU to GPU for acceleration + hessian = hessian.to(weight.device) + if len(weight.size()) == 4: + weight = weight.flatten(1) + + zero = torch.diag(hessian) == 0 + hessian[zero, zero] = 1 + weight[:, zero] = 0 + + damp = hessian_damp * torch.mean(torch.diag(hessian)) + cols = weight.size(1) + diag = torch.arange(cols) + hessian[diag, diag] += damp + + hessian_inv = invert(hessian) + + # remove the Hessian matrix to save GPU memory + del hessian + torch.cuda.empty_cache() + + return weight, hessian_inv + + +@torch.no_grad() +def create_sgpt_mask( + tensor: torch.Tensor, hessian: torch.Tensor, config: SearchConfig +) -> torch.Tensor: + """Create a sparse mask for the given tensor.""" + shape = tensor.size() + weight, hessian_inv = prepare(tensor, hessian, config["hessian_damp"]) + rows, cols = weight.size() + hessian_inv_diag = torch.diagonal(hessian_inv, dim1=0, dim2=1) + + is_nm_prune, n, m = get_nmprune_info(config["pattern"]) + col_bs = config["col_block_size"] + row_bs = config["row_block_size"] + # if row_bs is not specified, prune the whole weight block + if row_bs == -1: + row_bs = rows + + for r1 in range(0, rows, row_bs): + r2 = min(r1 + row_bs, rows) + # the mask of the weights not to be pruned + w_rows = weight[r1:r2].float() + + # pruning the weight block W[row:row+row_bs, i1:i1+col_bs] + for i1 in range(0, cols, col_bs): + i2 = min(i1 + col_bs, cols) + w_blk = w_rows[:, i1:i2].clone() + q_blk = torch.zeros_like(w_blk) + # the error of the weights to be pruned + delta_blk = torch.zeros_like(w_blk) + hinv_blk = hessian_inv[i1:i2, i1:i2] + hinv_diag_blk = hessian_inv_diag[i1:i2] + + errors_blk = (w_blk**2) / (hinv_diag_blk**2 + 1e-9) + if torch.isnan(errors_blk).any(): + print("nan in errors_blk.") + + mask_blk = torch.zeros_like(w_blk, dtype=torch.bool) + + for j in range(i2 - i1): + # compute the error of the weights to be pruned + w = w_blk[:, j] + d = hinv_diag_blk[j] + if is_nm_prune and j % m == 0: + errors_blk = (w_blk[:, j : j + m] ** 2) / (hinv_diag_blk[j : j + m] ** 2 + 1e-9) + mask_blk.scatter_( + 1, j + torch.topk(errors_blk, n, dim=1, largest=False)[1], True + ) + + q = w.clone() + q[mask_blk[:, j]] = 0 + q_blk[:, j] = q + + # update the remaining weights in the col_bs block to compensate the error caused by pruning W[:, j] + err = (w - q) / d + w_blk[:, j:] -= err.unsqueeze(1).matmul(hinv_blk[j, j:].unsqueeze(0)) + delta_blk[:, j] = err + + # compensate the error caused by pruning W[:, i: i + col_bs] with the weights update in W[:, i + col_bs:] + w_rows[:, i1:i2] = q_blk + w_rows[:, i2:] -= delta_blk.matmul(hessian_inv[i1:i2, i2:]) + if torch.isnan(w_rows[:, i2:]).any(): + print("nan") + + weight[r1:r2] = w_rows + + mask = weight != 0 + + return mask.view(shape) + + +class SparseGPTSearcher(BaseSparseSearcher): + """SparseGPT-based sparse mask searching algorithm.""" + + @property + def default_search_config(self) -> SearchConfig: + """Get the default config for the searcher.""" + return { + **super().default_search_config, + "col_block_size": 128, # column block size in sparsegpt + "row_block_size": -1, # row block size in sparsegpt + "hessian_damp": 0.1, # hessian damp in sparsegpt + "calib_size": 256, # calibration size for hessian matrix calculation + "device": "cuda", # device of hessian matrix + } + + def _check_weight_size(self, weight, mod_name) -> bool: + """Check if the weight size is supported by SparseGPT.""" + _, _, m = get_nmprune_info(self.config["pattern"]) + + # the column size must be divisible by m + if weight.size(0) % m != 0 or weight.size(1) % m != 0: + warnings.warn( + f"Skipping pruning {mod_name} of size={str(weight.size())} and" + f" type={str(weight.dtype)} for SparseGPT" + ) + return False + + return True + + def _compute_mask(self, module: SparseModule) -> torch.BoolTensor: + """Compute the mask (and weight update) for the given module.""" + return create_sgpt_mask(module.weight, module.hessian, self.config) + + @torch.no_grad() + def before_search(self): + """Register the forward hook to collect the hessian matrix.""" + super().before_search() + + handles = [] + for _, module in self._named_sparsifiable_modules(): + # setup and register the forward hook + self._setup_forward_hook(module) + handles.append(module.register_forward_hook(self._hook_compute_hessian)) + + print_rank_0(f"Collecting Hessian statistics for {len(handles)} modules.") + + # run a forward loop to collect the hessian matrix + assert self.forward_loop is not None, "Please provide `data_loader` or `forward_loop`!" + self.forward_loop(self.model) + + # remove the forward hooks + for handle in handles: + handle.remove() + + def after_search(self): + """Remove Hessian artifcats from network.""" + super().after_search() + for _, module in self._named_sparsifiable_modules(): + del module.hessian + del module.samples + + @staticmethod + def _is_memory_sufficient(device_id, threshold): + """Check if the memory usage on the CUDA device is below the threshold.""" + total_memory = torch.cuda.get_device_properties(device_id).total_memory + allocated_memory = torch.cuda.memory_allocated(device_id) + free_memory = total_memory - allocated_memory + return free_memory / total_memory > (1 - threshold) + + @classmethod + def _setup_forward_hook(cls, mod: SparseModule) -> None: + """Setup the attributes we need for our forward hook during the SparseGPT search.""" + # initialize the hessian matrix + if isinstance(mod, nn.Conv2d): + # For conv2d layers, the hessian matrix is calculated as X * X^T, where X is the + # flattened weight matrix. + cols = mod.weight.size(1) * mod.weight.size(2) * mod.weight.size(3) + else: + # For linear layers, the hessian matrix is calculated as X * X^T, where X is the + # weight matrix. + cols = mod.weight.size(1) + + target_device = mod.weight.device + # Hessian matrix is stored in the GPU memory by default + if target_device.type == "cuda" and cls._is_memory_sufficient(target_device.index, 0.8): + hessian = torch.zeros((cols, cols), dtype=torch.float32).to(target_device) + else: + hessian = torch.zeros((cols, cols), dtype=torch.float32).to("cpu") + + # store the hessian matrix and the number of samples + # TODO: this should probably be improved eventually!! + mod.hessian = hessian + mod.samples = 0 + + @classmethod + def _hook_compute_hessian(cls, mod: nn.Module, inp: torch.Tensor, out: torch.Tensor): + with torch.inference_mode(): + # TODO: move the hessian matrix to GPU memory if possible + if isinstance(inp, tuple): + inp = inp[0] + # use torch.float32 to avoid overflow in Hessian + if len(inp.shape) == 2: + inp = inp.unsqueeze(0) + tmp = inp.shape[0] + # nn.Linear and *ParallelLinear in mcore + if "Linear" in type(mod).__name__: + if len(inp.shape) == 3: + inp = inp.reshape((-1, inp.shape[-1])) + inp.t_() + if isinstance(mod, nn.Conv2d): + unfold = nn.Unfold( + mod.kernel_size, + dilation=mod.dilation, + stride=mod.stride, + ) + inp = unfold(inp) + inp = inp.permute([1, 0, 2]) + inp = inp.flatten(1) + mod.hessian *= mod.samples / (mod.samples + tmp) + mod.samples += tmp + inp = math.sqrt(2 / mod.samples) * inp.float() + + # the hessian matrix is calculated as X * X^T + target_device = mod.hessian.device + if mod.hessian.device.type == "cuda": + if cls._is_memory_sufficient(mod.hessian.device.index, 0.8): + mod.hessian += inp.matmul(inp.t()).to(mod.hessian.device) + else: + target_device = "cpu" + mod.hessian = mod.hessian.to("cpu") + mod.hessian += inp.matmul(inp.t()).to(target_device) + + assert not torch.isinf(mod.hessian).any(), "Hessian contains inf" diff --git a/modelopt/torch/sparsity/sparsification.py b/modelopt/torch/sparsity/sparsification.py new file mode 100644 index 0000000000..0773be813f --- /dev/null +++ b/modelopt/torch/sparsity/sparsification.py @@ -0,0 +1,123 @@ +# 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. + +"""High-level API to automatically sparsify your model with various algorithms.""" + +from typing import Any, Optional + +from torch import nn + +from modelopt.torch.opt.conversion import apply_mode, get_mode +from modelopt.torch.opt.mode import ModeLike +from modelopt.torch.opt.searcher import BaseSearcher, SearchConfig +from modelopt.torch.utils import unwrap_model + +from .mode import SparsityModeRegistry + +__all__ = ["sparsify", "export"] + + +def sparsify( + model: nn.Module, mode: ModeLike, config: Optional[SearchConfig] = None +) -> tuple[nn.Module, dict[str, Any]]: + """Sparsify a given model and search for they optimal sparsified weights. + + Args: + model: A standard model that contains standard building blocks to be sparsified in-place. + mode: A (list of) string(s) or Mode(s) or a list of tuples containing the mode and its + config indicating the desired mode(s) (and configurations) for the convert + process. Modes set up the model for different algorithms for model optimization. The + following modes are available: + + * :class:`"sparse_magnitude"`: + The ``model`` will be sparsified according to the magnitude of weights in each + layer. The mode's config is described in + :class:`SparseMagnitudeConfig`. + * :class:`"sparsegpt"`: + The ``model`` will be sparsified and weights are updated optimally using an Hessian + approximation of the loss function (see SparseGPT paper for details). The mode's + config is described in + :class:`SparseGPTConfig`. + + If the mode argument is specified as a dictionary, the keys should indicate the mode and + the values specify the per-mode configuration. If not provided, then default + configuration would be used. + + config: Additional optional arguments to configure the search. Currently, we support: + + * ``verbose``: Whether to print detailed search stats during search. + * ``forward_loop``: A ``Callable`` that takes a model as input and runs a forward loop + on it. It is recommended to choose the data loader used inside the forward loop + carefully to reduce the runtime. Cannot be provided at the same time as + ``data_loader`` and ``collect_func``. + * ``data_loader``: An iterator yielding batches of data for calibrating the + normalization layers in the model or compute gradient scores. It is recommended to use + the same data loader as for training but with significantly fewer iterations. Cannot + be provided at the same time as ``forward_loop``. + * ``collect_func``: A ``Callable`` that takes a batch of data from the data loader as + input and returns the input to ``model.forward()`` as described in + :meth:`run_forward_loop `. Cannot + be provided at the same time as ``forward_loop``. + + .. note:: + + Additional configuration options may be added by individual algorithms. Please + refer to the documentation of the individual algorithms for more information. + + Returns: A sparsified model + + .. note:: + + The given model is sparsified in-place. The returned model is thus a reference to the same + model instance as the input model. + """ + # apply sparsity to the model + model = apply_mode(model, mode, registry=SparsityModeRegistry) + + # retrieve searcher class + searcher_cls: type[BaseSearcher] = getattr(get_mode(model), "search_algorithm") + + # run search+sparsification algorithm + searcher = searcher_cls() + searcher.search(model, {}, tuple(), config) + + # return the sparsified model + return model + + +def export(model: nn.Module) -> nn.Module: + """Export a sparse dynamic model to a regular model. + + This should be done after the model is fine-tuned and the weights are fixed. + + .. warning:: + + After the call to ``export()``, the sparsity mask will no longer be enforced. This means any + future weight updates would destroy the sparsity pattern. If you want to continue training, + call ``export()`` after training is finished. + """ + # unwrap a DP/DDP model + model = unwrap_model( + model, + warn=True, + msg=( + f"Unwrapping a {type(model).__name__} model for export! Note that the export is" + " in-place and the model wrapper should be re-created after export since the wrapper" + " might not support changing parameters after initialization." + ), + ) + + # apply export mode and return model + return apply_mode(model, "export_sparse", registry=SparsityModeRegistry) diff --git a/modelopt/torch/speculative/__init__.py b/modelopt/torch/speculative/__init__.py new file mode 100644 index 0000000000..4e1334edbf --- /dev/null +++ b/modelopt/torch/speculative/__init__.py @@ -0,0 +1,20 @@ +# 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. + +"""Speculative Decoding Optimizations.""" + +from . import eagle, medusa, mode, plugins +from .config import * +from .speculative_decoding import * diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py new file mode 100644 index 0000000000..d1819be333 --- /dev/null +++ b/modelopt/torch/speculative/config.py @@ -0,0 +1,63 @@ +# 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. + +"""Configurations for speculative decoding modes.""" + +from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField + + +class MedusaConfig(ModeloptBaseConfig): + """Medusa config.""" + + medusa_num_heads: int = ModeloptField( + default=2, + description=("The number of medusa heads added to the model."), + ) + + medusa_num_layers: int = ModeloptField( + default=1, + description=("The number of ResBlocks used in medusa head."), + ) + + +class EagleConfig(ModeloptBaseConfig): + """Eagle config.""" + + eagle_num_layers: int = ModeloptField( + default=1, + description=("The number of decoder used in the eagle model."), + ) + + use_input_layernorm_in_first_layer: bool = ModeloptField( + default=True, description=("Whether to use input_layernorm in the first decoder layer.") + ) + + use_last_layernorm: bool = ModeloptField( + default=False, description=("Whether to use a final layernorm before lm_head.") + ) + + +class RedrafterConfig(ModeloptBaseConfig): + """Redrafter config.""" + + redrafter_predict_n_tokens: int = ModeloptField( + default=2, + description=("The number of tokens that redrafter will predict"), + ) + + redrafter_num_layers: int = ModeloptField( + default=1, + description=("The number of ResBlocks used in lm head."), + ) diff --git a/modelopt/torch/speculative/eagle/__init__.py b/modelopt/torch/speculative/eagle/__init__.py new file mode 100644 index 0000000000..91df775802 --- /dev/null +++ b/modelopt/torch/speculative/eagle/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""Eagle Optimization Method.""" + +from .conversion import * +from .eagle_model import * diff --git a/modelopt/torch/speculative/eagle/conversion.py b/modelopt/torch/speculative/eagle/conversion.py new file mode 100644 index 0000000000..db4bedc146 --- /dev/null +++ b/modelopt/torch/speculative/eagle/conversion.py @@ -0,0 +1,59 @@ +# 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. + +"""Eagle conversion/restore utilities.""" + +from torch import nn + +from modelopt.torch.opt.conversion import ModelLikeModule +from modelopt.torch.opt.dynamic import _DMRegistryCls +from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict + +from ..config import EagleConfig + +EagleDMRegistry = _DMRegistryCls(prefix="Eagle") # global instance for the registry + + +def convert_to_eagle_model(model: nn.Module, config: EagleConfig) -> ConvertReturnType: + """Convert the model to a eagle model as per `config`.""" + # initialize the true module if necessary + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + + original_cls = type(model) + if original_cls not in EagleDMRegistry: + for cls in EagleDMRegistry._registry: + if issubclass(original_cls, cls): + EagleDMRegistry.register({original_cls: "base_model_class"})(EagleDMRegistry[cls]) + break + + eagle_model = EagleDMRegistry.convert(model) + eagle_model.modify( + eagle_num_layers=config.eagle_num_layers, + use_input_layernorm_in_first_layer=config.use_input_layernorm_in_first_layer, + use_last_layernorm=config.use_last_layernorm, + ) + + # no metadata, all specifed via config. + metadata = {} + + return eagle_model, metadata + + +def restore_eagle_model(model: nn.Module, config: EagleConfig, metadata: MetadataDict) -> nn.Module: + """Function for restoring a previously convert model to a eagle model.""" + # the metadata should be empty + assert not metadata, "No metadata expected!" + + return convert_to_eagle_model(model, config)[0] diff --git a/modelopt/torch/speculative/eagle/eagle_model.py b/modelopt/torch/speculative/eagle/eagle_model.py new file mode 100644 index 0000000000..902408afc8 --- /dev/null +++ b/modelopt/torch/speculative/eagle/eagle_model.py @@ -0,0 +1,32 @@ +# 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. + +"""Eagle model to support eagle decoding.""" + +from modelopt.torch.opt.dynamic import DynamicModule + + +class EagleModel(DynamicModule): + """Base Eagle Model.""" + + def _setup(self): + self._register_temp_attribute("eagle_num_layers", 0) + self._register_temp_attribute("eagle_module", None) + + def modify(self, eagle_num_layers, use_input_layernorm_in_first_layer, use_last_layernorm): + """Base Eagle Model modify function. Child class should implement the details.""" + self.eagle_num_layers = eagle_num_layers + self.use_input_layernorm_in_first_layer = use_input_layernorm_in_first_layer + self.use_last_layernorm = use_last_layernorm diff --git a/modelopt/torch/speculative/eagle/utils.py b/modelopt/torch/speculative/eagle/utils.py new file mode 100644 index 0000000000..5caafa4d23 --- /dev/null +++ b/modelopt/torch/speculative/eagle/utils.py @@ -0,0 +1,388 @@ +# Adapted from https://github.com/SafeAILab/EAGLE/blob/7d065d0/eagle/model/cnets.py +# coding=utf-8 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. + +# 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. + +"""PyTorch LLaMA model.""" + +import math +from typing import Optional + +import torch +from torch import nn + + +# Copied from transformers.models.bart.modeling_bart._make_causal_mask +def make_causal_mask( + input_ids_shape: torch.Size, + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, +): + """Make causal mask used for bi-directional self-attention.""" + bsz, tgt_len = input_ids_shape + mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device) + mask_cond = torch.arange(mask.size(-1), device=device) + mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0) + mask = mask.to(dtype) + + if past_key_values_length > 0: + mask = torch.cat( + [torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1 + ) + return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) + + +# Copied from transformers.models.bart.modeling_bart._expand_mask +def expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): + """Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.""" + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). + + The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids): + """Apply rotary position embedding.""" + # The first two dimensions of cos and sin are always 1, so we can `squeeze` them. + cos = cos.squeeze(1).squeeze(0) # [seq_len, dim] + sin = sin.squeeze(1).squeeze(0) # [seq_len, dim] + cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] + sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class LlamaRotaryEmbedding(nn.Module): + """Llama Rotary Embedding.""" + + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + """Init function for LlamaRotaryEmbedding.""" + super().__init__() + + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) + self.register_buffer("inv_freq", inv_freq) + + # Build here to make `torch.jit.trace` work. + self._set_cos_sin_cache( + seq_len=max_position_embeddings, + device=self.inv_freq.device, + dtype=torch.get_default_dtype(), + ) + + def _set_cos_sin_cache(self, seq_len, device, dtype): + self.max_seq_len_cached = seq_len + t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) + + freqs = torch.einsum("i,j->ij", t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype)) + self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype)) + + def forward(self, x, seq_len=None): + """Forward function for LlamaRotaryEmbedding.""" + # x: [bs, num_attention_heads, seq_len, head_size] + if seq_len > self.max_seq_len_cached: + self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) + + return ( + self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype), + self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype), + ) + + +class LlamaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper.""" + + def __init__( + self, + hidden_size, + num_attention_heads, + num_key_value_heads, + max_position_embeddings, + rope_theta, + ): + """Init function for LlamaAttention.""" + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.num_key_value_heads = num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = max_position_embeddings + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" + f" and `num_heads`: {self.num_heads})." + ) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear( + self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False + ) + self.v_proj = nn.Linear( + self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False + ) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + self.rotary_emb = LlamaRotaryEmbedding( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=rope_theta, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, + output_attentions: bool = False, + use_cache: bool = False, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]: + """Forward function for LlamaAttention.""" + bsz, q_len, _ = hidden_states.size() + + 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.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose( + 1, 2 + ) + value_states = value_states.view( + bsz, q_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value[0].shape[-2] + cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) + query_states, key_states = apply_rotary_pos_emb( + query_states, key_states, cos, sin, position_ids + ) + + if past_key_value is not None: + # reuse k, v, self_attention + key_states = torch.cat([past_key_value[0], key_states], dim=2) + value_states = torch.cat([past_key_value[1], value_states], dim=2) # type: ignore + + past_key_value = (key_states, value_states) if use_cache else None + + # repeat k/v heads if n_kv_heads < n_heads + 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 attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( + query_states.dtype + ) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class LlamaRMSNorm(nn.Module): + """LlamaRMSNorm class.""" + + def __init__(self, hidden_size, eps=1e-6): + """LlamaRMSNorm is equivalent to T5LayerNorm.""" + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + """Forward function for LlamaRMSNorm.""" + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +class LlamaMLP(nn.Module): + """LlamaMLP class.""" + + def __init__(self, hidden_size, intermediate_size): + """Init function for LlamaMLP.""" + super().__init__() + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = torch.nn.SiLU() + + def forward(self, x): + """Forward function for LlamaMLP.""" + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + return down_proj + + +class LlamaDecoderLayer(nn.Module): + """LlamaDecoderLayer class.""" + + def __init__( + self, + index, + hidden_size, + intermediate_size=14336, + rms_norm_eps=1e-05, + num_attention_heads=32, + num_key_value_heads=8, + max_position_embeddings=131072, + rope_theta=500000.0, + use_input_layernorm_in_first_layer=True, + ): + """Init function for LlamaDecoderLayer.""" + super().__init__() + self.hidden_size = hidden_size + self.self_attn = LlamaAttention( + hidden_size, + num_attention_heads, + num_key_value_heads, + max_position_embeddings, + rope_theta, + ) + self.mlp = LlamaMLP(hidden_size, intermediate_size) + self.index = index + if use_input_layernorm_in_first_layer or self.index != 0: + self.input_layernorm = LlamaRMSNorm(hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = LlamaRMSNorm(hidden_size, eps=rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]: + """Forward function for LlamaDecoderLayer.""" + residual = hidden_states + + if hasattr(self, "input_layernorm"): + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs # type: ignore diff --git a/modelopt/torch/speculative/medusa/__init__.py b/modelopt/torch/speculative/medusa/__init__.py new file mode 100644 index 0000000000..ae4745d787 --- /dev/null +++ b/modelopt/torch/speculative/medusa/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""Medusa Optimization Method.""" + +from .conversion import * +from .medusa_model import * diff --git a/modelopt/torch/speculative/medusa/conversion.py b/modelopt/torch/speculative/medusa/conversion.py new file mode 100644 index 0000000000..df3e59de51 --- /dev/null +++ b/modelopt/torch/speculative/medusa/conversion.py @@ -0,0 +1,59 @@ +# 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. + +"""Medusa conversion/restore utilities.""" + +from torch import nn + +from modelopt.torch.opt.conversion import ModelLikeModule +from modelopt.torch.opt.dynamic import _DMRegistryCls +from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict + +from ..config import MedusaConfig + +MedusaDMRegistry = _DMRegistryCls(prefix="Medusa") # global instance for the registry + + +def convert_to_medusa_model(model: nn.Module, config: MedusaConfig) -> ConvertReturnType: + """Convert the model to a medusa model as per `config`.""" + # initialize the true module if necessary + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + + original_cls = type(model) + if original_cls not in MedusaDMRegistry: + for cls in MedusaDMRegistry._registry: + if issubclass(original_cls, cls): + MedusaDMRegistry.register({original_cls: "base_model_class"})(MedusaDMRegistry[cls]) + break + + medusa_model = MedusaDMRegistry.convert(model) + medusa_model.modify( + medusa_num_heads=config.medusa_num_heads, medusa_num_layers=config.medusa_num_layers + ) + + # no metadata, all specifed via config. + metadata = {} + + return medusa_model, metadata + + +def restore_medusa_model( + model: nn.Module, config: MedusaConfig, metadata: MetadataDict +) -> nn.Module: + """Function for restoring a previously convert model to a medusa model.""" + # the metadata should be empty + assert not metadata, "No metadata expected!" + + return convert_to_medusa_model(model, config)[0] diff --git a/modelopt/torch/speculative/medusa/medusa_model.py b/modelopt/torch/speculative/medusa/medusa_model.py new file mode 100644 index 0000000000..37b1bb422d --- /dev/null +++ b/modelopt/torch/speculative/medusa/medusa_model.py @@ -0,0 +1,32 @@ +# 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. + +"""Medusa model to support medusa decoding.""" + +from modelopt.torch.opt.dynamic import DynamicModule + + +class MedusaModel(DynamicModule): + """Base Medusa Model.""" + + def _setup(self): + self._register_temp_attribute("medusa_num_heads", 0) + self._register_temp_attribute("medusa_num_layers", 0) + self._register_temp_attribute("medusa_heads", None) + + def modify(self, medusa_num_heads=0, medusa_num_layers=0): + """Base Medusa Model modify function. Child class should implement the details.""" + self.medusa_num_heads = medusa_num_heads + self.medusa_num_layers = medusa_num_layers diff --git a/modelopt/torch/speculative/mode.py b/modelopt/torch/speculative/mode.py new file mode 100644 index 0000000000..d09f22f3f7 --- /dev/null +++ b/modelopt/torch/speculative/mode.py @@ -0,0 +1,115 @@ +# 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. + +"""This module contains the mode descriptor for the quantization mode.""" + +from modelopt.torch.opt.config import ModeloptBaseConfig +from modelopt.torch.opt.mode import ( + ConvertEntrypoint, + RestoreEntrypoint, + _ModeDescriptor, + _ModeRegistryCls, +) + +from .config import EagleConfig, MedusaConfig, RedrafterConfig +from .eagle.conversion import convert_to_eagle_model, restore_eagle_model +from .medusa.conversion import convert_to_medusa_model, restore_medusa_model +from .redrafter.conversion import convert_to_redrafter_model, restore_redrafter_model + +SpeculativeDecodingModeRegistry = _ModeRegistryCls() + + +@SpeculativeDecodingModeRegistry.register_mode +class MedusaModeDescriptor(_ModeDescriptor): + """Class to describe the ``"medusa"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "medusa" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return MedusaConfig + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return convert_to_medusa_model + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_medusa_model + + +@SpeculativeDecodingModeRegistry.register_mode +class EagleModeDescriptor(_ModeDescriptor): + """Class to describe the ``"eagle"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "eagle" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return EagleConfig + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return convert_to_eagle_model + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_eagle_model + + +@SpeculativeDecodingModeRegistry.register_mode +class RedrafterModeDescriptor(_ModeDescriptor): + """Class to describe the ``"redrafter"`` mode. + + The properties of this mode can be inspected via the source code. + """ + + @property + def name(self) -> str: + """Returns the value (str representation) of the mode.""" + return "redrafter" + + @property + def config_class(self) -> type[ModeloptBaseConfig]: + """Specifies the config class for the mode.""" + return RedrafterConfig + + @property + def convert(self) -> ConvertEntrypoint: + """The mode's entrypoint for converting a model.""" + return convert_to_redrafter_model + + @property + def restore(self) -> RestoreEntrypoint: + """The mode's entrypoint for restoring a model.""" + return restore_redrafter_model diff --git a/modelopt/torch/speculative/plugins/__init__.py b/modelopt/torch/speculative/plugins/__init__.py new file mode 100644 index 0000000000..2453382393 --- /dev/null +++ b/modelopt/torch/speculative/plugins/__init__.py @@ -0,0 +1,30 @@ +# 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. + +"""Handles speculative plugins for third-party modules. + +Please check out the source code of this module for examples of how plugins work and how you can +write your own one. Currently, we support plugins for + +- :meth:`transformers` +""" + +from modelopt.torch.utils import import_plugin + +with import_plugin("megatron"): + from .megatron import * + +with import_plugin("transformers"): + from .transformers import * diff --git a/modelopt/torch/speculative/plugins/megatron.py b/modelopt/torch/speculative/plugins/megatron.py new file mode 100644 index 0000000000..aeca0052a0 --- /dev/null +++ b/modelopt/torch/speculative/plugins/megatron.py @@ -0,0 +1,722 @@ +# 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. + +"""Plugin to add Medusa support for megatron-core GPT model.""" + +import copy +import warnings +from typing import Optional + +import megatron.core +import torch +import torch.nn.functional as F +from megatron.core import InferenceParams, tensor_parallel +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding +from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding +from megatron.core.models.gpt import GPTModel +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.parallel_state import get_data_parallel_rank, get_tensor_model_parallel_rank +from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.transformer_block import TransformerBlock +from megatron.core.transformer.utils import sharded_state_dict_default +from megatron.core.utils import make_tp_sharded_tensor_for_checkpoint +from packaging.version import Version + +from ..eagle.conversion import EagleDMRegistry +from ..eagle.eagle_model import EagleModel +from ..medusa.conversion import MedusaDMRegistry +from ..medusa.medusa_model import MedusaModel + + +def mcore_version_higher_than(target_version: str): + """Check if megatron-core is least this version.""" + return Version(megatron.core.__version__) > Version(target_version) + + +class MedusaLayer(MegatronModule): + """MedusaLayer impl following TensorRT-LLM's model definition. + + Medusa layer consists of a column parallel linear following a silu. + """ + + def __init__(self, config): + """Constructor. + + Args: + config: MCore transformer config + """ + super().__init__(config=config) + + self.activation_func = F.silu + + self.linear = tensor_parallel.ColumnParallelLinear( + config.hidden_size, + config.hidden_size, + config=config, + init_method=config.init_method, + bias=True, + skip_bias_add=False, + gather_output=True, + skip_weight_param_allocation=False, + ) + + def forward(self, x): + """Forward function.""" + y, _ = self.linear(x) + return x + self.activation_func(y), None + + def sharded_state_dict( + self, prefix: str = "", sharded_offsets: tuple = (), metadata: Optional[dict] = None + ) -> ShardedStateDict: + """Return MCore sharded_state_dict.""" + return self.linear.sharded_state_dict(f"{prefix}linear.", sharded_offsets, metadata) + + +class MedusaHead(MegatronModule): + """MedusaHead impl following TensorRT-LLM's model definition. + + https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/models/medusa/model.py + Medusa head consists of several MedusaLayers and an lm_head. + """ + + def __init__(self, config, vocab_size: int, num_layers: int = 1, parallel_output: bool = True): + """Constructor. + + Args: + config: MCore transformer config + vocab_size: vocabulary size + num_layers: number of Medusa layers + parallel_output: if False, then all_gather the logits + """ + super().__init__(config=config) + + self.medusa_layers = torch.nn.ModuleList([MedusaLayer(config) for _ in range(num_layers)]) + + self.lm_head = tensor_parallel.ColumnParallelLinear( + config.hidden_size, + vocab_size, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + gather_output=not parallel_output, + skip_weight_param_allocation=False, + ) + + def load_state_dict_post_hook(module, incompatible_keys): + incompatible_keys.missing_keys.clear() + incompatible_keys.unexpected_keys.clear() + + self.register_load_state_dict_post_hook(load_state_dict_post_hook) + + def forward(self, x): + """Forward function.""" + for layer in self.medusa_layers: + x, _ = layer(x) + return self.lm_head(x) + + def sharded_state_dict( + self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict = None + ) -> ShardedStateDict: + """Return MCore sharded_state_dict.""" + assert not sharded_offsets, "Unexpected sharded offsets" + sharded_state_dict = {} + layer_prefix = f"{prefix}medusa_layers." + for i, layer in enumerate(self.medusa_layers): + state_dict_prefix = f"{layer_prefix}{i}." + sharded_pp_offset = [] + layer_sharded_state_dict = layer.sharded_state_dict( + state_dict_prefix, sharded_pp_offset, metadata + ) + sharded_state_dict.update(layer_sharded_state_dict) + sharded_state_dict.update( + self.lm_head.sharded_state_dict(f"{prefix}lm_head.", sharded_offsets, metadata) + ) + return sharded_state_dict + + +@MedusaDMRegistry.register({GPTModel: "megatron.core.models.gpt.GPTModel"}) +class _DynamicMedusaGPTModel(MedusaModel): + """A ``megatron.core.models.gpt.GPTModel`` model with dynamic hyperparams.""" + + def _setup(self): + super()._setup() + self._register_temp_attribute("medusa_report_acc", True) + self._register_temp_attribute("medusa_freeze_base_model", True) + self._register_temp_attribute("calibration_mode", False) + + def modify( + self, + medusa_num_heads=0, + medusa_num_layers=0, + medusa_freeze_base_model=True, + medusa_report_acc=True, + ): + """Constructor. + + Args: + config: MedusaConfig that specifies the medusa head configuration as well as + weights of base model and medusa head. + """ + if self.config.pipeline_model_parallel_size > 1: + warnings.warn( + "Pipeline parallelism detected! _DynamicMedusaGPTModel only supports " + "pipeline parallelism during TensorRT-LLM checkpoint export." + ) + super().modify(medusa_num_heads=medusa_num_heads, medusa_num_layers=medusa_num_layers) + + self.medusa_report_acc = medusa_report_acc + self.medusa_freeze_base_model = medusa_freeze_base_model + + # Freeze all parameters + if self.medusa_freeze_base_model: + for name, param in self.named_parameters(): + param.requires_grad = False + + if self.post_process: + self.medusa_heads = torch.nn.ModuleList( + [ + MedusaHead(self.config, self.vocab_size, num_layers=self.medusa_num_layers) + for _ in range(self.medusa_num_heads) + ] + ) + + def forward(self, *args, labels: torch.Tensor = None, **kwargs): + """Forward pass of the Medusa GPTModel. + + Returns: + torch.Tensor: If labels are provided, then return lm_loss of all heads. Otherwise, + return the original logits. + """ + if self.post_process: + # Set the post_process to False such that the forward will return the hidden_state. + self.post_process = False + # Calling parent's forward to get hidden_states + hidden_states = GPTModel.forward(self, *args, labels=labels, **kwargs) + # Reset the post_process to True + self.post_process = True + else: + # return GPTModel.forward(self, *args, labels=labels, **kwargs) + return GPTModel.forward(self, *args, labels=None, **kwargs) + + # Original output logits + logits, _ = self.output_layer(hidden_states) + + report_acc = self.medusa_report_acc and labels is not None + + acc = [] + + # Medusa heads forward. We want to run through all the heads just to make sure all modules + # are exercised during calibration. + for i, head in enumerate(self.medusa_heads): + new_logits, _ = head(hidden_states) + + # If label is not provided, then this is the inference/generation case. We didn't + # implement fast decoding; hence we want to return the original logits untouched. + if labels is not None: + logits = torch.cat((logits, new_logits), dim=0) + + if report_acc: + seq_len = new_logits.shape[0] + gathered_logits = gather_from_tensor_model_parallel_region(new_logits) + medusa_top1 = gathered_logits.transpose(0, 1).argmax(dim=-1) + medusa_labels = labels[:, (i + 1) * seq_len : (i + 2) * seq_len] + top1_p = torch.eq(medusa_labels, medusa_top1).sum() / medusa_top1.numel() + acc.append(top1_p) + + # Return the original logits untouched. + if labels is None: + # [s b h] => [b s h] + return logits.transpose(0, 1).contiguous() + + if report_acc and get_tensor_model_parallel_rank() == 0: + print("Medusa Training Accuracy: {}".format(acc)) + + loss = self.compute_language_model_loss(labels, logits) + + return loss + + def sharded_state_dict( + self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict = None + ) -> ShardedStateDict: + """Override the shared_state_dict to take care medusa_heads.""" + assert not sharded_offsets, "Unexpected sharded offsets" + + sharded_state_dict = GPTModel.sharded_state_dict(self, prefix, sharded_offsets, metadata) + + if not hasattr(self, "medusa_heads") or self.medusa_heads is None: + return sharded_state_dict + + # This is a remedy for nn.ModuleList. GPTModel.sharded_state_dict() is calling into + # MegatronModule.sharded_state_dict() which requires all children to implement + # sharded_state_dict(). medusa_heads is an nn.ModuleList which only has state_dict() + # implemented. As a result, all the submodules will not be sharded. + # + # The remedy is to pop all medusa_heads* out and call the MedusaHead sharded_state_dict() + # again to populate the correct sharded_staet_dict. + extra_keys = [] + for key, val in sharded_state_dict.items(): + if "medusa_heads" in key: + extra_keys += [key] + for key in extra_keys: + sharded_state_dict.pop(key, None) + + layer_prefix = f"{prefix}medusa_heads." + for i, layer in enumerate(self.medusa_heads): + layer_sharded_state_dict = layer.sharded_state_dict(f"{layer_prefix}{i}.", [], metadata) + sharded_state_dict.update(layer_sharded_state_dict) + return sharded_state_dict + + +class EagleLanguageModelEmbedding(LanguageModelEmbedding): + """Allow last pp stage to also load the embedding.""" + + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: tuple[tuple[int, int, int]] = (), + metadata: Optional[dict] = None, + ) -> ShardedStateDict: + """Different from the default, we change the state_dict to have 1 replica at pp.""" + state_dict = self.state_dict(prefix="", keep_vars=True) + + weight_prefix = f"{prefix}word_embeddings.weight" + return { + weight_prefix: make_tp_sharded_tensor_for_checkpoint( + tensor=state_dict["word_embeddings.weight"], + key=weight_prefix, + allow_shape_mismatch=True, + prepend_offsets=sharded_offsets, + replica_id=(0, 1, get_data_parallel_rank(with_context_parallel=True)), + ) + } + + +class EagleTransformerBlock(TransformerBlock): + """Only store the EAGLE decoder in the last pp stage.""" + + def sharded_state_dict( + self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict = None + ) -> ShardedStateDict: + """Generate a sharded state dictionary for the transformer block. + + Args: + prefix (str, optional): Prefix to be added to all keys in the state dict. + Defaults to an empty string. + sharded_offsets (tuple, optional): Tuple of sharding offsets. + metadata (dict, optional): Additional metadata for sharding. + Can specify if layers are non-homogeneous. Defaults to None. + + Returns: + ShardedStateDict: A dictionary containing the sharded state of the model. + """ + assert not sharded_offsets, "Unexpected sharded offsets" + + sharded_state_dict = {} + + layer_prefix = f"{prefix}layers." + num_layers = self.config.num_layers + for global_layer_offset, layer in enumerate(self.layers): + # global_layer_offset = layer.layer_number - 1 # self.layer_number starts at 1 + state_dict_prefix = ( + f"{layer_prefix}{global_layer_offset}." # module list index in TransformerBlock + ) + sharded_prefix = layer_prefix + sharded_pp_offset = [ + (0, global_layer_offset, num_layers) + ] # PP sharding offset for ShardedTensors + + layer_sharded_state_dict = layer.sharded_state_dict( + state_dict_prefix, sharded_pp_offset, metadata + ) + replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix) + sharded_state_dict.update(layer_sharded_state_dict) + + # Add modules other than self.layers + for name, module in self.named_children(): + if module is not self.layers: + sharded_state_dict.update( + sharded_state_dict_default( + module, f"{prefix}{name}.", sharded_offsets, metadata + ) + ) + + return sharded_state_dict + + +class EagleModule(MegatronModule): + """EagleModule definition. + + EagleModule consists of an FC projection and additional decoder layers. + """ + + def __init__(self, config, num_layers: int, transformer_layer_spec, use_last_layernorm): + """Constructor. + + EagleModule is essentially a GPTModel except that it only exists in + the last pp stage. As a result, pre_process must be True (otherwise + the decoder expects the input is from the receive buffer). + post_process must be True to perform the final_layernorm. + + Args: + config: MCore transformer config + num_layers: number of Eagle layers + transformer_layer_spec: Megatron core mode spec + """ + super().__init__(config=config) + + eagle_config = copy.deepcopy(config) + eagle_config.num_layers = num_layers + eagle_config.pipeline_model_parallel_size = 1 + self.fc = tensor_parallel.ColumnParallelLinear( + config.hidden_size * 2, + config.hidden_size, + config=eagle_config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + gather_output=True, + skip_weight_param_allocation=False, + ) + # Eagle does not use the final_layernorm in decoder. It distills the base model + # final_layernorm to eagle_module hidden_states directly. + self.decoder = EagleTransformerBlock( + config=eagle_config, + spec=transformer_layer_spec, + post_layer_norm=use_last_layernorm, + pre_process=True, + post_process=True, + ) + + def forward( + self, + decoder_input: torch.Tensor, + attention_mask: torch.Tensor, + rotary_pos_emb: torch.Tensor, + inference_params: InferenceParams = None, + packed_seq_params: PackedSeqParams = None, + extra_block_kwargs: dict = None, + ) -> torch.Tensor: + """Forward function.""" + decoder_input, _ = self.fc(decoder_input) + + decoder_input_list = [decoder_input] + self.decoder.set_input_tensor(decoder_input_list[0]) + hidden_states = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + inference_params=inference_params, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + **(extra_block_kwargs or {}), + ) + return hidden_states + + +@EagleDMRegistry.register({GPTModel: "megatron.core.models.gpt.GPTModel"}) +class _DynamicEagleGPTModel(EagleModel): + """A ``megatron.core.models.gpt.GPTModel`` model with dynamic hyperparams.""" + + def _setup(self): + super()._setup() + self._register_temp_attribute("eagle_self_logit_distillation", False) + self._register_temp_attribute("eagle_freeze_base_model", True) + self._register_temp_attribute("calibration_mode", False) + + def modify( + self, + eagle_num_layers=0, + use_input_layernorm_in_first_layer=True, + use_last_layernorm=False, + eagle_self_logit_distillation=False, + eagle_freeze_base_model=True, + eagle_report_acc=True, + ): + if self.config.pipeline_model_parallel_size > 1: + warnings.warn( + "Pipeline parallelism detected! _DynamicEagleGPTModel only supports " + "pipeline parallelism during TensorRT-LLM checkpoint export." + ) + super().modify( + eagle_num_layers=eagle_num_layers, + use_input_layernorm_in_first_layer=use_input_layernorm_in_first_layer, + use_last_layernorm=use_last_layernorm, + ) + self.eagle_report_acc = eagle_report_acc + self.eagle_self_logit_distillation = eagle_self_logit_distillation + self.eagle_freeze_base_model = eagle_freeze_base_model + + if self.share_embeddings_and_output_weights: + raise ValueError("For EAGLE, args.share_embeddings_and_output_weights must be False") + + if self.position_embedding_type != "rope": + raise ValueError("For EAGLE, only rotary embedding is supported") + + if not self.pre_process and self.post_process: + self.embedding = EagleLanguageModelEmbedding( + config=self.config, + vocab_size=self.vocab_size, + max_sequence_length=self.max_sequence_length, + position_embedding_type=self.position_embedding_type, + ) + + # Freeze all parameters + if self.eagle_freeze_base_model: + for name, param in self.named_parameters(): + param.requires_grad = False + + # Only the last PP stage has the additional projection and decoder layer. + # This is to simplify the export. + if self.post_process: + self.eagle_module = EagleModule( + self.config, + self.eagle_num_layers, + self.transformer_layer_spec, + self.use_last_layernorm, + ) + + # Eagle loss functions + self.eagle_loss = torch.nn.SmoothL1Loss(reduction="none") + # self.kld = LogitsKLLoss() + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + decoder_input: torch.Tensor = None, + labels: torch.Tensor = None, + inference_params: InferenceParams = None, + packed_seq_params: PackedSeqParams = None, + extra_block_kwargs: dict = None, + **kwargs, + ) -> torch.Tensor: + # Word and rotary positional embeddings + if decoder_input is not None: + pass + elif self.pre_process: + decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids) + else: + # intermediate stage of pipeline + # decoder will get hidden_states from decoder.input_tensor + decoder_input = None + + if mcore_version_higher_than("0.9.0"): + extra_kwargs = {"packed_seq_params": None} + else: + extra_kwargs = {} + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + inference_params, + self.decoder, + decoder_input, + self.config, + **extra_kwargs, + ) + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len) + + # Run base modeld decoder forward. + hidden_states = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + inference_params=inference_params, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + **(extra_block_kwargs or {}), + ) + + # Typically, this is only the case when PP > 1. + if not self.post_process: + return hidden_states + + logits_sbh, _ = self.output_layer(hidden_states, weight=None) + + if inference_params is None or self.calibration_mode: + # EAGLE needs to know the newly generated token logits[-1, :, :] + logits_bsh = logits_sbh[:-1, :, :].transpose(0, 1) + gathered_logits_bsh = gather_from_tensor_model_parallel_region(logits_bsh) + prob_bsh = torch.softmax(gathered_logits_bsh, dim=-1) + eager_ids = prob_bsh.argmax(dim=-1) + + eagle_embeddings = self.embedding( + input_ids=eager_ids, + position_ids=position_ids[:, 1:], + ) + # cat([s - 1, b, h], [s - 1, b, h], dim-1) = [s - 1, b, 2h] + eagle_decoder_input = torch.cat((eagle_embeddings, hidden_states[:-1, :, :]), dim=-1) + eagle_hidden_states = self.eagle_module( + eagle_decoder_input, + attention_mask[:, :, 1:, 1:], # TODO (chenhany): this may needs some fix + rotary_pos_emb=rotary_pos_emb[1:, :, :, :], + inference_params=inference_params, + packed_seq_params=packed_seq_params, + **(extra_block_kwargs or {}), + ) + eagle_logits, _ = self.output_layer(eagle_hidden_states) + + # If labels are not provided, return the original logits. We only return after + # all eagle weights have been exercised for quantization calibration purpose. + if labels is None: + return logits_sbh.transpose(0, 1).contiguous() + + # Compute hidden state regression loss + regression_loss = ( + self.eagle_loss( + eagle_hidden_states, + hidden_states[1:, :, :], + ) + .mean(dim=-1) + .transpose(0, 1) + ) + + # Compute lm loss (classification loss) or KLDivergence + if self.eagle_self_logit_distillation: + lm_loss = self.kld(eagle_logits, logits_sbh[1:, :, :]) + else: + lm_loss = self.compute_language_model_loss(labels[:, 1:], eagle_logits) + + acc = [] + if self.eagle_report_acc: + with torch.no_grad(): + gathered_logits = gather_from_tensor_model_parallel_region(eagle_logits) + eagle_top1 = gathered_logits.transpose(0, 1).argmax(dim=-1) + top1_p = torch.eq(labels[:, 1:], eagle_top1).sum() / eagle_top1.numel() + acc.append(top1_p) + + if get_tensor_model_parallel_rank() == 0: + print("EAGLE Training Accuracy: {}".format(acc)) + + return lm_loss + 10.0 * regression_loss + + def tree_decode(self, input_ids: torch.Tensor, tree=None): + self.eval() + + choices = [] if tree is None else tree.choices + draft_len = len(choices) + num_additional_tokens = 0 if tree is None else tree.depth + + micro_batch_size, seq_len = input_ids.size() + att_mask_len = seq_len + 1 + draft_len + + # Prepare attention masks + attention_mask = torch.tril( + torch.ones( + (1, att_mask_len, att_mask_len), + device=input_ids.device, + ), + ).view(1, 1, att_mask_len, att_mask_len) + attention_mask = attention_mask < 0.5 + + if draft_len > 0: + draft_attention_mask = torch.tensor( + tree.attn_mask, + device=input_ids.device, + ).view(1, 1, draft_len, draft_len) + attention_mask[0, 0, -draft_len:, -draft_len:] = draft_attention_mask + + # Prepare position ids and rope + position_ids = torch.arange( + att_mask_len, + dtype=torch.long, + device=input_ids.device, + ) + rotary_pos_emb = self.rotary_pos_emb(seq_len + 1 + num_additional_tokens) + + if num_additional_tokens > 0: + rotary_pos_emb, additional_rotary_pos_emb = torch.split( + rotary_pos_emb, + (seq_len + 1, num_additional_tokens), + dim=0, + ) + for i, c in enumerate(choices): + position_ids[seq_len + 1 + i] = seq_len + len(c) + rotary_pos_emb = torch.cat( + (rotary_pos_emb, additional_rotary_pos_emb[[len(c) - 1], :, :, :]), dim=0 + ) + + position_ids = position_ids.unsqueeze(0) + + # Forward + decoder_input = self.embedding( + input_ids=input_ids, + position_ids=position_ids[:, :seq_len], + ) + hidden_states = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask[:, :, :seq_len, :seq_len], + inference_params=None, + rotary_pos_emb=rotary_pos_emb[:seq_len, :, :, :], + packed_seq_params=None, + ) + logits, _ = self.output_layer(hidden_states) + + # [s b h] => [b s h] + all_logprob = gather_from_tensor_model_parallel_region(logits[[-1], :, :].transpose(0, 1)) + all_logprob = torch.softmax(all_logprob, dim=-1) + top_vals, top_ids = all_logprob[:, -1, :].topk(1, dim=-1) + + if num_additional_tokens == 0: + return top_ids + + new_tokens = top_ids + eagle_ids = torch.cat((input_ids[:, 1:], top_ids), dim=-1) + eagle_rotary_pos_emb = rotary_pos_emb[1 : 1 + eagle_ids.shape[-1], :, :, :] + eagle_hidden_states = hidden_states + + for i in range(num_additional_tokens): + eagle_position_ids = position_ids[:, 1 : 1 + eagle_ids.shape[-1]] + eagle_attn_mask = attention_mask[ + :, :, 1 : 1 + eagle_ids.shape[-1], 1 : 1 + eagle_ids.shape[-1] + ] + eagle_rotary_pos_emb = rotary_pos_emb[1 : 1 + eagle_ids.shape[-1], :, :, :] + + eagle_decoder_input = self.embedding( + input_ids=eagle_ids, + position_ids=eagle_position_ids, + ) + eagle_decoder_input, _ = self.eagle_linear_proj( + torch.cat((eagle_decoder_input, eagle_hidden_states), dim=-1) + ) + + new_hidden_states = self.eagle_decoder( + hidden_states=eagle_decoder_input, + attention_mask=eagle_attn_mask, + inference_params=None, + rotary_pos_emb=eagle_rotary_pos_emb, + packed_seq_params=None, + ) + + eagle_logits, _ = self.output_layer(new_hidden_states) + + all_logprob = gather_from_tensor_model_parallel_region( + eagle_logits[tree.relative_ids[i], :, :].transpose(0, 1) + ) + all_logprob = torch.softmax(all_logprob, dim=-1) + + for idx, tk in zip(tree.relative_ids[i], tree.top_k[i]): + if tk == 0: + continue + top_vals, top_ids = all_logprob[:, idx, :].topk(tk, dim=-1) + + new_tokens = torch.cat((new_tokens, top_ids), dim=-1) + eagle_ids = torch.cat((eagle_ids, top_ids), dim=-1) + + for _ in range(tk): + eagle_hidden_states = torch.cat( + (eagle_hidden_states, new_hidden_states[[idx], :, :]), dim=0 + ) + + return new_tokens diff --git a/modelopt/torch/speculative/plugins/transformers.py b/modelopt/torch/speculative/plugins/transformers.py new file mode 100644 index 0000000000..4cb3e94f93 --- /dev/null +++ b/modelopt/torch/speculative/plugins/transformers.py @@ -0,0 +1,620 @@ +# Adapted from: https://github.com/ctlllll/axolotl/blob/f86767e/src/axolotl/monkeypatch/medusa_utils.py +# +# 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. + +# 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. + +"""Support medusa for huggingface models.""" + +import contextlib +from typing import Any, Optional + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss +from transformers import Cache, DynamicCache, PreTrainedModel +from transformers.trainer_pt_utils import LabelSmoother +from transformers.utils import ModelOutput + +from ..eagle.conversion import EagleDMRegistry +from ..eagle.eagle_model import EagleModel +from ..eagle.utils import LlamaDecoderLayer, LlamaRMSNorm, expand_mask, make_causal_mask +from ..medusa.conversion import MedusaDMRegistry +from ..medusa.medusa_model import MedusaModel +from ..redrafter.conversion import RedrafterDMRegistry +from ..redrafter.redrafter_model import RedrafterModel +from ..utils import ResBlock + +IGNORE_TOKEN_ID = LabelSmoother.ignore_index + + +@MedusaDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) +class HFMedusaModel(MedusaModel): + """Medusa Model Class for huggingface models.""" + + def modify(self, medusa_num_heads=0, medusa_num_layers=0): + """Constructor. + + Args: + medusa_num_heads: number of medusa heads. + medusa_num_layers: number of ResBlock layers in each head. + """ + super().modify(medusa_num_heads=medusa_num_heads, medusa_num_layers=medusa_num_layers) + self.config.medusa = { + "num_medusa_heads": medusa_num_heads, + "num_medusa_layers": medusa_num_layers, + } + + hidden_size = self.lm_head.weight.shape[-1] + vocab_size = self.lm_head.weight.shape[0] + + # Create a list of Medusa heads + self.medusa_heads = nn.ModuleList( + [ + nn.Sequential( + *([ResBlock(hidden_size)] * self.medusa_num_layers), + nn.Linear(hidden_size, vocab_size, bias=False), + ) + for _ in range(self.medusa_num_heads) + ] + ) + + # Ensure medusa_head's dtype and device align with the base_model + self.medusa_heads.to(self.lm_head.weight.dtype).to(self.lm_head.weight.device) + self.medusa_heads.device = self.lm_head.weight.device + if hasattr(self, "hf_device_map") and "lm_head" in self.hf_device_map: + self.hf_device_map["medusa_heads"] = self.hf_device_map["lm_head"] + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, + freeze_base_model: bool = True, + medusa_heads_coefficient: Optional[float] = 0.2, + medusa_decay_coefficient: Optional[float] = 0.8, + ) -> Any: + """Forward pass of the MedusaModel. + + Returns: + torch.Tensor: A tensor containing predictions from all Medusa heads. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Pass input through the base model + with torch.no_grad() if freeze_base_model else contextlib.nullcontext(): + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + + medusa_logits = [self.medusa_heads[i](hidden_states) for i in range(self.medusa_num_heads)] + + if labels is not None: + loss = 0 + loss_fct = CrossEntropyLoss() + # Base model loss + if not freeze_base_model: + loss_logits = logits.view(-1, logits.shape[-1]) + loss_labels = labels.view(-1) + base_model_loss = loss_fct(loss_logits, loss_labels) + loss += base_model_loss + # Medusa loss + for i in range(self.medusa_num_heads): + labels = labels[..., 1:].contiguous() + loss_logits = medusa_logits[i][:, : -(1 + i)].contiguous() + loss_logits = loss_logits.view(-1, loss_logits.shape[-1]) + loss_labels = labels.view(-1) + loss += ( + loss_fct(loss_logits, loss_labels) + * medusa_decay_coefficient**i + * medusa_heads_coefficient + ) + else: + loss = None + + if not return_dict: + output = (logits, medusa_logits) + return (loss,) + output if loss is not None else output + + return ModelOutput( + loss=loss, + logits=logits, + medusa_logits=medusa_logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +class EagleModule(nn.Module): + """Eagle module used in EAGLE model.""" + + def __init__(self, eagle_config, bias=True): + """Init function for EagleModule.""" + super().__init__() + + self.fc = nn.Linear(2 * eagle_config["hidden_size"], eagle_config["hidden_size"], bias=bias) + self.layers = nn.ModuleList( + [ + LlamaDecoderLayer( + index, + eagle_config["hidden_size"], + eagle_config["intermediate_size"], + eagle_config["rms_norm_eps"], + eagle_config["num_attention_heads"], + eagle_config["num_key_value_heads"], + eagle_config["max_position_embeddings"], + eagle_config["rope_theta"], + eagle_config["use_input_layernorm_in_first_layer"], + ) + for index in range(eagle_config["num_hidden_layers"]) + ] + ) + if eagle_config["use_last_layernorm"]: + self.norm = LlamaRMSNorm(eagle_config["hidden_size"], eagle_config["rms_norm_eps"]) + + def forward( + self, + hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor, + lm_head: nn.Module, + attention_mask: Optional[torch.Tensor] = None, + loss_mask: Optional[torch.Tensor] = None, + logits: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = False, + ): + """Forward function for EagleModule.""" + batch_size, seq_length, _ = hidden_states.shape + seq_length_with_past = seq_length + past_key_values_length = 0 + + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + if position_ids is None: + device = hidden_states.device if hidden_states is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, + seq_length + past_key_values_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + if attention_mask is None: + attention_mask = torch.ones( + (batch_size, seq_length_with_past), dtype=torch.bool, device=hidden_states.device + ) + attention_mask = self._prepare_decoder_attention_mask( + attention_mask, (batch_size, seq_length), hidden_states, past_key_values_length + ) + + inputs_embeds = inputs_embeds.to(hidden_states.dtype).to(hidden_states.device) + hidden_states = self.fc(torch.cat((inputs_embeds, hidden_states), dim=-1)) + + next_decoder_cache = () + for idx, decoder_layer in enumerate(self.layers): + past_key_value = past_key_values[idx] if past_key_values is not None else None + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) + + if hasattr(self, "norm"): + hidden_states = self.norm(hidden_states) + + logits = lm_head(hidden_states).to(hidden_states.device) + + return hidden_states, logits, next_decoder_cache + + def _prepare_decoder_attention_mask( + self, attention_mask, input_shape, inputs_embeds, past_key_values_length + ): + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + combined_attention_mask = None + if input_shape[-1] > 1: + combined_attention_mask = make_causal_mask( + input_shape, + # inputs_embeds.dtype, + torch.float32, # [MODIFIED] force to cast to float32 + device=inputs_embeds.device, + past_key_values_length=past_key_values_length, + ) + + if attention_mask is not None: + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = expand_mask( + attention_mask, torch.float32, tgt_len=input_shape[-1] + ).to(inputs_embeds.device) + combined_attention_mask = ( + expanded_attn_mask + if combined_attention_mask is None + else expanded_attn_mask + combined_attention_mask + ) + + return combined_attention_mask + + +@EagleDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) +class HFEagleModel(EagleModel): + """Eagle Model Class for huggingface models.""" + + def modify(self, eagle_num_layers, use_input_layernorm_in_first_layer, use_last_layernorm): + """Constructor. + + Args: + config: The config for eagle decoder layers. + """ + super().modify( + eagle_num_layers=eagle_num_layers, + use_input_layernorm_in_first_layer=use_input_layernorm_in_first_layer, + use_last_layernorm=use_last_layernorm, + ) + + self.config.eagle = { + "num_hidden_layers": eagle_num_layers, + "num_attention_heads": self.config.num_attention_heads, + "intermediate_size": self.config.intermediate_size, + "hidden_size": self.config.hidden_size, + "num_key_value_heads": self.config.num_key_value_heads, + "rms_norm_eps": self.config.rms_norm_eps, + "max_position_embeddings": self.config.max_position_embeddings, + "rope_theta": self.config.rope_theta, + "use_input_layernorm_in_first_layer": use_input_layernorm_in_first_layer, + "use_last_layernorm": use_last_layernorm, + } + self.eagle_module = EagleModule(self.config.eagle) + + if hasattr(self.model.layers[-1].self_attn, "o_proj"): + device = self.model.layers[-1].self_attn.o_proj.weight.device + elif hasattr(self.model.layers[-1].self_attn, "q_proj"): + device = self.model.layers[-1].self_attn.q_proj.weight.device + elif hasattr(self.model.layers[-1].self_attn, "qkv_proj"): + device = self.model.layers[-1].self_attn.qkv_proj.weight.device + self.eagle_module.to(self.dtype).to(device) + + # Make sure self.model.embed_tokens and self.lm_head are frozen + for param in self.model.embed_tokens.parameters(): + param.requires_grad = False + for param in self.lm_head.parameters(): + param.requires_grad = False + + def forward( + self, + input_ids: torch.LongTensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, + loss_mask: torch.Tensor = None, + freeze_base_model: Optional[bool] = True, + classification_loss_coefficient: Optional[float] = 0.1, + regression_loss_coefficient: Optional[float] = 1.0, + ) -> Any: + """Forward pass of the EagleModel. + + Returns: + hidden_states: The hidden state from the base model. + logits: logits from the base model. + eagle_hidden_states: The hidden state from eagle_module. + eagle_logits: logits from the eagle_module. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + if past_key_values is not None and hasattr(past_key_values, "eagle_cache"): + eagle_cache = past_key_values.eagle_cache + else: + eagle_cache = None + + with torch.no_grad() if freeze_base_model else contextlib.nullcontext(): + outputs = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + labels=None, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=True, + return_dict=True, + cache_position=cache_position, + ) + past_key_values = outputs.past_key_values + if not isinstance(past_key_values, Cache): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + hidden_states = outputs.hidden_states[-1] + logits = outputs.logits + + # Shift left 1 token for eagle inputs + zeropadding = torch.zeros( + input_ids.shape[0], 1, dtype=input_ids.dtype, device=input_ids.device + ) + eagle_input_ids = torch.cat((input_ids[:, 1:], zeropadding), dim=1) + if attention_mask is not None: + zeropadding = torch.zeros( + attention_mask.shape[0], 1, dtype=attention_mask.dtype, device=attention_mask.device + ) + attention_mask = torch.cat((attention_mask[:, 1:], zeropadding), dim=1) + + with torch.no_grad(): + inputs_embeds = self.model.embed_tokens(eagle_input_ids) + eagle_hidden_states, eagle_logits, eagle_cache = self.eagle_module( + hidden_states, + inputs_embeds, + self.lm_head, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=eagle_cache, + use_cache=True, + output_attentions=output_attentions, + ) + if not isinstance(eagle_cache, Cache): + eagle_cache = DynamicCache.from_legacy_cache(eagle_cache) + past_key_values.eagle_cache = eagle_cache + + loss = None + if not freeze_base_model and labels is not None: + loss_fct = CrossEntropyLoss() + loss_logits = logits.view(-1, logits.shape[-1]) + labels = labels.view(-1) + base_model_loss = loss_fct(loss_logits, labels) + loss = base_model_loss + + if loss_mask is not None: + # Shift hidden_states and logits to align with eagle counterparts + zeropadding = torch.zeros( + hidden_states.shape[0], + 1, + hidden_states.shape[2], + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + hidden_states = torch.cat((hidden_states[:, 1:], zeropadding), dim=1).detach() + zeropadding = torch.zeros( + logits.shape[0], 1, logits.shape[2], dtype=logits.dtype, device=logits.device + ) + base_model_logits = torch.cat((logits[:, 1:], zeropadding), dim=1).detach() + + regression_loss, classification_loss = self._eagle_loss( + hidden_states, base_model_logits, eagle_hidden_states, eagle_logits, loss_mask + ) + eagle_loss = ( + regression_loss_coefficient * regression_loss + + classification_loss_coefficient * classification_loss + ) + if loss is None: + loss = eagle_loss + else: + loss += eagle_loss + + if not return_dict: + output = (logits, eagle_logits) + return (loss,) + output if loss is not None else output + + return ModelOutput( + loss=loss, + logits=logits, + eagle_logits=eagle_logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def _eagle_loss(self, hidden_states, logits, eagle_hidden_states, eagle_logits, loss_mask): + """Function for EAGLE loss computing.""" + loss_mask = loss_mask[:, :, None] + criterion = nn.SmoothL1Loss(reduction="none") + classification_loss = nn.Softmax(dim=2)(logits) * nn.LogSoftmax(dim=2)(eagle_logits) + classification_loss = -torch.sum(torch.sum(loss_mask * classification_loss, 2)) / ( + loss_mask.sum() + 1e-5 + ) + regression_loss = criterion(eagle_hidden_states, hidden_states) + regression_loss = torch.sum(torch.mean(loss_mask * regression_loss, 2)) / ( + loss_mask.sum() + 1e-5 + ) + return regression_loss, classification_loss + + +class Drafter(nn.Module): + """Wrapper class for the drafter in Redrafter.""" + + def __init__(self, redrafter_num_layers, hidden_size, vocab_size): + """Init function for drafter.""" + super().__init__() + + self.lm_head = torch.nn.Sequential( + *([ResBlock(2 * hidden_size) for _ in range(redrafter_num_layers)]), + torch.nn.Linear(in_features=2 * hidden_size, out_features=vocab_size, bias=False), + ) + + self.rnn_u = torch.nn.Linear(hidden_size, hidden_size, bias=True) + self.rnn_w = torch.nn.Linear(hidden_size, hidden_size, bias=False) + + def forward( + self, input_embs: torch.Tensor, cumsum_input_embs: torch.Tensor, hidden_states: torch.Tensor + ): + """Forward pass of the Drafter.""" + input_embs = torch.roll(input_embs, -1, dims=1) + o = self.rnn_u(cumsum_input_embs) + cumsum_input_embs = nn.SiLU()(o + self.rnn_w(input_embs)) + h = torch.cat((hidden_states, cumsum_input_embs), -1) + logits = self.lm_head(h) + + return input_embs, cumsum_input_embs, logits + + +@RedrafterDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"}) +class HFRedrafterModel(RedrafterModel): + """Redrafter Model Class for huggingface models.""" + + def modify(self, redrafter_predict_n_tokens=0, redrafter_num_layers=0): + """Constructor. + + Args: + redrafter_predict_n_tokens: number of tokens the drafter will predict. + redrafter_num_layers: number of ResBlock layers in lm head. + """ + super().modify( + redrafter_predict_n_tokens=redrafter_predict_n_tokens, + redrafter_num_layers=redrafter_num_layers, + ) + + hidden_size = self.lm_head.weight.shape[-1] + vocab_size = self.lm_head.weight.shape[0] + + self.config.redrafter = { + "num_draft_layers": redrafter_num_layers, + "hidden_size": hidden_size, + "exit_dim": 2 * hidden_size, + "rnn": True, + } + + self.drafter = Drafter(self.redrafter_predict_n_tokens, hidden_size, vocab_size) + + # Ensure drafter's dtype and device align with the base_model + self.drafter.to(self.lm_head.weight.dtype).to(self.lm_head.weight.device) + self.drafter.device = self.lm_head.weight.device + if hasattr(self, "hf_device_map") and "lm_head" in self.hf_device_map: + self.hf_device_map["drafter"] = self.hf_device_map["lm_head"] + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, + freeze_base_model: bool = True, + ) -> Any: + """Forward pass of the RedrafterModel.""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Pass input through the base model + with torch.no_grad() if freeze_base_model else contextlib.nullcontext(): + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + + redrafter_logits = [] + input_embs = self.model.embed_tokens(input_ids) + cumsum_input_embs = torch.zeros_like( + input_embs, dtype=input_embs.dtype, device=input_embs.device + ) + + for _ in range(self.redrafter_predict_n_tokens): + input_embs, cumsum_input_embs, drafter_logits = self.drafter( + input_embs, cumsum_input_embs, hidden_states + ) + redrafter_logits.append(drafter_logits) + + if labels is not None: + loss = 0 + loss_fct = CrossEntropyLoss() + # Base model loss + if not freeze_base_model: + loss_logits = logits.view(-1, logits.shape[-1]) + loss_labels = labels.view(-1) + base_model_loss = loss_fct(loss_logits, loss_labels) + loss += base_model_loss + # Drafter loss + for i in range(self.redrafter_predict_n_tokens): + labels = labels[..., 1:].contiguous() + loss_logits = redrafter_logits[i][:, : -(1 + i)].contiguous() + loss_logits = loss_logits.view(-1, loss_logits.shape[-1]) + loss_labels = labels.view(-1) + loss += loss_fct(loss_logits, loss_labels) + else: + loss = None + + if not return_dict: + output = (logits, redrafter_logits) + return (loss,) + output if loss is not None else output + + return ModelOutput( + loss=loss, + logits=logits, + redrafter_logits=redrafter_logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/modelopt/torch/speculative/redrafter/__init__.py b/modelopt/torch/speculative/redrafter/__init__.py new file mode 100644 index 0000000000..7c7bd2a90c --- /dev/null +++ b/modelopt/torch/speculative/redrafter/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""Redrafter Optimization Method.""" + +from .conversion import * +from .redrafter_model import * diff --git a/modelopt/torch/speculative/redrafter/conversion.py b/modelopt/torch/speculative/redrafter/conversion.py new file mode 100644 index 0000000000..d58f970885 --- /dev/null +++ b/modelopt/torch/speculative/redrafter/conversion.py @@ -0,0 +1,62 @@ +# 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. + +"""Redrafter conversion/restore utilities.""" + +from torch import nn + +from modelopt.torch.opt.conversion import ModelLikeModule +from modelopt.torch.opt.dynamic import _DMRegistryCls +from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict + +from ..config import RedrafterConfig + +RedrafterDMRegistry = _DMRegistryCls(prefix="Redrafter") # global instance for the registry + + +def convert_to_redrafter_model(model: nn.Module, config: RedrafterConfig) -> ConvertReturnType: + """Convert the model to a redrafter model as per `config`.""" + # initialize the true module if necessary + model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + + original_cls = type(model) + if original_cls not in RedrafterDMRegistry: + for cls in RedrafterDMRegistry._registry: + if issubclass(original_cls, cls): + RedrafterDMRegistry.register({original_cls: "base_model_class"})( + RedrafterDMRegistry[cls] + ) + break + + redrafter_model = RedrafterDMRegistry.convert(model) + redrafter_model.modify( + redrafter_predict_n_tokens=config.redrafter_predict_n_tokens, + redrafter_num_layers=config.redrafter_num_layers, + ) + + # no metadata, all specifed via config. + metadata = {} + + return redrafter_model, metadata + + +def restore_redrafter_model( + model: nn.Module, config: RedrafterConfig, metadata: MetadataDict +) -> nn.Module: + """Function for restoring a previously convert model to a redrafter model.""" + # the metadata should be empty + assert not metadata, "No metadata expected!" + + return convert_to_redrafter_model(model, config)[0] diff --git a/modelopt/torch/speculative/redrafter/redrafter_model.py b/modelopt/torch/speculative/redrafter/redrafter_model.py new file mode 100644 index 0000000000..df4365622a --- /dev/null +++ b/modelopt/torch/speculative/redrafter/redrafter_model.py @@ -0,0 +1,32 @@ +# 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. + +"""Redrafter model to support redrafter decoding.""" + +from modelopt.torch.opt.dynamic import DynamicModule + + +class RedrafterModel(DynamicModule): + """Base Redrafter Model.""" + + def _setup(self): + self._register_temp_attribute("redrafter_predict_n_tokens", 0) + self._register_temp_attribute("redrafter_num_layers", 0) + self._register_temp_attribute("drafter", None) + + def modify(self, redrafter_predict_n_tokens=0, redrafter_num_layers=0): + """Base Redrafter Model modify function. Child class should implement the details.""" + self.redrafter_predict_n_tokens = redrafter_predict_n_tokens + self.redrafter_num_layers = redrafter_num_layers diff --git a/modelopt/torch/speculative/speculative_decoding.py b/modelopt/torch/speculative/speculative_decoding.py new file mode 100644 index 0000000000..9e9c191537 --- /dev/null +++ b/modelopt/torch/speculative/speculative_decoding.py @@ -0,0 +1,55 @@ +# 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. + +"""User-facing API for converting a model into a `modelopt.torch.speculative.MedusaModel`.""" + +import torch.nn as nn + +from modelopt.torch.opt.conversion import apply_mode +from modelopt.torch.opt.mode import ModeLike + +from .mode import SpeculativeDecodingModeRegistry + +__all__ = ["convert"] + + +def convert(model: nn.Module, mode: ModeLike) -> nn.Module: + """Main conversion function to turn a base model into a speculative decoding model. + + Args: + model: The base model to be used. + mode: A (list of) string(s) or Mode(s) or a list of tuples containing the mode and its + config indicating the desired mode(s) (and configurations) for the convert + process. Modes set up the model for different algorithms for model optimization. The + following modes are available: + + * :class:`"medusa"`: The + ``model`` will be converted into a medusa model with added medusa head. + The mode's config is described in + :class:`MedusaConfig`. + * :class:`"eagle"`: The + ``model`` will be converted into a eagle model with added eagle weights. + The mode's config is described in + :class:`EagleConfig`. + + If the mode argument is specified as a dictionary, the keys should indicate the mode and + the values specify the per-mode configuration. + + Returns: + An instance of :class:`MedusaModel