trending/Model-Optimizer
1
0
Fork 0
mirror of https://github.com/NVIDIA/Model-Optimizer.git synced 2026-10-02 03:14:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update files on Github (#258)

Browse Source 
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
This commit is contained in:
Keval Morabia
2025-08-17 02:44:16 +05:30
committed by GitHub
parent c19ef08798
commit 37085354c9
95 changed files with 4362 additions and 1133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG-Windows.rst
+8
View File
@@ -2,6 +2,14 @@
Model Optimizer Changelog (Windows)
===================================
0.33 (2025-07-21)
^^^^^^^^^^^^^^^^^
**New Features**
- TensorRT Model Optimizer for Windows now supports `NvTensorRtRtx <https://onnxruntime.ai/docs/execution-providers/TensorRTRTX-ExecutionProvider.html>`_ execution-provider.
0.27 (2025-04-30)
^^^^^^^^^^^^^^^^^
CHANGELOG.rst
+4 -1
View File
@@ -8,13 +8,15 @@ Model Optimizer Changelog (Linux)
**Deprecations**
- Deprecate ``torch<2.5`` support.
- Deprecate ``torch<2.6`` support.
**New Features**
- (Experimental) Add quantization support for custom TensorRT op in ONNX models.
- Add support for Minifinetuning (MFT; https://arxiv.org/abs/2506.15702) self-corrective distillation, which enables training on small datasets with severely mitigated catastrophic forgetting.
- Add tree decoding support for Megatron Eagle models.
- For most VLMs, we now explicitly disable quant on the vision part so we add them to the excluded_modules during HF export.
- Add support for ``hidden_size`` and ``num_layers`` pruning for Megatron Core Mamba models in ``mcore_gpt_minitron`` mode.
0.33 (2025-07-14)
^^^^^^^^^^^^^^^^^
@@ -36,6 +38,7 @@ Model Optimizer Changelog (Linux)
- ModelOpt now supports quantization of tensor-parallel sharded Huggingface transformer models. This requires ``transformers>=4.52.0``.
- Support quantization of FSDP2 wrapped models and add FSDP2 support in the ``llm_qat`` example.
- Add NeMo 2 Simplified Flow examples for quantization aware training/distillation (QAT/QAD), speculative decoding, pruning & distillation.
- Fix a Qwen3 MOE model export issue.
0.31 (2025-06-04)
^^^^^^^^^^^^^^^^^
docs/source/getting_started/_installation_for_Linux.rst
+1 -1
View File
@@ -16,7 +16,7 @@ Latest Model Optimizer (``nvidia-modelopt``) currently has the following system
+-------------------------+-----------------------------+
| CUDA | >=12.0 |
+-------------------------+-----------------------------+
| PyTorch | >=2.4 |
| PyTorch | >=2.6 |
+-------------------------+-----------------------------+
| TensorRT-LLM (Optional) | 0.20 |
+-------------------------+-----------------------------+
docs/source/guides/6_save_load.rst
+7
View File
@@ -166,6 +166,13 @@ Here is an example of how to enable ModelOpt save/restore with the Huggingface A
# Save the ModelOpt-modified model architecture and weights using Huggingface APIs
model.save_pretrained(f"ModelOpt_{model_path}")
By default, the modelopt state is saved in the same directory as the model weights.
You can disable this by setting the ``save_modelopt_state`` to ``False`` in the ``save_pretrained`` API, as shown below:
.. code-block:: python
model.save_pretrained(f"ModelOpt_{model_path}", save_modelopt_state=False)
The model saved as above can be restored using the Huggingface ``from_pretrained`` API.
Do not forget to call :meth:`mto.enable_huggingface_checkpointing() <modelopt.torch.opt.plugins.huggingface.enable_huggingface_checkpointing>`
before loading the model. This needs to be done only once in the program.
docs/source/guides/8_autocast.rst
+1 -1
View File
@@ -2,7 +2,7 @@ AutoCast (ONNX)
###############
AutoCast is a tool for converting FP32 ONNX models to mixed precision FP32-FP16 or FP32-BF16 models.
While casting FP32 to FP6/BF16, some nodes might be more sensitive to effecting accuracy.
While casting FP32 to FP16/BF16, some nodes might be more sensitive to effecting accuracy.
AutoCast intelligently selects nodes to keep in FP32 precision to maintain model accuracy while benefiting from
reduced precision on the rest of the nodes. AutoCast automatically injects cast operations around the selected
nodes.
examples/llm_ptq/example_utils.py
+16 -58
View File
@@ -18,15 +18,10 @@ import sys
from typing import Any
import torch
import transformers
from accelerate import infer_auto_device_map, init_empty_weights
from accelerate.utils import get_max_memory
from transformers import (
AutoConfig,
AutoModelForCausalLM,
AutoProcessor,
AutoTokenizer,
Llama4ForConditionalGeneration,
)
from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer
from modelopt.torch.utils.image_processor import MllamaImageProcessor
@@ -148,7 +143,7 @@ def get_model(
if device == "cpu":
device_map = "cpu"
config_kwargs = {"trust_remote_code": trust_remote_code}
config_kwargs = {"trust_remote_code": trust_remote_code} if trust_remote_code else {}
if attn_implementation is not None:
config_kwargs["attn_implementation"] = attn_implementation
@@ -182,61 +177,24 @@ def get_model(
max_memory = {key: value * gpu_mem_percentage for key, value in max_memory.items()}
model_kwargs["max_memory"] = max_memory
if hf_config.model_type == "bart":
# device_map "auto" and "cuda" triggers error regarding meta tensor from safetensors
device_map = None
if is_speculative(hf_config):
model = AutoModelForCausalLM.from_pretrained(
ckpt_path,
device_map=device_map,
**model_kwargs,
)
elif hf_config.model_type == "llava":
from transformers import LlavaForConditionalGeneration
hf_llava = LlavaForConditionalGeneration.from_pretrained(
ckpt_path, device_map=device_map, **model_kwargs
)
model = hf_llava.language_model
elif hf_config.model_type == "t5":
from transformers import AutoModelForSeq2SeqLM
model = AutoModelForSeq2SeqLM.from_pretrained(
ckpt_path, device_map=device_map, **model_kwargs
)
elif hf_config.model_type == "bart":
from transformers import AutoModelForSeq2SeqLM
# device_map "auto" and "cuda" triggers error regarding meta tensor from safetensors
model = AutoModelForSeq2SeqLM.from_pretrained(
ckpt_path, device_map=None, **model_kwargs
).to(device)
elif hf_config.model_type == "whisper":
from transformers import WhisperForConditionalGeneration
model = WhisperForConditionalGeneration.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,
)
elif hf_config.model_type == "mllama":
from transformers import MllamaForConditionalGeneration
model = MllamaForConditionalGeneration.from_pretrained(
ckpt_path,
device_map=device_map,
**model_kwargs,
)
elif hf_config.model_type == "llama4":
model = Llama4ForConditionalGeneration.from_pretrained(
ckpt_path,
device_map=device_map,
**model_kwargs,
)
else:
architecture = hf_config.architectures[0]
assert hasattr(transformers, architecture), (
f"Architecture {architecture} not found in transformers: {transformers.__version__}"
)
auto_model_module = getattr(transformers, architecture)
with init_empty_weights():
# When computing the device_map, assuming half precision by default,
# unless specified by the hf_config.
@@ -246,7 +204,7 @@ def get_model(
# DeciLMForCausalLM does not support max_memory argument
if "architectures" in hf_config and "DeciLMForCausalLM" in hf_config.architectures:
model_kwargs2.pop("max_memory", None)
model = AutoModelForCausalLM.from_config(
model = auto_model_module._from_config(
hf_config,
**model_kwargs2,
)
@@ -269,7 +227,7 @@ def get_model(
)
model_kwargs["max_memory"] = max_memory
model = AutoModelForCausalLM.from_pretrained(
model = auto_model_module.from_pretrained(
ckpt_path,
device_map=device_map,
**model_kwargs,
examples/llm_ptq/hf_ptq.py
+44 -30
View File
@@ -46,6 +46,7 @@ from modelopt.torch.utils.dataset_utils import (
create_forward_loop,
get_dataset_dataloader,
get_max_batch_size,
get_supported_datasets,
)
from modelopt.torch.utils.image_processor import MllamaImageProcessor
from modelopt.torch.utils.memory_monitor import launch_memory_monitor
@@ -195,6 +196,9 @@ def main(args):
# launch a memory monitor to read the currently used GPU memory.
launch_memory_monitor()
# Force eager execution for all model types.
torch.compiler.set_stance("force_eager")
# Check that only one quantization format is provided for non auto_quant case
if not args.auto_quantize_bits:
assert len(args.qformat.split(",")) == 1, (
@@ -267,14 +271,6 @@ def main(args):
full_model = model
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,
model_type,
@@ -283,20 +279,12 @@ def main(args):
attn_implementation=args.attn_implementation,
)
elif model_type == "whisper":
if args.dataset is None:
args.dataset = "peoples_speech"
warnings.warn(
"Currently only the peoples_speech dataset is supported for the whisper model. "
"Overriding dataset to peoples_speech."
)
elif args.dataset != "peoples_speech":
raise ValueError("Only the peoples_speech dataset is supported for the whisper model.")
processor = get_processor(
args.pyt_ckpt_path, model_type, device, trust_remote_code=args.trust_remote_code
)
else:
if args.dataset is None:
args.dataset = "cnn_dailymail"
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
@@ -305,16 +293,31 @@ def main(args):
# We only quantize the language model for VLMs other than the type supported above.
if hasattr(model, "language_model"):
assert model_type == "llama4", (
"Only llama4 should reach here. Please uncomment this check if you are modelopt developers."
)
parent_model = model # llama4 case
if isinstance(type(model).__dict__.get("language_model"), property):
assert hasattr(model, "model") and hasattr(model.model, "language_model"), (
"Expected language_model in model.model, but attribute not found. "
"This may indicate an unsupported model structure."
)
parent_model = model.model # gemma3, qwen2.5 VL case
disabled_quant_cfg = {
"quant_cfg": {"default": {"enable": False}},
"algorithm": "max",
}
for name, child in parent_model.named_children():
# Apply disabled quant to all children except language_model so we can exclude them during HF export.
if name != "language_model":
mtq.quantize(child, disabled_quant_cfg, forward_loop=None)
model = model.language_model
if args.sparsity_fmt != "dense":
if args.batch_size == 0:
# Sparse algorithm takes more GPU memory so we reduce the batch_size by 4.
args.batch_size = max(get_max_batch_size(model) // 4, 1)
args.batch_size = min(args.batch_size, args.calib_size)
args.batch_size = min(args.batch_size, sum(args.calib_size))
print(f"Use calib batch_size {args.batch_size}")
@@ -373,7 +376,7 @@ def main(args):
sample_input_single_batch=sample_input_single_batch,
enable_grad=run_auto_quant,
)
args.batch_size = min(args.batch_size, args.calib_size)
args.batch_size = min(args.batch_size, sum(args.calib_size))
print(f"Use calib batch_size {args.batch_size}")
@@ -383,17 +386,17 @@ def main(args):
"The MllamaImageProcessor must be set."
)
calib_dataloader = get_vlm_dataset_dataloader(
dataset_name=args.dataset,
dataset_name=args.dataset[0] if args.dataset else "scienceqa",
processor=processor,
batch_size=args.batch_size,
num_samples=args.calib_size,
num_samples=args.calib_size[0],
)
elif model_type == "whisper":
assert processor is not None and isinstance(processor, WhisperProcessor), (
"The AutoProcessor must be set."
)
calib_dataloader, first_text = get_speech_dataset_dataloader(
dataset_name=args.dataset,
dataset_name=args.dataset[0] if args.dataset else "peoples_speech",
processor=processor,
batch_size=args.batch_size,
num_samples=args.calib_size,
@@ -454,7 +457,7 @@ def main(args):
"input_features" if model_type == "whisper" else "input_ids"
][0:1]
try:
generated_ids_before_ptq = model.generate(input_ids, max_new_tokens=100)
generated_ids_before_ptq = full_model.generate(input_ids, max_new_tokens=100)
except Exception as e:
print(
"Error during model generation. Please check if your transformers version is "
@@ -472,7 +475,8 @@ def main(args):
torch.cuda.empty_cache()
generated_ids_after_ptq = None
if model_type != "llama4":
generated_ids_after_ptq = model.generate(input_ids, max_new_tokens=100)
# Our fake quantizer may not be fully compatible with torch.compile.
generated_ids_after_ptq = full_model.generate(input_ids, max_new_tokens=100)
else:
warnings.warn(
"Llama4 Maverick generation after quantization has a bug. Skipping generation sample."
@@ -600,15 +604,23 @@ if __name__ == "__main__":
default=0,
)
parser.add_argument(
"--calib_size", help="Number of samples for calibration.", type=int, default=512
"--calib_size",
help=(
"Number of samples for calibration. If a comma separated list of values is provided, "
"each value will be used as the calibration size for the corresponding dataset."
),
type=str,
default="512",
)
parser.add_argument("--export_path", default="exported_model")
parser.add_argument(
"--dataset",
help="name of dataset.",
help=(
f"name of a dataset, or a comma separated list of datasets. "
f"dataset choices are {get_supported_datasets()}"
),
type=str,
default=None,
choices=["magpie", "cnn_dailymail", "pile", "pg19", "wikipedia"],
)
parser.add_argument("--inference_tensor_parallel", type=int, default=1)
parser.add_argument("--inference_pipeline_parallel", type=int, default=1)
@@ -695,4 +707,6 @@ if __name__ == "__main__":
args = parser.parse_args()
args.dataset = args.dataset.split(",") if args.dataset else None
args.calib_size = [int(num_sample) for num_sample in args.calib_size.split(",")]
main(args)
examples/llm_ptq/run_tensorrt_llm.py
+1 -1
View File
@@ -66,7 +66,7 @@ def run(args):
print("TensorRT-LLM example outputs:")
llm = LLM(args.engine_dir, tokenizer=tokenizer)
llm = LLM(args.engine_dir, tokenizer=tokenizer, max_batch_size=len(input_texts))
torch.cuda.cudart().cudaProfilerStart()
outputs = llm.generate_text(input_texts, args.max_output_len)
torch.cuda.cudart().cudaProfilerStop()
examples/llm_qat/launch.sh
+4 -3
View File
@@ -166,13 +166,14 @@ if [[ "${DISTILL}" == "True" ]]; then
FSDP_ARGS="$FSDP_ARGS --fsdp_cpu_ram_efficient_loading False"
fi
# real quantization does not work with FSDP
if [[ "${COMPRESS,,}" == "true" ]]; then
echo "Compression is not supported with FSDP. Disabling FSDP."
# real quantization does not work with FSDP, only works with FSDP2
if [[ "${COMPRESS,,}" == "true" && "${USE_FSDP2,,}" != "true" ]]; then
echo "Compression is not supported with FSDP. Disabling FSDP and using DDP."
FSDP_ARGS=""
CONFIG_FILE="ddp.yaml"
fi
CMD="accelerate launch --config-file accelerate_config/$CONFIG_FILE $FSDP_ARGS \
main.py \
--model_name_or_path $MODEL \
examples/llm_qat/main.py
+7 -1
View File
@@ -189,6 +189,10 @@ def train():
)
tokenizer.pad_token_id = tokenizer.eos_token_id
# We set model.config.use_cache to False for training when gradient_checkpointing=False.
# Currently useful for FSDP2 to allow for setting activation_checkpointing=True in the config file.åå
model.config.use_cache = False
print_rank_0("Loading dataset...")
data_module = make_supervised_data_module(
dataset=data_args.dataset,
@@ -243,7 +247,9 @@ def train():
distill_kwargs["distill_config"] = distill_config
trainer_cls = QADTrainer if training_args.distill else QATTrainer
training_args.lora_config = get_lora_config()
if training_args.lora:
training_args.lora_config = get_lora_config()
trainer = trainer_cls(
model=model,
processing_class=tokenizer,
examples/onnx_ptq/torch_quant_to_onnx.py
+1 -1
View File
@@ -138,7 +138,7 @@ def main():
# Quantize model
quantized_model = quantize_model(model, config, data_loader)
use_autocast = args.quantize_mode == "mxfp8"
use_autocast = args.quantize_mode != "mxfp8"
# Export to ONNX
export_to_onnx(
examples/speculative_decoding/ar_validate.py
+82
View File
@@ -0,0 +1,82 @@
# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
from accelerate import Accelerator
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
import modelopt.torch.opt as mto
from modelopt.torch.speculative.plugins.transformers import HFARValidation
mto.enable_huggingface_checkpointing()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, required=True, help="Path to model directory")
parser.add_argument("--steps", type=int, default=1, help="Steps for AR validation")
parser.add_argument(
"--osl", type=int, default=100, help="Output sequence length for AR validation"
)
parser.add_argument(
"--num_samples", type=int, default=20, help="Number of MT-Bench samples to use"
)
args = parser.parse_args()
accelerator = Accelerator()
# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(args.model_path, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(args.model_path)
model.eval()
model = accelerator.prepare(model)
validator = HFARValidation(model, tokenizer)
# Load MT-Bench prompts from HuggingFace
ds = load_dataset("HuggingFaceH4/mt_bench_prompts")["train"]
num_samples = min(args.num_samples, len(ds))
ars = []
for i in range(num_samples):
prompt = ds[i]["prompt"][0]
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(accelerator.device)
# Apply chat template to the prompt, continuing with assistant response
if hasattr(tokenizer, "apply_chat_template"):
chat_messages = [
{"role": "user", "content": prompt},
]
prompt = tokenizer.apply_chat_template(
chat_messages, tokenize=False, add_generation_prompt=True
)
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(accelerator.device)
# validate AR
_, ar = validator.validate(args.osl, input_ids=input_ids, steps=args.steps)
ars.append(ar)
if accelerator.is_main_process:
print(f"[{i + 1}/{num_samples}] Prompt: {prompt[:60]}... | AR: {ar:.4f}")
if ars and accelerator.is_main_process:
avg_ar = sum(ars) / len(ars)
print("\n==== AR Validation Results on MT-Bench ====")
print(f"Number of samples: {len(ars)}")
print(f"Output Sequence Length: {args.osl}")
print(f"Steps: {args.steps}")
print(f"Average AR: {avg_ar:.4f}")
if __name__ == "__main__":
main()
examples/speculative_decoding/calibrate_draft_vocab.py
+59
View File
@@ -0,0 +1,59 @@
# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import json
import os
import torch
from transformers import AutoTokenizer
from modelopt.torch.speculative.utils import calibrate_frequent_vocab
def main():
parser = argparse.ArgumentParser(description="Calibrate draft vocab and save to .pt file")
parser.add_argument("--model", type=str, required=True, help="Model name or path for tokenizer")
parser.add_argument("--data", type=str, required=True, help="Path to training data (jsonl)")
parser.add_argument("--draft_vocab_size", type=int, required=True, help="Draft vocab size")
parser.add_argument(
"--calibrate_size",
type=int,
default=None,
help="Number of samples to use for calibration. If None, use all dataset.",
)
parser.add_argument(
"--save_dir", type=str, default="draft_vocab_cache", help="Path to save .pt file"
)
args = parser.parse_args()
print("Calibrating vocab...")
tokenizer = AutoTokenizer.from_pretrained(args.model)
with open(args.data) as f:
conversations = [json.loads(line)["conversations"] for line in f]
if args.calibrate_size:
conversations = conversations[: args.calibrate_size]
conversations = [item for sublist in conversations for item in sublist]
d2t = calibrate_frequent_vocab(tokenizer, conversations, args.draft_vocab_size)
model_name = os.path.basename(os.path.normpath(args.model))
vocab_path = os.path.join(args.save_dir, model_name, "d2t.pt")
os.makedirs(os.path.dirname(vocab_path), exist_ok=True)
torch.save(d2t, vocab_path)
print(f"Saved calibrated vocab to {vocab_path}")
if __name__ == "__main__":
main()
examples/speculative_decoding/launch.sh
+12 -2
View File
@@ -18,6 +18,10 @@ set -eo pipefail
while [ $# -gt 0 ]; do
case "$1" in
--training_seq_len*)
if [[ "$1" != *=* ]]; then shift; fi
TRAINING_SEQ_LEN="${1#*=}"
;;
--model*)
if [[ "$1" != *=* ]]; then shift; fi
MODEL="${1#*=}"
@@ -62,6 +66,10 @@ while [ $# -gt 0 ]; do
if [[ "$1" != *=* ]]; then shift; fi
EAGLE_NUM_LAYERS="${1#*=}"
;;
--draft_vocab_size*)
if [[ "$1" != *=* ]]; then shift; fi
DRAFT_VOCAB_SIZE="${1#*=}"
;;
--fsdp_transformer_layer_cls_to_wrap*)
if [[ "$1" != *=* ]]; then shift; fi
FSDP_TRANSFORMER_LAYER_CLS_TO_WRAP="${1#*=}"
@@ -99,16 +107,18 @@ TRAIN_BS=${TRAIN_BS:-4}
MEDUSA_NUM_HEADS=${MEDUSA_NUM_HEADS:-1}
MEDUSA_NUM_LAYERS=${MEDUSA_NUM_LAYERS:-1}
EAGLE_NUM_LAYERS=${EAGLE_NUM_LAYERS:-1}
DRAFT_VOCAB_SIZE=${DRAFT_VOCAB_SIZE:-0}
REDRAFTER_TOKENS=${REDRAFTER_TOKENS:-1}
REDRAFTER_NUM_LAYERS=${REDRAFTER_NUM_LAYERS:-1}
FSDP_TRANSFORMER_LAYER_CLS_TO_WRAP=${FSDP_TRANSFORMER_LAYER_CLS_TO_WRAP:-"LlamaDecoderLayer"}
NUM_GPU=${NUM_GPU:-1}
DO_EVAL=${DO_EVAL:-"True"}
TRAINING_SEQ_LEN=${TRAINING_SEQ_LEN:-2048}
if [[ "$MODE" == "medusa" ]]; then
SPECULATIVE_ARGS="--medusa_num_heads $MEDUSA_NUM_HEADS --medusa_num_layers $MEDUSA_NUM_LAYERS"
elif [[ "$MODE" == "eagle" ]]; then
SPECULATIVE_ARGS="--eagle_num_layers $EAGLE_NUM_LAYERS"
SPECULATIVE_ARGS="--eagle_num_layers $EAGLE_NUM_LAYERS --draft_vocab_size $DRAFT_VOCAB_SIZE"
else
echo "Only medusa and eagle supported for now!"
exit 1
@@ -125,7 +135,7 @@ fi
CMD="accelerate launch $MULTI_GPU --mixed_precision bf16 main.py \
--mode $MODE \
--model_name_or_path $MODEL \
--model_max_length 2048 \
--training_seq_len $TRAINING_SEQ_LEN \
--dataloader_drop_last True \
--bf16 True \
--output_dir $OUTPUT_DIR \
examples/speculative_decoding/main.py
+49 -3
View File
@@ -29,6 +29,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
from dataclasses import dataclass, field
from typing import Literal
@@ -42,6 +43,7 @@ from transformers.trainer_utils import get_last_checkpoint
import modelopt.torch.opt as mto
import modelopt.torch.speculative as mtsp
from modelopt.torch.speculative.utils import calibrate_frequent_vocab
from modelopt.torch.utils import print_rank_0
torch.manual_seed(0)
@@ -60,12 +62,20 @@ class DataArguments:
)
eval_data_path: str = field(default=None, metadata={"help": "Path to the evaluation data."})
lazy_preprocess: bool = True
draft_vocab_cache_dir: str = field(
default="draft_vocab_cache",
metadata={"help": "Path to the d2t cache directory."},
)
calibrate_size: int = field(
default=None,
metadata={"help": "Size of the calibration data. If None, use entire training set."},
)
@dataclass
class TrainingArguments(transformers.TrainingArguments):
cache_dir: str | None = field(default=None)
model_max_length: int = field(
training_seq_len: int = field(
default=2048,
metadata={
"help": (
@@ -88,7 +98,9 @@ class MedusaArguments:
class EagleArguments:
eagle_num_layers: int | None = field(default=1)
use_input_layernorm_in_first_layer: bool | None = field(default=True)
use_last_layernorm: bool | None = field(default=False)
use_last_layernorm: bool | None = field(default=True)
use_aux_hidden_state: bool | None = field(default=True)
draft_vocab_size: int | None = field(default=32000)
def train():
@@ -127,7 +139,7 @@ def train():
)
tokenizer = transformers.AutoTokenizer.from_pretrained(
model_args.model_name_or_path,
model_max_length=training_args.model_max_length,
model_max_length=training_args.training_seq_len,
)
if tokenizer.chat_template is None:
tokenizer.chat_template = (
@@ -149,8 +161,42 @@ def train():
"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,
"use_aux_hidden_state": eagle_args.use_aux_hidden_state,
"draft_vocab_size": eagle_args.draft_vocab_size,
}
mtsp.convert(model, [("eagle", config)])
if eagle_args.draft_vocab_size > 0 and (
not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
):
model_name = os.path.basename(os.path.normpath(model_args.model_name_or_path))
vocab_cache_path = os.path.join(
data_args.draft_vocab_cache_dir, model_name, "d2t.pt"
)
if os.path.exists(vocab_cache_path):
vocab_cache = torch.load(vocab_cache_path)
if len(vocab_cache) == eagle_args.draft_vocab_size:
model.eagle_module.d2t = vocab_cache
print_rank_0(f"Loaded draft vocab cache from {vocab_cache_path}.")
else:
print_rank_0(
"No matching draft vocab cache found, calibrating vocab using training set..."
)
with open(data_args.data_path) as f:
calibrate_conversations = [json.loads(line)["conversations"] for line in f]
if data_args.calibrate_size:
calibrate_conversations = calibrate_conversations[
: data_args.calibrate_size
]
calibrate_conversations = [
item for sublist in calibrate_conversations for item in sublist
]
model.eagle_module.d2t = calibrate_frequent_vocab(
tokenizer, calibrate_conversations, eagle_args.draft_vocab_size
)
else:
raise Exception(f"{training_args.mode} is not supported!")
examples/speculative_decoding/server_generate.py
+1 -1
View File
@@ -91,7 +91,7 @@ def generate_data(messages, idx, system_prompt):
else:
raise ValueError(f"Message format not recognized: {message}")
if role != "user":
if role not in ["user", "human"]:
return
output_messages.append(
{
examples/windows/onnx_ptq/genai_llm/README.md
+8
View File
@@ -116,3 +116,11 @@ Please refer to [support matrix](https://nvidia.github.io/TensorRT-Model-Optimiz
1. **Check Input Model**
During INT4 AWQ execution, the input onnx model (one mentioned in `--onnx_path` argument) will be run with onnxruntime (ORT) for calibration (using ORT EP mentioned in `--calibration_eps` argument). So, make sure that input onnx model is running fine with the specified ORT EP.
1. **Config availability for calibration with NvTensorRtRtx EP**
Note that while using `NvTensorRtRtx` for INT4 AWQ quantization, profile (min/max/opt ranges) of input-shapes of the model is created internally using the details from the model's config (e.g. config.json in HuggingFace model card). This input-shapes-profile is used during onnxruntime session creation. Make sure that config.json is available in the model-directory if `model_name` is a local model path (instead of HuggingFace model-name).
1. **Error - Invalid Position-IDs input to the ONNX model**
The ONNX models produced using ONNX GenerativeAI (GenAI) have different IO bindings for models produced using different execution-providers (EPs). For instance, model built with DML EP has position-ids input in the ONNX model but models builts using CUDA EP or NvTensorRtRtx EP don't have position-ids inputs. So, set `add_position_ids` command-line argument to `true` or `false` depending on the base model, or set that value (hard-code) in the quantize script if required.
examples/windows/onnx_ptq/genai_llm/quantize.py
+64
View File
@@ -33,6 +33,59 @@ logging.getLogger().setLevel(logging.INFO)
pt_to_np = {"torch.int64": np.int64, "torch.float32": np.float32, "torch.float16": np.float16}
def prepare_input_shapes_string(
batch_size, seq_len, past_seq_len, num_layers, num_kv_heads, head_dim
):
shapes = ""
shapes += f"input_ids:{batch_size}x{seq_len}"
shapes += f",attention_mask:{batch_size}x{seq_len}"
for i in range(num_layers):
key_name = f"past_key_values.{i}.key"
value_name = f"past_key_values.{i}.value"
shapes += f",{key_name}:{batch_size}x{num_kv_heads}x{past_seq_len}x{head_dim}"
shapes += f",{value_name}:{batch_size}x{num_kv_heads}x{past_seq_len}x{head_dim}"
return shapes
def get_input_shapes_profile(model_name_or_path):
config = AutoConfig.from_pretrained(model_name_or_path)
head_dim = config.hidden_size // config.num_attention_heads
if hasattr(config, "head_dim") and config.head_dim is not None:
head_dim = config.head_dim
num_kv_heads = config.num_key_value_heads
num_layers = config.num_hidden_layers
min_shapes = prepare_input_shapes_string(1, 1, 0, num_layers, num_kv_heads, head_dim)
max_shapes = prepare_input_shapes_string(1, 1024, 1024, num_layers, num_kv_heads, head_dim)
opt_shapes = prepare_input_shapes_string(1, 512, 512, num_layers, num_kv_heads, head_dim)
return min_shapes, max_shapes, opt_shapes
def make_input_shapes_profile_for_ep_list(ep_list, model_name_or_path):
# Input-shapes-profile will be used in provider-options for ORT session creation.
# Provider options (even if {}) are needed for all EPs when we provide for any one of them.
# Using empty shapes_profile for non-NvTensorRtRtx EPs.
input_shapes_profile_sequence = []
for ep in ep_list:
if ep == "NvTensorRtRtx":
min_shapes, max_shapes, opt_shapes = get_input_shapes_profile(model_name_or_path)
input_shapes_profile = {
"nv_profile_min_shapes": min_shapes,
"nv_profile_max_shapes": max_shapes,
"nv_profile_opt_shapes": opt_shapes,
}
input_shapes_profile_sequence.append(input_shapes_profile)
else:
input_shapes_profile_sequence.append({})
return input_shapes_profile_sequence
def make_model_input(
config,
input_ids_arg,
@@ -341,6 +394,16 @@ def main(args):
args.trust_remote_code,
)
input_shapes_profile_data = None
if "NvTensorRtRtx" in args.calibration_eps and (args.algo not in ["rtn", "rtn_dq"]):
# NvTensorRtRtx EP uses (min, max, opt) profile for dynamic shapes in the model's inputs.
input_shapes_profile_data = make_input_shapes_profile_for_ep_list(
args.calibration_eps, args.model_name
)
print(
f"\n--Quantize-Script-- input_shapes_profile is None? - {input_shapes_profile_data is None}\n"
)
t = time.time()
logging.info("\nQuantizing the model....\n")
quantized_onnx_model = quantize_int4(
@@ -350,6 +413,7 @@ def main(args):
calibration_eps=args.calibration_eps,
use_zero_point=args.use_zero_point,
block_size=args.block_size,
input_shapes_profile=input_shapes_profile_data,
awqlite_alpha_step=args.awqlite_alpha_step,
awqlite_run_per_subgraph=args.awqlite_run_per_subgraph,
awqlite_fuse_nodes=args.awqlite_fuse_nodes,
modelopt/deploy/llm/generate.py
+29 -7
View File
@@ -25,6 +25,7 @@ import torch
from packaging.version import Version
from tensorrt_llm import SamplingParams
from tensorrt_llm.bindings.executor import DecodingConfig
from tensorrt_llm.llmapi import CudaGraphConfig
from tensorrt_llm.llmapi import KvCacheConfig as TRT_KvCacheConfig
from tensorrt_llm.llmapi.llm import LLM as TRT_LLM
from tensorrt_llm.llmapi.tokenizer import TokenizerBase, TransformersTokenizer
@@ -49,11 +50,12 @@ class LLM(TRT_LLM):
"""A wrapper over the ``tensorrt_llm.llmapi.llm.LLM`` for LLM profiling and validation."""
def _build_trt_llm_from_config(
self, config, engine_dir, tokenizer, kv_cache_config, medusa_choices
self, config, engine_dir, tokenizer, kv_cache_config, medusa_choices, max_batch_size
):
build_config = config["build_config"]
world_size = config.get("pretrained_config", {}).get("mapping", {}).get("world_size", 1)
max_tokens_kv_cache = build_config["max_seq_len"] * build_config["max_batch_size"]
max_batch_size = max(max_batch_size, build_config["max_batch_size"])
max_tokens_kv_cache = build_config["max_seq_len"] * max_batch_size
trt_kv_cache_config = TRT_KvCacheConfig(enable_block_reuse=False)
@@ -87,7 +89,9 @@ class LLM(TRT_LLM):
**kwargs,
)
def _build_torch_llm_from_config(self, checkpoint_dir, tokenizer, tp, trust_remote_code):
def _build_torch_llm_from_config(
self, checkpoint_dir, tokenizer, tp, trust_remote_code, max_batch_size
):
kwargs = {}
if tokenizer is not None:
kwargs["tokenizer"] = tokenizer
@@ -100,6 +104,15 @@ class LLM(TRT_LLM):
enable_block_reuse=False, free_gpu_memory_fraction=0.85
)
cuda_graph_config = None
if max_batch_size > 0:
cuda_graph_config = CudaGraphConfig(
batch_sizes=[2**i for i in range(int((max_batch_size - 1).bit_length()))]
+ [max_batch_size],
max_batch_size=max_batch_size,
enable_padding=True,
)
super().__init__(
backend="pytorch",
model=checkpoint_dir,
@@ -108,8 +121,7 @@ class LLM(TRT_LLM):
enable_chunked_prefill=True,
kv_cache_config=trt_kv_cache_config,
# pytorch backend configs
use_cuda_graph=True,
cuda_graph_padding_enabled=True,
cuda_graph_config=cuda_graph_config,
**kwargs,
)
@@ -121,6 +133,7 @@ class LLM(TRT_LLM):
medusa_choices: Any = None,
tp: int = 0,
trust_remote_code: bool = False,
max_batch_size: int = 0,
):
"""Initializes the LLM runner class.
@@ -132,6 +145,8 @@ class LLM(TRT_LLM):
medusa_choices: The medusa choices for the decoding config.
tp: the tensor parallel size (for the torch backend). If 0, it will be set to the number of GPUs.
trust_remote_code: whether to trust the remote code (for the torch backend).
max_batch_size: Max batch size for the LLM backend. If 0, it will be set to the max batch size
in the engine config.
"""
assert Version(tensorrt_llm.__version__) >= Version("0.17.0")
@@ -140,7 +155,12 @@ class LLM(TRT_LLM):
if "build_config" in config:
self._build_trt_llm_from_config(
config, checkpoint_dir, tokenizer, kv_cache_config, medusa_choices
config,
checkpoint_dir,
tokenizer,
kv_cache_config,
medusa_choices,
max_batch_size,
)
self._is_torch = False
@@ -152,7 +172,9 @@ class LLM(TRT_LLM):
"medusa_choices is not supported with the torch llmapi"
)
self._build_torch_llm_from_config(checkpoint_dir, tokenizer, tp, trust_remote_code)
self._build_torch_llm_from_config(
checkpoint_dir, tokenizer, tp, trust_remote_code, max_batch_size
)
self._is_torch = True
self._max_seq_len = config["max_position_embeddings"]
self._max_beam_width = 1
modelopt/onnx/autocast/convert.py
+2 -1
View File
@@ -169,9 +169,10 @@ def convert_to_f16(
"""
assert low_precision_type in ["fp16", "bf16"], "low_precision_type must be either fp16 or bf16"
# Opset 21 is needed for NVFP4 quantization support (DQ with 'block_size' attribute)
sanitizer = GraphSanitizer(
model,
min_opset=19,
min_opset=21,
trt_plugins=trt_plugins,
max_ir_version=LATEST_IR_VERSION_SUPPORTED_BY_ORT,
)
modelopt/onnx/autocast/precisionconverter.py
+12 -6
View File
@@ -210,6 +210,8 @@ class PrecisionConverter:
def _get_np_type(node, inp, opset=onnx.defs.onnx_opset_version()):
if node.op == "Cast":
return helper.tensor_dtype_to_np_dtype(node.attrs["to"])
elif node.op == "DequantizeLinear":
return node.inputs[1].dtype # scale type
elif not inp.dtype or inp.dtype == onnx.TensorProto.UNDEFINED:
return None
elif node.op not in self.custom_ops:
@@ -226,12 +228,16 @@ class PrecisionConverter:
return None
def _can_propagate_type(from_type, to_type):
from_type_onnx = helper.np_dtype_to_tensor_dtype(from_type)
to_type_onnx = helper.np_dtype_to_tensor_dtype(to_type)
return (
from_type_onnx in [*ONNX_TYPES, onnx.TensorProto.UNDEFINED]
and to_type_onnx in ONNX_TYPES
)
try:
from_type_onnx = helper.np_dtype_to_tensor_dtype(from_type)
to_type_onnx = helper.np_dtype_to_tensor_dtype(to_type)
return (
from_type_onnx in [*ONNX_TYPES, onnx.TensorProto.UNDEFINED]
and to_type_onnx in ONNX_TYPES
)
except Exception as e:
logger.warning(f"Failed to check if type can be propagated: {e}")
return False
def _propagate_cast_type_through_nodes(node, np_type, iter=1):
# Return if node is of cast type (from iter=2)
modelopt/onnx/quantization/graph_utils.py
+1
View File
@@ -354,6 +354,7 @@ def filter_quantizable_kgen_heads(
"BatchNormalization",
"GlobalAveragePool",
"MaxPool",
"Mul", # Example: VoVNet
]
for partition in kgen_partitions:
modelopt/onnx/quantization/int4.py
+13 -3
View File
@@ -21,6 +21,7 @@ import math
import os
import tempfile
import time
from collections.abc import Sequence
from typing import Any, cast
import numpy
@@ -424,6 +425,7 @@ def _quantize_awq_clip(
block_size: int,
force_fp16: bool = False,
nodes_to_exclude: list[str] = [],
input_shapes_profile: Sequence[dict[str, str]] | None = None,
**kwargs: Any,
) -> onnx.ModelProto:
"""Quantizes `onnx_model` using the Activation aware quantization a.k.a AWQ algorithm."""
@@ -453,7 +455,7 @@ def _quantize_awq_clip(
logger.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)
session = create_inference_session(augmented_onnx_path, calibration_eps, input_shapes_profile)
inputs = []
for inp_d in data_reader:
inputs.append(inp_d)
@@ -907,6 +909,7 @@ def _quantize_awq_lite(
enable_weight_clipping: bool = False,
use_zero_point: bool = False,
nodes_to_exclude: list[str] = [],
input_shapes_profile: Sequence[dict[str, str]] | None = None,
**kwargs: Any,
) -> onnx.ModelProto:
"""Quantizes `onnx_model` using the Activation aware quantization a.k.a AWQ algorithm."""
@@ -953,7 +956,7 @@ def _quantize_awq_lite(
logger.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)
session = create_inference_session(augmented_onnx_path, calibration_eps, input_shapes_profile)
inputs = []
for inp_d in data_reader:
inputs.append(inp_d)
@@ -1218,6 +1221,7 @@ def quantize(
block_size: int | None = None,
nodes_to_exclude: list[str] | None = [r"/lm_head"],
log_level: str = "INFO",
input_shapes_profile: Sequence[dict[str, str]] | None = None,
**kwargs: Any,
) -> onnx.ModelProto:
"""Applies INT4 Weight-Only-Quantization (WoQ) to an ONNX model.
@@ -1250,6 +1254,11 @@ def quantize(
.. note::
By default, ``lm-head`` node is NOT quantized.
log_level: The logging level to use (default: logging.INFO)
input_shapes_profile:
The profile of shapes of inputs to the ONNX model - might be needed by some execution providers like
TensorrtExecutionProvider and NvTensorRTRTXExecutionProvider. Used in onnxruntime session creation.
Default value is None.
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].
@@ -1260,7 +1269,6 @@ def quantize(
Default: 0.5.
- **awqclip_bsz_col** (int): Batch size for processing the column dimension in awq-clip.
Default: 1024.
log_level: The logging level to use (default: logging.INFO)
**Returns**: A quantized ONNX model in ONNX ModelProto format.
"""
configure_logging(level=log_level.upper())
@@ -1318,6 +1326,7 @@ def quantize(
nodes_to_exclude=nodes_to_exclude,
use_zero_point=use_zero_point,
enable_weight_clipping=do_weight_clipping,
input_shapes_profile=input_shapes_profile,
**kwargs,
)
elif calibration_method in ["awq_clip", "awq_clip_trt"]:
@@ -1328,6 +1337,7 @@ def quantize(
calibration_eps,
block_size,
nodes_to_exclude=nodes_to_exclude,
input_shapes_profile=input_shapes_profile,
**kwargs,
)
else:
modelopt/onnx/quantization/ort_utils.py
+24 -2
View File
@@ -18,6 +18,7 @@
import glob
import os
import platform
from collections.abc import Sequence
import onnxruntime as ort
from onnxruntime.quantization.operators.qdq_base_operator import QDQOperatorBase
@@ -209,17 +210,38 @@ def update_trt_ep_support(
return trt_plugins
def create_inference_session(onnx_path_or_model: str | bytes, calibration_eps: list[str]):
def create_inference_session(
onnx_path_or_model: str | bytes,
calibration_eps: list[str],
input_shapes_profile: Sequence[dict[str, str]] | None = None,
):
"""Create an ORT InferenceSession."""
logger.info("Creating ORT InferenceSession")
sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
if input_shapes_profile is not None:
# Input-shapes-profile is used by NvTensorRtRtx EP and also usable by TRT EP.
# Input-shapes-profile is passed in provider-options which require that length of
# provider-options equals length of providers.
assert len(input_shapes_profile) == len(calibration_eps), (
"Number of calibration EPs and number of input-shapes-profile don't match"
)
for i in range(len(input_shapes_profile)):
if len(input_shapes_profile[i]) > 0:
logger.debug(
f"Found non-empty input-shapes-profile for calibration-EP: {calibration_eps[i]}"
)
for k, v in input_shapes_profile[i].items():
logger.debug(
f"Input-Shapes-Profile: EP: {calibration_eps[i]}, key: {k}, value: {v}"
)
providers = _prepare_ep_list(calibration_eps)
logger.debug(f"Created session with providers: {providers}")
logger.debug(f"Creating session with providers: {providers}")
return ort.InferenceSession(
onnx_path_or_model,
sess_options=sess_options,
providers=providers,
provider_options=None if input_shapes_profile is None else input_shapes_profile,
)
modelopt/onnx/quantization/quantize.py
+37 -1
View File
@@ -35,6 +35,7 @@ import os
import platform
import shutil
import tempfile
from collections.abc import Sequence
from typing import Any
import onnx
@@ -226,6 +227,7 @@ def quantize(
passes: list[str] = ["concat_elimination"],
simplify: bool = False,
calibrate_per_node: bool = False,
input_shapes_profile: Sequence[dict[str, str]] | None = None,
**kwargs: Any,
) -> None:
"""Quantizes the provided ONNX model.
@@ -246,7 +248,7 @@ def quantize(
Input shapes used for calibration process.
calibration_eps:
Priority order for the execution providers (EP) to calibrate the model.
Any subset of ['trt', 'cuda:x', 'dml:x', 'cpu'], where 'x' is the device id.
Any subset of ['NvTensorRtRtx', 'trt', 'cuda:x', 'dml:x', 'cpu'], where 'x' is the device id.
.. note::
If a custom op is detected in the model, 'trt' will automatically be added to the EP list.
@@ -303,6 +305,39 @@ def quantize(
calibrate_per_node:
Calibrate the model node by node instead of calibrating the entire model. This allowes calibration with
a lower system memory with the cost of longer calibration time.
input_shapes_profile:
This is a sequence of shapes-profile for each EP in calibration_eps. Some EPs like NvTensorRtRtx use these
shapes profile for optimized engine generation for those input shapes. Length of this parameters should
equal length of calibration_eps (i.e. one profile data per EP in calibration_eps, in that order).
A shapes-profile comprises of "min", "max", and "opt" values for the shapes of model inputs
(esp. dynamic shapes). Consider following example snippets for shape-profile data-format of some EPs.
input_shape_profile_for_NvTensorRtrRtx_EP = {
"nv_profile_min_shapes": "input1:dim1xdim2...,input2:dim1xdim2...,...",
"nv_profile_max_shapes": "input1:dim1xdim2...,input2:dim1xdim2...,...",
"nv_profile_opt_shapes": "input1:dim1xdim2...,input2:dim1xdim2...,...",
}
input_shape_profile_for_TensorRT_EP = {
"trt_profile_min_shapes": "input1:dim1xdim2...,input2:dim1xdim2...,...",
"trt_profile_max_shapes": "input1:dim1xdim2...,input2:dim1xdim2...,...",
"trt_profile_opt_shapes": "input1:dim1xdim2...,input2:dim1xdim2...,...",
}
For EPs that don't require such shapes profile (e.g. CPU EP, CUDA EP, DML EP), empty profile {} can be used.
For example, if calibration_eps are ["NvTensorRtRtx", "cpu"], then input_shapes_profile can be set to:
- [input_shapes_profile_for_NvTensorRtRtx_EP, {}]
If None of the calibration_eps require any such shapes profile for model inputs, then nothing needs to be
set for this "input_shapes_profile" parameter.
Default value is None.
kwargs:
Additional keyword arguments for int4 quantization, including:
- awqlite_alpha_step (float): Alpha step for lite, range [0, 1].
@@ -437,6 +472,7 @@ def quantize(
nodes_to_exclude=nodes_to_exclude,
use_zero_point=use_zero_point,
log_level=log_level,
input_shapes_profile=input_shapes_profile,
**kwargs,
)
else:
modelopt/onnx/trt_utils.py
+65 -19
View File
@@ -22,7 +22,12 @@ import onnx
import onnx_graphsurgeon as gs
from modelopt.onnx.logging_config import logger
from modelopt.onnx.utils import get_dynamic_graph_inputs, parse_shapes_spec, save_onnx
from modelopt.onnx.utils import (
get_dynamic_graph_inputs,
get_tensor_by_name,
parse_shapes_spec,
save_onnx,
)
try:
import tensorrt as trt
@@ -108,35 +113,79 @@ def get_custom_layers(
return custom_layers, all_tensor_info
def infer_types_shapes(graph: gs.Graph, all_tensor_info: dict) -> None:
"""Updates tensor shapes in ORT graph.
def infer_types_shapes(model: onnx.ModelProto, all_tensor_info: dict) -> onnx.ModelProto:
"""Updates tensor shapes in ONNX graph.
Args:
graph: ONNX model's GS graph.
model: ONNX model.
all_tensor_info: Dictionary containing tensors information.
Returns:
None. In-memory modification of graph.
onnx.ModelProto: ONNX model with inferred types and shapes.
"""
logger.debug("Inferring types and shapes for graph tensors")
def _map_trt_to_python_type(trt_type: trt.DataType):
def _map_trt_to_onnx_type(trt_type: trt.DataType):
trt_to_onnx_dtype_mapping = {
trt.float32: onnx.TensorProto.FLOAT,
trt.float16: onnx.TensorProto.FLOAT16,
trt.bfloat16: onnx.TensorProto.BFLOAT16,
trt.int4: onnx.TensorProto.INT4,
trt.int8: onnx.TensorProto.INT8,
trt.uint8: onnx.TensorProto.UINT8,
trt.int32: onnx.TensorProto.INT32,
trt.int64: onnx.TensorProto.INT64,
trt.bool: onnx.TensorProto.BOOL,
trt.fp8: onnx.TensorProto.FLOAT8E4M3FN,
trt.fp4: onnx.TensorProto.FLOAT4E2M1,
}
try:
return trt.nptype(trt_type)
return trt_to_onnx_dtype_mapping[trt_type]
except TypeError as e:
logger.warning(f"{e}. TRT datatype: {trt_type}. Setting to None")
return None
updated_tensors = 0
for node in graph.nodes:
for out in node.outputs:
if out.name in all_tensor_info:
out.shape = all_tensor_info[out.name]["shape"]
out.dtype = out.dtype or _map_trt_to_python_type(all_tensor_info[out.name]["dtype"])
updated_tensors += 1
def _create_tensor_shape_proto_from_np_arr(np_arr):
new_shape_proto = onnx.TensorShapeProto()
for dim_val in np_arr:
dim = onnx.TensorShapeProto.Dimension()
setattr(dim, "dim_param" if isinstance(dim_val, str) else "dim_value", dim_val)
new_shape_proto.dim.append(dim)
return new_shape_proto
logger.info(f"Updated {updated_tensors} tensors with type and shape information")
for node in model.graph.node:
for out in node.output:
if out not in all_tensor_info:
continue
tensor = get_tensor_by_name(model, out)
if isinstance(tensor, onnx.ValueInfoProto):
if not tensor.type.tensor_type.elem_type:
tensor.type.tensor_type.elem_type = _map_trt_to_onnx_type(
all_tensor_info[tensor.name]["dtype"]
)
if all_tensor_info[tensor.name]["shape"]:
tensor.type.tensor_type.shape.CopyFrom(
_create_tensor_shape_proto_from_np_arr(
all_tensor_info[tensor.name]["shape"]
)
)
elif tensor is None:
tensor = onnx.helper.make_tensor_value_info(
name=out,
elem_type=_map_trt_to_onnx_type(all_tensor_info[out]["dtype"]),
shape=all_tensor_info[out]["shape"],
)
model.graph.value_info.append(tensor)
logger.info("Updated tensors with type and shape information")
# Topologically sort graph
graph = gs.import_onnx(model)
graph.cleanup().toposort()
model = gs.export_onnx(graph)
return model
def set_trt_plugin_domain(model: onnx.ModelProto, custom_ops: list[str]) -> onnx.ModelProto:
@@ -188,10 +237,7 @@ def infer_types_shapes_tensorrt(
_, all_tensor_info = get_custom_layers(model, trt_plugins, strongly_typed)
# Ensure that all tensors in the graph have type and shape info
graph = gs.import_onnx(model)
infer_types_shapes(graph, all_tensor_info)
model = gs.export_onnx(graph)
return model
return infer_types_shapes(model, all_tensor_info)
def load_onnx_model(
modelopt/onnx/utils.py
+18 -6
View File
@@ -25,7 +25,7 @@ from typing import Any
import numpy as np
import onnx
import onnx_graphsurgeon as gs
from onnx import ValueInfoProto, numpy_helper
from onnx import TensorProto, ValueInfoProto, numpy_helper
from onnx.helper import get_attribute_value
from onnx_graphsurgeon import Constant, Node, Variable
@@ -287,9 +287,16 @@ def _convert_types_to_np(types: dict[str, int] | list[int] | int) -> Any:
return onnx.helper.tensor_dtype_to_np_dtype(types)
def get_tensor_by_name(onnx_model: onnx.ModelProto, tensor_name: str) -> ValueInfoProto | None:
def get_tensor_by_name(
onnx_model: onnx.ModelProto, tensor_name: str
) -> ValueInfoProto | TensorProto | None:
"""This function returns a tensor from its name.
This function searches for a tensor in the model's:
1. Value info (shape/type info, no data)
2. Initializers (TensorProto, contains actual data)
3. Inputs and outputs
Args:
onnx_model: ONNX model.
tensor_name: tensor name.
@@ -297,10 +304,15 @@ def get_tensor_by_name(onnx_model: onnx.ModelProto, tensor_name: str) -> ValueIn
Returns:
tensor
"""
for tensor in onnx_model.graph.value_info:
if tensor.name == tensor_name:
return tensor
return None
tensor_val = next(
(tens for tens in onnx_model.graph.value_info if tens.name == tensor_name), None
)
tensor_init = next(
(tens for tens in onnx_model.graph.initializer if tens.name == tensor_name), None
)
tensor_inp = next((tens for tens in onnx_model.graph.input if tens.name == tensor_name), None)
tensor_out = next((tens for tens in onnx_model.graph.output if tens.name == tensor_name), None)
return tensor_val or tensor_init or tensor_inp or tensor_out
def gen_random_inputs(
modelopt/torch/__init__.py
+2 -2
View File
@@ -22,9 +22,9 @@ from torch import __version__ as _torch_version
from . import distill, nas, opt, prune, quantization, sparsity, speculative, utils
if _Version(_torch_version) < _Version("2.6"):
if _Version(_torch_version) < _Version("2.7"):
_warnings.warn(
"nvidia-modelopt will drop torch<2.6 support in a future release.", DeprecationWarning
"nvidia-modelopt will drop torch<2.7 support in a future release.", DeprecationWarning
)
# Since `hf` dependencies are optional and users have pre-installed transformers, we need to ensure
modelopt/torch/_deploy/utils/torch_onnx.py
+14 -23
View File
@@ -410,28 +410,19 @@ def get_onnx_bytes_and_metadata(
else nullcontext()
)
with torch.inference_mode(), autocast, quantizer_context:
if not dynamo_export or Version(torch.__version__) >= Version("2.6"):
additional_kwargs = {}
if not dynamo_export and Version(torch.__version__) >= Version("2.8"):
additional_kwargs["dynamic_axes"] = dynamic_axes
torch.onnx.export(
model,
dummy_input,
onnx_save_path,
input_names=input_names,
output_names=output_names,
opset_version=onnx_opset,
dynamo=dynamo_export,
**additional_kwargs,
)
else: # torch < 2.6 with dynamo export
export_options = torch.onnx.ExportOptions(dynamic_shapes=True)
dummy_input_args, dummy_input_kwargs = split_args_kwargs(dummy_input)
if dummy_input_kwargs is None:
dummy_input_kwargs = {}
torch.onnx.dynamo_export(
model, *dummy_input_args, export_options=export_options, **dummy_input_kwargs
).save(onnx_save_path)
additional_kwargs = {}
if not dynamo_export and Version(torch.__version__) >= Version("2.8"):
additional_kwargs["dynamic_axes"] = dynamic_axes
torch.onnx.export(
model,
dummy_input,
onnx_save_path,
input_names=input_names,
output_names=output_names,
opset_version=onnx_opset,
dynamo=dynamo_export,
**additional_kwargs,
)
# Check that export worked
assert len(os.listdir(onnx_path)) > 0, "Torch to onnx export failed."
@@ -463,7 +454,7 @@ def get_onnx_bytes_and_metadata(
onnx_opt_graph = qdq_to_dq(onnx_opt_graph)
if weights_dtype == "float16":
if is_fp4_quantized(model) or is_mxfp8_quantized(model):
if not use_autocast:
onnx_opt_graph = convert_float_to_float16(
onnx_opt_graph, keep_io_types=False, disable_shape_infer=True
)
modelopt/torch/export/layer_utils.py
+27 -24
View File
@@ -85,19 +85,26 @@ with import_plugin("megatron", verbose=False):
def get_experts_list(module: torch.nn.Module, model_type: str):
"""Returns list of grouped experts by linear name for given module."""
experts_list = []
# Define linear layer names for different model types
if "mixtralforcausallm" in model_type:
experts_list.extend(
[
[
_get_mixtral_expert(module.experts, i, linear_name)
for i in range(len(module.experts))
]
for linear_name in ["w1", "w2", "w3"]
]
)
linear_names = ["w1", "w2", "w3"]
elif any(
qwen_variant in model_type
for qwen_variant in ["qwenmoeforcausallm", "qwen2moeforcausallm", "qwen3moeforcausallm"]
):
linear_names = ["gate_proj", "down_proj", "up_proj"]
else:
raise NotImplementedError(f" {model_type} not supported")
# Common logic for all supported model types
experts_list.extend(
[
[_get_expert_attr(module.experts, i, linear_name) for i in range(len(module.experts))]
for linear_name in linear_names
]
)
return experts_list
@@ -892,17 +899,13 @@ def build_mlp_config(
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
# ...
def _get_expert_attr(experts: nn.Module, export_id: int, linear_name: str):
# Generic expert attribute accessor.
# Works for most MoE models that store experts as a list/ModuleList where
# each expert has linear layers as direct attributes:
# experts[0].w1, experts[0].w2, experts[0].w3 (Mixtral)
# experts[0].gate_proj, experts[0].down_proj, experts[0].up_proj (Qwen)
# experts[0].linear_fc1, experts[0].linear_fc2 (Llama MCore)
return getattr(experts[export_id], linear_name)
@@ -1205,7 +1208,7 @@ def build_moe_config(module: nn.Module, decoder_type) -> MOEConfig:
module.experts.local_experts,
["linear_fc1", "linear_fc2"],
len(module.experts.local_experts),
_get_mixtral_expert,
_get_expert_attr,
)
# For Mcore model, experts.fc.weight needs to be flipped along axis = 1
mid_point = experts.fc.weight.shape[1] // 2
@@ -1222,7 +1225,7 @@ def build_moe_config(module: nn.Module, decoder_type) -> MOEConfig:
module.experts,
["w1", "w2", "w3"],
len(module.experts),
_get_mixtral_expert,
_get_expert_attr,
)
elif decoder_type == "dbrx":
experts.fc, experts.proj = build_stacked_experts(
@@ -1236,7 +1239,7 @@ def build_moe_config(module: nn.Module, decoder_type) -> MOEConfig:
module.experts,
["gate_proj", "down_proj", "up_proj"],
len(module.experts),
_get_mixtral_expert,
_get_expert_attr,
)
else:
raise NotImplementedError(f"{decoder_type} not supported")
@@ -1698,7 +1701,7 @@ def update_experts_avg_prequant_scale(experts: nn.Module):
"""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
get_func = _get_expert_attr
num_experts = len(experts.experts)
experts = experts.experts
elif "dbrx" in type(experts).__name__.lower():
modelopt/torch/export/model_config.py
+1
View File
@@ -32,6 +32,7 @@ QUANTIZATION_INT8_SQ = "int8_sq"
QUANTIZATION_INT4_AWQ = "int4_awq"
QUANTIZATION_W4A8_AWQ = "w4a8_awq"
QUANTIZATION_NVFP4 = "nvfp4"
QUANTIZATION_MXFP4 = "mxfp4"
QUANTIZATION_W4A8_MXFP4_FP8 = "w4a8_mxfp4_fp8"
QUANTIZATION_NVFP4_AWQ = "nvfp4_awq"
QUANTIZATION_FP8_PB_REAL = "fp8_pb_real"
modelopt/torch/export/model_config_export.py
+1 -4
View File
@@ -545,11 +545,8 @@ def export_tensorrt_llm_checkpoint(
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."
" can be saved with torch.save for further inspection."
)
raise e
modelopt/torch/export/plugins/mcore_common.py
+8 -1
View File
@@ -31,7 +31,12 @@ from .mcore_nemotron import (
nemotron_h_causal_lm_export,
nemotron_h_causal_lm_import,
)
from .mcore_qwen import qwen3_causal_lm_export, qwen3_causal_lm_import
from .mcore_qwen import (
qwen3_causal_lm_export,
qwen3_causal_lm_import,
qwen25_causal_lm_export,
qwen25_causal_lm_import,
)
all_mcore_hf_export_mapping: dict[str, Any] = {
"DeepseekV2ForCausalLM": deepseek_causal_lm_export,
@@ -44,6 +49,7 @@ all_mcore_hf_export_mapping: dict[str, Any] = {
"LlamaForCausalLMEagle3": eagle3_llama_causal_lm_export,
"Qwen3ForCausalLM": qwen3_causal_lm_export,
"Qwen3MoeForCausalLM": qwen3_causal_lm_export,
"Qwen2ForCausalLM": qwen25_causal_lm_export,
}
all_mcore_hf_import_mapping: dict[str, Any] = {
@@ -54,4 +60,5 @@ all_mcore_hf_import_mapping: dict[str, Any] = {
"NemotronHForCausalLM": nemotron_h_causal_lm_import,
"Qwen3ForCausalLM": qwen3_causal_lm_import,
"Qwen3MoeForCausalLM": qwen3_causal_lm_import,
"Qwen2ForCausalLM": qwen25_causal_lm_import,
}
modelopt/torch/export/plugins/mcore_custom.py
+7 -1
View File
@@ -279,7 +279,10 @@ def _get_safetensor_slices(
tensor_slice = f.get_slice(key)
assert tensor_slice is not None
shape = tensor_slice.get_shape()
assert len(shape) in (2, 3), f"Shape {shape} is not supported!"
assert len(shape) in (1, 2, 3), f"Shape {shape} is not supported!"
# 1 for bias case
# 3 for packed MoE case
# MCore tensor parallel model sharding
sharding_dim = parallel_config.sharding_dim
parallel_group = parallel_config.parallel_group
@@ -339,6 +342,9 @@ def _get_safetensor_slices(
raise ValueError(
f"Unsupported sharding_dim: {sharding_dim} for shape: {shape}"
)
elif len(shape) == 1:
# For bias case
tensor = tensor_slice[rank_offset : rank_offset + per_rank_size]
else:
raise ValueError(f"Unsupported shape: {shape}")
return tensor
modelopt/torch/export/plugins/mcore_qwen.py
+28
View File
@@ -69,3 +69,31 @@ qwen3_causal_lm_export: dict[str, CustomModuleMapping] = {
"local_experts.linear_fc1": GatedMLPSlicing("model.layers.{}.mlp.experts.{}."),
"local_experts.linear_fc2": NameRemapping("model.layers.{}.mlp.experts.{}.down_proj."),
}
qwen25_causal_lm_import: dict[str, CustomModuleMapping] = {
"word_embeddings": NameRemapping("model.embed_tokens.", COL_TP),
"final_layernorm": NameRemapping("model.norm.", REPLICATE),
"output_layer": NameRemapping("lm_head.", COL_TP),
# Attention
"input_layernorm": NameRemapping("model.layers.{}.input_layernorm.", REPLICATE),
"linear_qkv": QKVMerging("model.layers.{}.self_attn.", COL_TP),
"linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj.", ROW_TP),
# MLP
"pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm.", REPLICATE),
"linear_fc1": GatedMLPMerging("model.layers.{}.mlp.", COL_TP),
"linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj.", ROW_TP),
}
qwen25_causal_lm_export: dict[str, CustomModuleMapping] = {
"word_embeddings": NameRemapping("model.embed_tokens."),
"final_layernorm": NameRemapping("model.norm."),
"output_layer": NameRemapping("lm_head."),
# Attention
"input_layernorm": NameRemapping("model.layers.{}.input_layernorm."),
"linear_qkv": QKVSlicing("model.layers.{}.self_attn."),
"linear_proj": NameRemapping("model.layers.{}.self_attn.o_proj."),
# MLP
"pre_mlp_layernorm": NameRemapping("model.layers.{}.post_attention_layernorm."),
"linear_fc1": GatedMLPSlicing("model.layers.{}.mlp."),
"linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj."),
}
modelopt/torch/export/plugins/megatron_importer.py
+28 -1
View File
@@ -313,6 +313,34 @@ class GPTModelImporter:
state_dict["weight"] = tensor.reshape(-1, hidden_size)
# Handle bias merging
bias = module.state_dict().get("bias", None)
if bias is not None:
q_bias = self._get_safetensor(
prefix + q_proj_name + ".bias", parallel_config=parallel_config
)
k_bias = self._get_safetensor(
prefix + k_proj_name + ".bias", parallel_config=parallel_config
)
v_bias = self._get_safetensor(
prefix + v_proj_name + ".bias", parallel_config=parallel_config
)
# Reshape separate biases to match the head structure
q_bias = q_bias.reshape(-1, head_size)
k_bias = k_bias.reshape(-1, head_size)
v_bias = v_bias.reshape(-1, head_size)
# Create target bias tensor with the same structure as the fused QKV
bias_tensor = bias.detach().clone().reshape([qkv_total_dim, head_size])
# Merge biases using the same slicing logic as weights
bias_tensor[q_slice] = q_bias.to(dtype=bias_tensor.dtype).to(device=bias_tensor.device)
bias_tensor[k_slice] = k_bias.to(dtype=bias_tensor.dtype).to(device=bias_tensor.device)
bias_tensor[v_slice] = v_bias.to(dtype=bias_tensor.dtype).to(device=bias_tensor.device)
state_dict["bias"] = bias_tensor.reshape(-1)
module.load_state_dict(state_dict)
def _unpack_name_remapping(
@@ -469,7 +497,6 @@ class GPTModelImporter:
# Output layer
if hasattr(model, "output_layer") and not model.share_embeddings_and_output_weights:
self.rules["output_layer"](model.output_layer)
# MTP
if hasattr(model, "mtp"):
# MTP is the last layer in DeepSeek V3/R1
modelopt/torch/export/quant_utils.py
+124 -232
View File
@@ -32,7 +32,11 @@ from modelopt.torch.quantization.qtensor import (
NVFP4QTensor,
QTensorWrapper,
)
from modelopt.torch.quantization.utils import is_quantized_linear
from modelopt.torch.quantization.utils import (
QuantizerAttrNames,
quantizer_attr_names,
weight_attr_names,
)
from ..quantization.nn import SequentialQuantizer, TensorQuantizer
from .model_config import (
@@ -46,6 +50,7 @@ from .model_config import (
QUANTIZATION_FP8_PC_PT,
QUANTIZATION_INT4_AWQ,
QUANTIZATION_INT8_SQ,
QUANTIZATION_MXFP4,
QUANTIZATION_NONE,
QUANTIZATION_NVFP4,
QUANTIZATION_NVFP4_AWQ,
@@ -190,68 +195,82 @@ def get_scaling_factor(quantizer: TensorQuantizer) -> torch.Tensor:
return scaling_factor
def get_activation_scaling_factor(module: nn.Module) -> torch.Tensor:
def get_activation_scaling_factor(
module: nn.Module, input_quantizer_name: str = "input_quantizer"
) -> torch.Tensor:
"""Returns the activation scaling factor."""
# If NVFP4, return activation scaling factor from NVFP4QTensor
input_quantizer = getattr(module, input_quantizer_name, None)
if input_quantizer is None:
return None
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
)
]:
return NVFP4QTensor.get_activation_scaling_factor(input_quantizer)
return get_scaling_factor(input_quantizer)
def get_weight_scaling_factor(module: nn.Module) -> torch.Tensor:
def get_weight_scaling_factor(module: nn.Module, weight_name: str = "weight") -> 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])
weight: nn.Parameter = getattr(module, weight_name)
weight_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr(
module, quantizer_attr_names(weight_name).weight_quantizer, None
)
if weight_quantizer is None:
return None
if isinstance(weight_quantizer, SequentialQuantizer):
return get_scaling_factor(weight_quantizer[0])
quantization_format = get_quantization_format(module)
# If NVFP4, we need to return quantized per_block scaling factors
if get_quantization_format(module) in [
if quantization_format 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).to(
module.weight.device
weight,
weight_quantizer.block_sizes[-1],
NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer).to(
weight.device
),
)[0]
if get_quantization_format(module) == QUANTIZATION_W4A8_MXFP4_FP8:
return MXFP4QTensor.quantize(
module.weight, block_size=module.weight_quantizer.block_sizes[-1]
)[1].reshape(*module.weight.shape[:-1], -1)
return (
get_scaling_factor(module.weight_quantizer) if hasattr(module, "weight_quantizer") else None
)
if quantization_format in [QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_MXFP4]:
return MXFP4QTensor.quantize(weight, block_size=weight_quantizer.block_sizes[-1])[
1
].reshape(*weight.shape[:-1], -1)
return get_scaling_factor(weight_quantizer)
def get_weight_scaling_factor_2(module: nn.Module) -> torch.Tensor:
def get_weight_scaling_factor_2(module: nn.Module, weight_name: str = "weight") -> torch.Tensor:
"""Returns the secondary weight scaling factor."""
weight_quantizer = getattr(module, quantizer_attr_names(weight_name).weight_quantizer, None)
if weight_quantizer is None:
return None
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 NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer)
# SequentialQuantizer is required
if not isinstance(weight_quantizer, SequentialQuantizer) or not weight_quantizer[-1].is_enabled:
return None
assert len(module.weight_quantizer) == 2, (
assert len(weight_quantizer) == 2, (
"modelopt only supports 2 sequential quantization layers for now"
)
return get_scaling_factor(module.weight_quantizer[-1])
return get_scaling_factor(weight_quantizer[-1])
def get_prequant_scaling_factor(module: nn.Module) -> torch.Tensor:
@@ -343,12 +362,12 @@ def get_kv_cache_dtype(modules: list[nn.Module] | nn.Module) -> str | None:
return QUANTIZATION_NONE
def get_weight_block_size(module: nn.Module) -> int:
def get_weight_block_size(module: nn.Module, weight_name: str = "weight") -> int:
"""Returns the weight block size."""
if not hasattr(module, "weight_quantizer"):
return 0
weight_quantizer = getattr(module, quantizer_attr_names(weight_name).weight_quantizer, None)
weight_quantizer = module.weight_quantizer
if weight_quantizer is None:
return 0
if isinstance(weight_quantizer, SequentialQuantizer):
weight_quantizer = weight_quantizer[0]
@@ -370,86 +389,89 @@ def get_quantization_format(module) -> str | None:
The first non-None quantization string is returned.
"""
def _get_quantization_from_linear_layer(layer):
if not hasattr(layer, "weight_quantizer") or not layer.weight_quantizer.is_enabled:
def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames):
weight_quantizer = getattr(layer, quantizer_attr_names.weight_quantizer, None)
input_quantizer = getattr(layer, quantizer_attr_names.input_quantizer, None)
if weight_quantizer is None or not weight_quantizer.is_enabled:
return QUANTIZATION_NONE
w_quantizer = layer.weight_quantizer
# Handle SequentialQuantizer
if isinstance(w_quantizer, SequentialQuantizer):
if isinstance(weight_quantizer, SequentialQuantizer):
assert (
len(w_quantizer) == 2
and w_quantizer[0].num_bits == 4
and w_quantizer[1].num_bits == (4, 3)
len(weight_quantizer) == 2
and weight_quantizer[0].num_bits == 4
and weight_quantizer[1].num_bits == (4, 3)
), "Unsupported SequentialQuantizer configuration"
assert (
w_quantizer[0].block_sizes
and len(w_quantizer[0].block_sizes) > 0
and w_quantizer[0].block_sizes[-1] > 0
weight_quantizer[0].block_sizes
and len(weight_quantizer[0].block_sizes) > 0
and weight_quantizer[0].block_sizes[-1] > 0
), "Invalid block_sizes for SequentialQuantizer"
return QUANTIZATION_W4A8_AWQ
# Handle individual num_bits cases
if w_quantizer.num_bits == 4:
assert len(w_quantizer.block_sizes) > 0 and w_quantizer.block_sizes[-1] > 0, (
if weight_quantizer.num_bits == 4:
assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, (
"Invalid block_sizes for INT4 quantizer"
)
return QUANTIZATION_INT4_AWQ
if w_quantizer.num_bits == 8:
if weight_quantizer.num_bits == 8:
return QUANTIZATION_INT8_SQ
if w_quantizer.num_bits == (4, 3):
if w_quantizer.block_sizes:
assert w_quantizer.block_sizes[-1] > 0, "Invalid block_sizes for FP8 quantizer"
if w_quantizer.fake_quant:
if weight_quantizer.num_bits == (4, 3):
if weight_quantizer.block_sizes:
assert weight_quantizer.block_sizes[-1] > 0, "Invalid block_sizes for FP8 quantizer"
if weight_quantizer.fake_quant:
return QUANTIZATION_FP8_PB_WO
else:
return QUANTIZATION_FP8_PB_REAL
if w_quantizer.axis == 0:
if weight_quantizer.axis == 0:
return QUANTIZATION_FP8_PC_PT
return QUANTIZATION_FP8
if w_quantizer.num_bits == (2, 1):
if hasattr(layer, "input_quantizer") and hasattr(
layer.input_quantizer, "_pre_quant_scale"
):
if weight_quantizer.num_bits == (2, 1):
# FP4 formats are all block quantization
block_sizes = getattr(weight_quantizer, "block_sizes")
scale_bits = block_sizes.get("scale_bits")
if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"):
return QUANTIZATION_NVFP4_AWQ
if getattr(layer, "fused_with_layernorm", False):
return QUANTIZATION_NVFP4_AWQ
block_sizes = getattr(layer.weight_quantizer, "block_sizes", None)
scale_bits = block_sizes.get("scale_bits", None) if block_sizes else None
assert input_quantizer is not None, (
f"input_quantizer is None for {quantizer_attr_names}"
)
if (
layer.weight_quantizer.is_enabled
and block_sizes
and block_sizes.get("type", "static") == "dynamic"
and scale_bits
block_sizes.get("type", "static") == "dynamic"
and scale_bits == (8, 0)
and layer.input_quantizer.is_enabled
and layer.input_quantizer.num_bits == (4, 3)
and layer.input_quantizer.block_sizes is None
and input_quantizer.is_enabled
and input_quantizer.num_bits == (4, 3)
and input_quantizer.block_sizes is None
):
return QUANTIZATION_W4A8_MXFP4_FP8
return QUANTIZATION_NVFP4
if scale_bits == (4, 3):
return QUANTIZATION_NVFP4
elif scale_bits == (8, 0):
return QUANTIZATION_MXFP4
# Raise error for unsupported num_bits
raise NotImplementedError(f"Unsupported quantizer with num_bits: {w_quantizer.num_bits}")
raise NotImplementedError(
f"Unsupported quantizer with num_bits: {weight_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
for weight_name in weight_attr_names(module):
quantization = _get_quantization_from_layer(module, quantizer_attr_names(weight_name))
if quantization != QUANTIZATION_NONE:
return quantization
for _, layer in module.named_children():
format = get_quantization_format(layer)
if format != QUANTIZATION_NONE:
return format
return QUANTIZATION_NONE
@@ -703,7 +725,7 @@ def to_quantized_weight(
else weights_scaling_factor2,
)[0]._quantized_data
if quantization == QUANTIZATION_W4A8_MXFP4_FP8:
if quantization in [QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_MXFP4]:
return MXFP4QTensor.quantize(weight, block_size=block_size)[0]._quantized_data
raise NotImplementedError(f"quantization format {quantization} not supported")
@@ -806,6 +828,24 @@ def postprocess_state_dict(state_dict: dict, maxbound: float, quantization: str
):
keys_to_delete.append(key)
# Check for tied weights and remove duplicates
seen_tensors = {}
# Remove any tied weights if found.
for key, value in post_state_dict.items():
if isinstance(value, torch.Tensor):
# Use tensor data pointer to identify tied weights
tensor_id = value.data_ptr()
if tensor_id in seen_tensors:
# This is a tied weight, mark for deletion and warn
keys_to_delete.append(key)
logger.warning(
f"Found tied weight: '{key}' is tied to '{seen_tensors[tensor_id]}'. "
f"Removing duplicate '{key}' from the exported state dict."
)
else:
seen_tensors[tensor_id] = key
for key in keys_to_delete:
del post_state_dict[key]
@@ -966,151 +1006,3 @@ def get_quant_config(named_modules: nn.Module | dict[str, nn.Module]) -> dict[st
quant_config["quantization"]["kv_cache_quant_algo"] = kv_cache_format
return quant_config
def quantize_llama4_experts_for_hf_export(module: nn.Module):
"""Quantize the experts in the Llama4 model."""
from transformers.models.llama4.modeling_llama4 import Llama4TextExperts
assert isinstance(module, Llama4TextExperts), "Module is not a Llama4TextExperts"
assert module.gate_up_proj_weight_quantizer.is_enabled
assert module.down_proj_weight_quantizer.is_enabled
assert module.gate_up_proj_input_quantizer.is_enabled
assert module.down_proj_input_quantizer.is_enabled
# Handle uncalibrated input quantizers that have None amax values
input_quantizers = [
module.gate_up_proj_input_quantizer,
module.down_proj_input_quantizer,
]
# Only handle amax for non-dynamic quantizers
non_dynamic_quantizers = [q for q in input_quantizers if not getattr(q, "_dynamic", False)]
if non_dynamic_quantizers:
# Find the maximum amax value from non-None quantizers
valid_amax_values = [
quantizer.amax for quantizer in non_dynamic_quantizers if quantizer.amax is not None
]
device = module.gate_up_proj.device
# If all quantizers have None amax, set a default value
if not valid_amax_values:
default_amax = torch.tensor(1.0, dtype=torch.float32, device=device)
warn(
"All input quantizers have None amax values. Setting default amax to 1.0. "
"This typically occurs when experts are not activated during calibration. "
"Consider increasing your calibration dataset size to ensure all experts are exercised."
)
for quantizer in non_dynamic_quantizers:
if quantizer.amax is None:
quantizer.amax = default_amax.clone()
else:
# Set None amax values to the maximum of existing values
max_amax = torch.max(torch.stack(valid_amax_values))
if max_amax.device != device:
max_amax = max_amax.to(device)
for quantizer in non_dynamic_quantizers:
if quantizer.amax is None:
warn(
f"Missing amax value for input quantizer. Setting it to {max_amax.item()} for export. "
"This typically occurs when certain experts are not activated during calibration. "
"Consider increasing your calibration dataset size to ensure all experts are exercised."
)
quantizer.amax = max_amax.clone()
for weight_name in ["gate_up_proj", "down_proj"]:
weight = getattr(module, weight_name)
weight_quantizer = getattr(module, f"{weight_name}_weight_quantizer")
if weight_quantizer.num_bits == (4, 3):
assert not weight_quantizer.block_sizes
weight_scale = weight_quantizer.amax.to(torch.float32) / weight_quantizer.maxbound
module.register_buffer(
f"{weight_name}_weight_scale",
weight_scale,
)
setattr(
module,
weight_name,
nn.Parameter(
(weight / weight_scale.to(weight.dtype).to(weight.device)).to(
torch.float8_e4m3fn
),
requires_grad=False,
),
)
elif weight_quantizer.num_bits == (2, 1):
# Maverick export can go OOM on the GPU. So just move to the CPU for weights compression.
weight = weight.to("cpu")
weight_scale_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(
weight_quantizer
).to("cpu")
module.register_buffer(
f"{weight_name}_weight_scale_2",
weight_scale_2,
)
block_size = weight_quantizer.block_sizes[-1]
# For bmm, the weight shape is (num_experts, input_dim, output_dim), so let's first transpose
# the weight to (num_experts, output_dim, input_dim) before calculating scaling factor and quantization.
weight = weight.transpose(-2, -1)
weight_scale = NVFP4QTensor.get_weights_scaling_factor(
weight,
block_size=block_size,
weights_scaling_factor_2=weight_scale_2,
)[0]
quantized_weights = to_quantized_weight(
weight,
weight_scale,
quantization=QUANTIZATION_NVFP4,
weights_scaling_factor2=weight_scale_2,
block_size=block_size,
)
# After quantization, we transpose the weight and scales back to the original order.
quantized_weights = quantized_weights.transpose(-2, -1)
weight_scale = weight_scale.transpose(-2, -1)
module.register_buffer(
f"{weight_name}_weight_scale",
weight_scale,
)
setattr(
module,
weight_name,
nn.Parameter(quantized_weights, requires_grad=False),
)
for input_name in ["gate_up_proj", "down_proj"]:
input_quantizer = getattr(module, f"{input_name}_input_quantizer")
# Skip processing for dynamic quantization since it doesn't have fixed amax
if getattr(input_quantizer, "_dynamic", False):
continue
if input_quantizer.num_bits == (4, 3):
assert not input_quantizer.block_sizes
input_scale = input_quantizer.amax.to(torch.float32) / input_quantizer.maxbound
module.register_buffer(
f"{input_name}_input_scale",
input_scale,
)
elif input_quantizer.num_bits == (2, 1):
input_scale_2 = NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(
input_quantizer
)
module.register_buffer(
f"{input_name}_input_scale",
input_scale_2,
)
modelopt/torch/export/unified_export_hf.py
+200 -101
View File
@@ -17,6 +17,7 @@
import collections.abc
import json
import re
import tempfile
import warnings
from collections import defaultdict
@@ -27,7 +28,8 @@ import torch
import torch.nn as nn
from modelopt.torch.quantization import set_quantizer_by_cfg_context
from modelopt.torch.quantization.nn import SequentialQuantizer
from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
from modelopt.torch.quantization.utils import quantizer_attr_names
from .convert_hf_config import convert_hf_quant_config_format
from .layer_utils import (
@@ -59,7 +61,6 @@ from .quant_utils import (
get_weight_scaling_factor_2,
postprocess_state_dict,
preprocess_linear_fusion,
quantize_llama4_experts_for_hf_export,
to_quantized_weight,
)
@@ -97,7 +98,12 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module):
handles = []
model_type = type(model).__name__.lower()
fused_linears = {}
module_names = set()
for name, module in model.named_modules():
module_names.add(name)
# 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)
@@ -151,6 +157,7 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module):
]:
# Fuse modules that have the same input
preprocess_linear_fusion(modules)
fused_linears[modules[0].name] = [module.name for module in modules]
# Fuse layernorms
if (
@@ -161,6 +168,185 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module):
# Pre quant scale of modules is already updated to avg_pre_quant_scale
fuse_prequant_layernorm(output_to_layernorm[tensor], modules)
# The dummy forward may not be able to activate all the experts.
# Process experts by naming rules like experts.0, experts.1, etc.
for name, modules_fused in fused_linears.items():
if re.search(r"experts?\.\d+", name):
expert_id = 0
while True:
new_expert_name = re.sub(r"(experts?\.)\d+", rf"\g<1>{expert_id}", name, count=1)
if new_expert_name in fused_linears:
expert_id += 1
continue
if new_expert_name not in module_names:
break
new_expert_modules = []
for name_fused in modules_fused:
new_expert_name = re.sub(r"(experts?\.)\d+", rf"\g<1>{expert_id}", name_fused)
assert new_expert_name in module_names
new_expert_modules.append(model.get_submodule(new_expert_name))
preprocess_linear_fusion(new_expert_modules)
expert_id += 1
def _export_quantized_weight(
sub_module: nn.Module, dtype: torch.dtype, weight_name: str = "weight"
):
"""For the given weight attr of the sub_module, export the quantization info of it.
The export includes converting weight tensor to correct quantized values and quantized dtype,
and registering scaling factors.
"""
quantization_format = get_quantization_format(sub_module)
if quantization_format == QUANTIZATION_NONE:
return
block_size = get_weight_block_size(sub_module, weight_name)
quantizer_attrs = quantizer_attr_names(weight_name)
weight: nn.Parameter = getattr(sub_module, weight_name)
weight_quantizer: TensorQuantizer | SequentialQuantizer = getattr(
sub_module, quantizer_attrs.weight_quantizer
)
input_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr(
sub_module, quantizer_attrs.input_quantizer, None
)
output_quantizer: TensorQuantizer | SequentialQuantizer | None = getattr(
sub_module, quantizer_attrs.output_quantizer, None
)
if quantization_format == QUANTIZATION_FP8:
# Convert amax to float32
weight_quantizer._amax = weight_quantizer._amax.to(torch.float32)
if weight_quantizer._amax.dim() == 1:
# Per-tensor amax
weight_scaling_factor = torch.tensor(
weight_quantizer.amax.item() / weight_quantizer.maxbound
)
else:
# Per-channel amax
weight_scaling_factor = torch.tensor(weight_quantizer.amax / weight_quantizer.maxbound)
sub_module.register_buffer(
quantizer_attrs.weight_scale,
weight_scaling_factor,
)
if hasattr(input_quantizer, "_amax"):
assert input_quantizer is not None
input_quantizer._amax = input_quantizer._amax.to(torch.float32)
sub_module.register_buffer(
quantizer_attrs.input_scale,
get_activation_scaling_factor(
sub_module, input_quantizer_name=quantizer_attrs.input_quantizer
).squeeze(),
)
if hasattr(output_quantizer, "_amax"):
assert output_quantizer is not None
output_quantizer._amax = output_quantizer._amax.to(torch.float32)
else:
# Register weight_scale and input_scale
if quantization_format == QUANTIZATION_FP8_PB_REAL:
sub_module.register_buffer(
quantizer_attrs.weight_scale,
weight_quantizer._scale.to(torch.float32),
)
del weight_quantizer._scale
else:
sub_module.register_buffer(
quantizer_attrs.weight_scale, get_weight_scaling_factor(sub_module, weight_name)
)
if (
input_quantizer is not None
and "disabled" not in repr(input_quantizer)
and input_quantizer.amax is not None
):
sub_module.register_buffer(
quantizer_attrs.input_scale,
get_activation_scaling_factor(
sub_module, input_quantizer_name=quantizer_attrs.input_quantizer
).squeeze(),
)
if quantization_format in [
QUANTIZATION_NVFP4_AWQ,
QUANTIZATION_NVFP4,
QUANTIZATION_W4A8_AWQ,
]:
# Register weight_scale_2
sub_module.register_buffer(
quantizer_attrs.weight_scale_2,
get_weight_scaling_factor_2(sub_module, weight_name).squeeze(),
)
weight_scale: torch.Tensor | None = getattr(sub_module, quantizer_attrs.weight_scale, None)
weight_scale_2: torch.Tensor | None = getattr(sub_module, quantizer_attrs.weight_scale_2, None)
quantized_weight = to_quantized_weight(
weight.to(dtype),
weight_scale,
quantization_format,
weight_scale_2,
block_size,
)
setattr(sub_module, weight_name, nn.Parameter(quantized_weight, requires_grad=False))
def _handle_llama4_experts_amax(module: nn.Module):
"""Handle the amax values for the experts in the Llama4 model."""
# Handle uncalibrated input quantizers that have None amax values
input_quantizers = [
module.gate_up_proj_input_quantizer,
module.down_proj_input_quantizer,
]
# Only handle enabled input quantizers
enabled_input_quantizers = [q for q in input_quantizers if q.is_enabled]
# Only handle amax for non-dynamic quantizers
non_dynamic_quantizers = [
q for q in enabled_input_quantizers if not getattr(q, "_dynamic", False)
]
if non_dynamic_quantizers:
# Find the maximum amax value from non-None quantizers
valid_amax_values = [
quantizer.amax for quantizer in non_dynamic_quantizers if quantizer.amax is not None
]
device = module.gate_up_proj.device
# If all quantizers have None amax, set a default value
if not valid_amax_values:
default_amax = torch.tensor(1.0, dtype=torch.float32, device=device)
warnings.warn(
"All input quantizers have None amax values. Setting default amax to 1.0. "
"This typically occurs when experts are not activated during calibration. "
"Consider increasing your calibration dataset size to ensure all experts are exercised."
)
for quantizer in non_dynamic_quantizers:
if quantizer.amax is None:
quantizer.amax = default_amax.clone()
else:
# Set None amax values to the maximum of existing values
max_amax = torch.max(torch.stack(valid_amax_values))
if max_amax.device != device:
max_amax = max_amax.to(device)
for quantizer in non_dynamic_quantizers:
if quantizer.amax is None:
warnings.warn(
f"Missing amax value for input quantizer. Setting it to {max_amax.item()} for export. "
"This typically occurs when certain experts are not activated during calibration. "
"Consider increasing your calibration dataset size to ensure all experts are exercised."
)
quantizer.amax = max_amax.clone()
def _export_hf_checkpoint(
model: nn.Module, dtype: torch.dtype | None = None
@@ -268,98 +454,14 @@ def _export_hf_checkpoint(
has_quantized_layers = False
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)
# Track if any layer is quantized
if quantization_format != QUANTIZATION_NONE:
has_quantized_layers = True
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).squeeze(),
)
if quantization_format not in [QUANTIZATION_FP8, QUANTIZATION_NONE]:
# Register weight_scale and input_scale
if quantization_format == QUANTIZATION_FP8_PB_REAL:
sub_module.register_buffer(
"weight_scale",
sub_module.weight_quantizer._scale.to(torch.float32),
)
del sub_module.weight_quantizer._scale
else:
sub_module.register_buffer(
"weight_scale", get_weight_scaling_factor(sub_module)
)
# Remove size-1 dimensions for blocked fp8 scales
sub_module.weight_scale.squeeze()
if (
hasattr(sub_module, "input_quantizer")
and "disabled" not in repr(sub_module.input_quantizer)
and sub_module.input_quantizer.amax is not None
):
sub_module.register_buffer(
"input_scale", get_activation_scaling_factor(sub_module).squeeze()
)
# 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,
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)
elif "Llama4TextExperts" in type(sub_module).__name__:
quantize_llama4_experts_for_hf_export(sub_module)
if get_quantization_format(sub_module) != QUANTIZATION_NONE:
has_quantized_layers = True
if is_quantlinear(sub_module):
_export_quantized_weight(sub_module, dtype)
elif "Llama4TextExperts" in type(sub_module).__name__:
_handle_llama4_experts_amax(sub_module)
for weight_name in ["gate_up_proj", "down_proj"]:
_export_quantized_weight(sub_module, dtype, weight_name)
quantized_state_dict = model.state_dict()
@@ -400,9 +502,9 @@ def export_hf_checkpoint(
hf_quant_config = convert_hf_quant_config_format(hf_quant_config)
# Save model
if not save_modelopt_state:
model._disable_modelopt_save = True
model.save_pretrained(export_dir, state_dict=post_state_dict)
model.save_pretrained(
export_dir, state_dict=post_state_dict, save_modelopt_state=save_modelopt_state
)
original_config = f"{export_dir}/config.json"
config_data = {}
@@ -416,11 +518,8 @@ def export_hf_checkpoint(
json.dump(config_data, file, indent=4)
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."
" can be saved with torch.save for further inspection."
)
raise e
modelopt/torch/export/unified_export_megatron.py
+33 -3
View File
@@ -501,8 +501,8 @@ class GPTModelExporter:
return all_rules
def _get_weight_scales(self, quantized_state: dict[str, Any], qformat: str):
weight_scale = quantized_state.get("weight_scale")
weight_scale_2 = quantized_state.get("weight_scale_2")
weight_scale = quantized_state.pop("weight_scale", None)
weight_scale_2 = quantized_state.pop("weight_scale_2", None)
if weight_scale is not None:
weight_scale = weight_scale.clone().detach()
@@ -539,6 +539,12 @@ class GPTModelExporter:
weight_scale_2,
block_size,
)
self._state_dict[prefix + "weight_scale"] = weight_scale.detach().clone()
if weight_scale_2 is not None:
if len(weight_scale_2.shape) > 0:
raise ValueError("weight_scale_2 must be a scalar!")
self._state_dict[prefix + "weight_scale_2"] = weight_scale_2.detach().clone()
for key, val in name_to_value.items():
if key == "output_scale" and skip_output_scale:
@@ -589,6 +595,12 @@ class GPTModelExporter:
self._state_dict[gate_proj_prefix + "weight_scale"] = gate_proj_weight_scale
self._state_dict[up_proj_prefix + "weight_scale"] = up_proj_weight_scale
if weight_scale_2 is not None:
if len(weight_scale_2.shape) > 0:
raise ValueError("weight_scale_2 must be a scalar!")
self._state_dict[gate_proj_prefix + "weight_scale_2"] = weight_scale_2.detach().clone()
self._state_dict[up_proj_prefix + "weight_scale_2"] = weight_scale_2.detach().clone()
# weight and weight_scale have been pop out.
for key, val in name_to_value.items():
gate_proj_key = gate_proj_prefix + key
@@ -678,7 +690,11 @@ class GPTModelExporter:
]
else:
# per-tensor scaling
proj_weight_scales = [weight_scale.detach().clone()] * 3
proj_weight_scales = [
weight_scale.detach().clone(),
weight_scale.detach().clone(),
weight_scale.detach().clone(),
]
for weight, scale, key in zip(proj_weights, proj_weight_scales, proj_keys):
quantized_weight = to_quantized_weight(
@@ -691,6 +707,12 @@ class GPTModelExporter:
self._state_dict[key] = quantized_weight
self._state_dict[key + "_scale"] = scale
if weight_scale_2 is not None:
if len(weight_scale_2.shape) > 0:
raise ValueError("weight_scale_2 must be a scalar!")
for weight, scale, key in zip(proj_weights, proj_weight_scales, proj_keys):
self._state_dict[key + "_scale_2"] = weight_scale_2.detach().clone()
# weight and weight_scale have been pop out.
for key, val in name_to_value.items():
q_proj_key = q_proj_prefix + key
@@ -699,6 +721,14 @@ class GPTModelExporter:
if key == "output_scale":
self._state_dict[prefix + k_scale_name] = val.detach().clone()
self._state_dict[prefix + v_scale_name] = val.detach().clone()
elif key == "bias":
# Slice bias similar to weight
bias = val.detach().clone()
bias = bias.reshape([qkv_total_dim, head_size])
proj_biases = [bias[s].reshape(-1) for s in slices]
proj_bias_keys = [q_proj_prefix + key, k_proj_prefix + key, v_proj_prefix + key]
for bias_tensor, bias_key in zip(proj_biases, proj_bias_keys):
self._state_dict[bias_key] = bias_tensor
else:
self._state_dict[q_proj_key] = val.detach().clone()
self._state_dict[k_proj_key] = val.detach().clone()
modelopt/torch/nas/hparams/concat.py
+1 -1
View File
@@ -45,7 +45,7 @@ class ConcatTracedHp(TracedHp):
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
n_max = 5e7 # takes 4s
# if we have less than n_max combinations, we can iterate over all of them
if n_combos <= n_max:
modelopt/torch/nas/plugins/megatron.py
+381 -1
View File
@@ -15,6 +15,7 @@
"""Plugin to add NAS/Pruning support for megatron-core GPT model."""
from collections.abc import Callable, Sequence
from typing import Any
from warnings import warn
@@ -48,8 +49,10 @@ from megatron.core.transformer.mlp import MLP
from megatron.core.transformer.transformer_layer import TransformerLayer
from modelopt.torch.opt.dynamic import DynamicModule
from modelopt.torch.opt.hparam import HPType
from modelopt.torch.opt.searcher import ConstraintsDict
from modelopt.torch.opt.utils import named_hparams
from modelopt.torch.trace import Symbol
from modelopt.torch.utils import distributed as dist
from modelopt.torch.utils import (
get_module_device,
@@ -83,7 +86,10 @@ except ImportError:
HAS_TE = False
try:
import mamba_ssm # noqa: F401
from megatron.core.models.mamba import MambaModel
from megatron.core.ssm.mamba_layer import MambaLayer
from megatron.core.ssm.mamba_mixer import ExtendedRMSNorm, MambaMixer
SUPPORTED_MODELS[MambaModel] = "megatron.core.models.mamba.MambaModel"
@@ -723,6 +729,367 @@ class _DynamicTransformerLayer(DynamicModule, MambaTransformerLayerMixin):
self.mlp.freeze()
class MambaNumHeadsHp(TracedHp):
"""An hparam for Mamba's num_heads.
Need special handling for active_slice property to trim heads within each group.
"""
def __init__(
self, choices: Sequence[HPType], original: HPType | None = None, ngroups: int = 1
) -> None:
super().__init__(choices, original)
self.ngroups = ngroups
@property
def active_slice(self) -> TracedHp.ActiveSlice:
"""Return the currently active sorted indices by trimming heads within each group."""
if self._slice_order is None:
if self.active == self.max:
return slice(self.active)
slice_order = torch.arange(self.max)
else:
slice_order = self._slice_order
target_nheads_per_group = self.active // self.ngroups
return slice_order.view(self.ngroups, -1)[:, :target_nheads_per_group].flatten() # type: ignore[misc]
class MambaDInnerHp(TracedHp):
"""An hparam for Mamba's d_inner.
Mamba's d_inner is a multiplication of mamba_num_heads and mamba_head_dim hparams.
"""
def __init__(self, mamba_num_heads: MambaNumHeadsHp, mamba_head_dim: TracedHp) -> None:
"""Initialize the Mamba d_inner hparam."""
self._mamba_num_heads = mamba_num_heads
self._mamba_head_dim = mamba_head_dim
choices = self._get_choices()
original = mamba_num_heads.original * mamba_head_dim.original
super().__init__(choices, original)
self._is_configurable = False
self._importance_estimators = None
@property # type: ignore[misc]
def active(self) -> int:
"""Return the active value of the hparam."""
assert isinstance(self._mamba_num_heads.active, int)
assert isinstance(self._mamba_head_dim.active, int)
return self._mamba_num_heads.active * self._mamba_head_dim.active
@property
def active_slice(self) -> TracedHp.ActiveSlice:
"""Return the currently active sorted indices or slice corresponding to the active value."""
num_heads_active_slice = self._mamba_num_heads.active_slice
head_dim_active_slice = self._mamba_head_dim.active_slice
if isinstance(num_heads_active_slice, slice):
num_heads_active_slice = torch.LongTensor(range(num_heads_active_slice.stop))
if isinstance(head_dim_active_slice, slice):
head_dim_active_slice = torch.LongTensor(range(head_dim_active_slice.stop))
indices = torch.arange(self.max).view(self._mamba_num_heads.max, self._mamba_head_dim.max)
active_slice = indices[num_heads_active_slice, :][:, head_dim_active_slice].flatten()
# check if active_slice corresponds to the vanilla slice
if torch.equal(active_slice, torch.arange(self.max)):
return slice(self.max)
return active_slice
def _get_choices(self) -> Sequence[HPType]:
return sorted(
{
num_heads * head_dim
for num_heads in self._mamba_num_heads.choices
for head_dim in self._mamba_head_dim.choices
}
)
def reset_choices(self) -> None:
"""Reset the choices of the Mamba d_inner hparam using updated choices of mamba_num_heads and mamba_head_dim."""
self._choices = self._get_choices()
@property # type: ignore[misc]
def choices(self) -> Sequence[HPType]:
"""Return available choices."""
return self._get_choices()
def _resolve_dependencies(
self, sym: Symbol, get_hp: Callable[[Symbol], TracedHp]
) -> dict[Symbol, TracedHp]:
raise NotImplementedError("MambaDInnerHp does not support `_resolve_dependencies`!")
class _DynamicExtendedRMSNorm(DynamicModule):
"""A ``megatron.core.ssm.mamba_mixer.ExtendedRMSNorm`` (GroupNorm) layer with dynamic hyperparams.
Very similar to _DynamicGroupNorm but with group_size dynamic attribute instead of num_groups.
Will be registered to DMRegistry if Mamba is available.
"""
def _setup(self):
# register hidden_size as hyperparameter
orig_hidden_size = self.weight.shape[0]
num_groups = orig_hidden_size // self.group_size
choices = [
c
for c in range(num_groups, orig_hidden_size + 1)
if c % num_groups == 0 and c % self.group_size == 0
]
self._register_hparam("hidden_size", TracedHp(choices, original=orig_hidden_size))
# register num_groups as a dynamic attribute so group size is same
self._register_temp_attribute("_num_groups", num_groups)
self._register_dynamic_attribute("group_size", self._get_group_size)
# register dynamic attributes
dyn_attrs = ["weight", "bias"]
for attr in dyn_attrs:
self._register_dynamic_attribute(attr, self._cut_to_active_hidden_size)
@staticmethod
def _get_group_size(mod: "_DynamicExtendedRMSNorm", value: int) -> int:
return mod.hidden_size // mod._num_groups
@staticmethod
def _cut_to_active_hidden_size(mod: "_DynamicExtendedRMSNorm", value: torch.Tensor | None):
return get_sliced_tensor(mod, value, "hidden_size")
class _DynamicMambaMixer(DynamicModule):
"""A ``megatron.core.ssm.mamba_mixer.MambaMixer`` layer with dynamic hyperparams.
Will be registered to DMRegistry if Mamba is available.
"""
def _setup(self):
assert self.d_inner == self.nheads * self.headdim, "d_inner must be nheads * headdim"
# Register hyperparameters for Mamba heads and head dimensions
# NOTE: d_model will be overwritten in set_hidden_size_hp to model's hidden_size hp
# along with related hparams (in_proj.input_size, norm.hidden_size, out_proj.output_size)
d_model = TracedHp(list(range(1, self.d_model + 1)))
mamba_num_heads = MambaNumHeadsHp(list(range(1, self.nheads + 1)), ngroups=self.ngroups)
mamba_head_dim = TracedHp(list(range(1, self.headdim + 1)))
d_inner = MambaDInnerHp(mamba_num_heads, mamba_head_dim)
bc = TracedHp([2 * self.ngroups * self.d_state]) # not configurable
self._register_hparam("d_model", d_model)
self._register_hparam("d_inner", d_inner)
self._register_hparam("mamba_num_heads", mamba_num_heads)
self._register_hparam("mamba_head_dim", mamba_head_dim)
self._register_hparam("bc", bc)
self._register_dynamic_attribute("d_inner_local", lambda mod, val: self.d_inner)
# Register dynamic attributes
self._register_dynamic_attribute("nheads", lambda mod, val: self.mamba_num_heads)
self._register_dynamic_attribute("nheads_local", lambda mod, val: self.nheads)
self._register_dynamic_attribute("headdim", lambda mod, val: self.mamba_head_dim)
# Convert to dynamic modules
self.in_proj = DMRegistry.convert(self.in_proj)
self.in_proj.output_size = build_concat_hp(
[d_inner, d_inner, bc, mamba_num_heads]
) # z, x, B, C, dt
conv_dim = build_concat_hp([d_inner, bc]) # z, B, C
self.conv1d = DMRegistry.convert(self.conv1d)
self.conv1d.in_channels = conv_dim
self.conv1d.out_channels = conv_dim
ks = self.conv1d.get_hparam("kernel_size")
ks.choices = [ks.original]
if self.rmsnorm:
self.norm = DMRegistry.convert(self.norm)
self.norm.hidden_size = d_inner
self.out_proj = DMRegistry.convert(self.out_proj)
self.out_proj.input_size = d_inner
# Register dynamic attributes for Mamba-specific parameters
self._register_dynamic_attribute("dt_bias", self._get_dt_bias_A_log_D)
self._register_dynamic_attribute("A_log", self._get_dt_bias_A_log_D)
self._register_dynamic_attribute("D", self._get_dt_bias_A_log_D)
assert not self.D_has_hdim, "D_has_hdim is not supported yet"
# Register importance estimator for mamba heads
self._register_temp_attribute("_activations", None)
self.hook_handle = self.in_proj.register_forward_hook(self._mamba_in_proj_forward_hook)
mamba_num_heads.register_importance(self._estimate_head_importance)
mamba_head_dim.register_importance(self._estimate_head_dim_importance)
@staticmethod
def _get_dt_bias_A_log_D(mod: "_DynamicMambaMixer", data: torch.Tensor) -> torch.Tensor: # noqa: N802
"""Return the sliced data based on mamba_num_heads's active_slice."""
return get_sliced_tensor(mod, data, "mamba_num_heads")
def _estimate_head_and_head_dim_rankings(self):
"""Get the rankings of Mamba heads and head dimensions.
Returns:
head_ranking: Ranking of Mamba heads of shape [mamba_num_heads.max]
head_dim_ranking: Ranking of Mamba head dimensions of shape [mamba_head_dim.max]
"""
scores = self._activations
assert scores is not None, "No activations collected for importance estimation."
max_nheads: int = self.get_hparam("mamba_num_heads").max
max_headdim: int = self.get_hparam("mamba_head_dim").max
max_d_inner: int = self.get_hparam("d_inner").max
target_headdim: int = self.headdim
nheads_per_group: int = max_nheads // self.ngroups
# While there can be many ways of computing the ranking out of z, x, and dt,
# based on ablations in the paper, using `x` is the best way to compute the ranking.
x_indices = torch.arange(max_d_inner, 2 * max_d_inner)
scores_x = scores[x_indices] # shape = [max_d_inner] i.e. [max_nheads * max_headdim]
# Get ranking of all head and target head dimensions (same for each head)
all_head_dim_importance = torch.linalg.vector_norm( # shape = [max_headdim]
scores_x.view(max_nheads, max_headdim), ord=2, dim=0
)
all_head_dim_ranking = all_head_dim_importance.argsort(descending=True).cpu()
target_head_dim_ranking = all_head_dim_ranking[:target_headdim]
# Get ranking of all heads with target head dimensions
target_head_dim_indices_per_head = torch.cat( # shape = [max_nheads * target_headdim]
[i * max_headdim + target_head_dim_ranking for i in range(max_nheads)]
)
# Get ranking of heads (sorted within their group)
groupwise_head_importance = torch.linalg.vector_norm( # shape = [ngroups, nheads_per_group]
scores_x[target_head_dim_indices_per_head].view(
self.ngroups, nheads_per_group, target_headdim
),
ord=2,
dim=2,
)
groupwise_head_ranking = groupwise_head_importance.argsort(dim=1, descending=True).cpu()
group_offsets = torch.arange(self.ngroups).unsqueeze(1) * nheads_per_group
all_head_ranking = (groupwise_head_ranking + group_offsets).flatten()
return all_head_ranking, all_head_dim_ranking
def _estimate_head_importance(self):
"""Get the importance of Mamba heads for sort_parameters()."""
head_ranking, _ = self._estimate_head_and_head_dim_rankings()
print_rank_0("Overriding mamba_num_heads.importance to ranking for simplicity.")
# [HACK] Return ranking instead of importance but disable argsort
# so it skips further sorting and returns same ranking inside sort_parameters()
# NOTE: Trimming should also happen within each group. This is handled in MambaNumHeadsHp.
head_ranking.argsort = lambda *args, **kwargs: head_ranking
return head_ranking
def _estimate_head_dim_importance(self):
"""Get the importance of Mamba head dimensions for sort_parameters()."""
_, head_dim_ranking = self._estimate_head_and_head_dim_rankings()
print_rank_0(
"Overriding mamba_head_dim.importance to correctly rank per group for simplicity."
)
# [HACK] Return ranking instead of importance but disable argsort
# so it skips further sorting and returns same ranking inside sort_parameters()
head_dim_ranking.argsort = lambda *args, **kwargs: head_dim_ranking
return head_dim_ranking
def _mamba_in_proj_forward_hook(self, module, input, output) -> None:
"""Hook to collect activations for importance estimation.
Activations are computed as mean over seq_len and then squared and summed over batch_size.
If we take the square root of the sum, we get the L2 norm of the activations.
"""
# Gather output [seq_len, batch_size, output_size] over all TP regions
# NOTE: This is not used at the moment since we restrict to TP=1
output = gather_from_tensor_model_parallel_region(output[0]).detach()
# Dont aggregate activations from non-max subnets (e.g. from profiling)
if output.shape[-1] != self.in_proj.get_hparam("output_size").max:
return
output = output.to(torch.float32) # use full precision to avoid overflow
activations = output.abs().mean(dim=0) # [batch_size, output_size]
activations = activations.pow(2).sum(dim=0) # [output_size]
if self._activations is None:
self._activations = activations
else:
self._activations += activations
def export(self) -> torch.nn.Module:
"""Export the dynamic module to a torch.nn.Module."""
self.hook_handle.remove()
self.in_proj.export()
self.out_proj.export()
self.conv1d.export()
if self.rmsnorm:
self.norm.export()
super().export()
return self
class _DynamicMambaLayer(DynamicModule, MambaTransformerLayerMixin):
"""A ``megatron.core.ssm.mamba_layer.MambaLayer`` layer with dynamic hyperparams.
Will be registered to DMRegistry if Mamba is available.
"""
def _setup(self):
# Convert to dynamic module
self.mixer = DMRegistry.convert(self.mixer)
self.norm = DMRegistry.convert(self.norm)
self._setup_mixin()
def set_hidden_size_hp(self, hidden_size: TracedHp) -> None:
"""Set the hidden size hyperparameter for the layer."""
self.mixer.d_model = hidden_size
self.mixer.in_proj.input_size = hidden_size
self.mixer.out_proj.output_size = hidden_size
self.norm.num_features = hidden_size
self._register_temp_attribute("max_hidden_size", hidden_size.max)
def modify(
self,
*,
mamba_num_heads_divisor: int = 1,
mamba_head_dim_divisor: int = 1,
**kwargs, # Unused hparams
) -> None:
"""Modify Mamba hyperparameters."""
# Modify MambaMixer hparams
for hp_name, divisor in [
("mamba_num_heads", mamba_num_heads_divisor),
("mamba_head_dim", mamba_head_dim_divisor),
]:
hp = self.mixer.get_hparam(hp_name)
choices = {int(make_divisible(c, divisor)) for c in hp.choices} # type: ignore[arg-type]
hp.choices = list(set(hp.choices) & choices | {hp.original})
def export(self):
"""Export the dynamic module to a torch.nn.Module."""
self._export_mixin()
self.mixer.export()
self.norm.export()
super().export()
return self
def freeze(self):
"""Freeze the hyperparameters."""
self.mixer.freeze()
super().freeze()
if HAS_MAMBA:
DMRegistry.register({ExtendedRMSNorm: "megatron.core.ssm.mamba_mixer.ExtendedRMSNorm"})(
_DynamicExtendedRMSNorm
)
DMRegistry.register({MambaMixer: "megatron.core.ssm.mamba_mixer.MambaMixer"})(
_DynamicMambaMixer
)
DMRegistry.register({MambaLayer: "megatron.core.ssm.mamba_layer.MambaLayer"})(
_DynamicMambaLayer
)
@DMRegistry.register(SUPPORTED_MODELS)
class _DynamicMCoreLanguageModel(DynamicModule):
"""A ``megatron.core.models.gpt.GPTModel`` model with dynamic hyperparams."""
@@ -737,7 +1104,9 @@ class _DynamicMCoreLanguageModel(DynamicModule):
assert self.config.expert_model_parallel_size == 1, "Expert parallel is not supported."
assert self.pre_process == is_pipeline_first_stage()
assert self.post_process == is_pipeline_last_stage()
assert self.position_embedding_type == "rope", "Only rope position embedding is supported."
assert self.position_embedding_type in ["rope", "none"], (
f"Only rope position embedding is supported, got {self.position_embedding_type}."
)
# Register num_layers hparam for depth pruning
self._register_hparam("num_layers", TracedHp(list(range(1, self.config.num_layers + 1))))
@@ -789,6 +1158,10 @@ class _DynamicMCoreLanguageModel(DynamicModule):
self._emb_layernorm_forward_hook
)
)
elif HAS_MAMBA and isinstance(layer, MambaLayer):
self.hook_handles.append(
layer.norm.register_forward_hook(self._emb_layernorm_forward_hook)
)
hidden_size.register_importance(self._estimate_hidden_size_importance) # type: ignore[union-attr]
def _emb_layernorm_forward_hook(self, module, input, output) -> None:
@@ -833,6 +1206,8 @@ class _DynamicMCoreLanguageModel(DynamicModule):
num_heads_per_group_divisor: int = 1,
num_query_groups_divisor: int = 1,
ffn_hidden_size_divisor: int = 1,
mamba_num_heads_divisor: int = 1,
mamba_head_dim_divisor: int = 1,
):
"""Modify the dynamic choices of the module according to provided keyword arguments.
@@ -841,6 +1216,8 @@ class _DynamicMCoreLanguageModel(DynamicModule):
num_heads_per_group_divisor: The divisor of the self-attention num_heads_per_group.
num_query_groups_divisor: The divisor of the self-attention num_query_groups.
ffn_hidden_size_divisor: The divisor of the mlp ffn_hidden_size.
mamba_num_heads_divisor: The divisor of the mamba num_heads.
mamba_head_dim_divisor: The divisor of the mamba head_dim.
"""
hp = self.get_hparam("hidden_size")
choices = {int(make_divisible(c, hidden_size_divisor)) for c in hp.choices} # type: ignore[arg-type]
@@ -851,6 +1228,8 @@ class _DynamicMCoreLanguageModel(DynamicModule):
num_heads_per_group_divisor=num_heads_per_group_divisor,
num_query_groups_divisor=num_query_groups_divisor,
ffn_hidden_size_divisor=ffn_hidden_size_divisor,
mamba_num_heads_divisor=mamba_num_heads_divisor,
mamba_head_dim_divisor=mamba_head_dim_divisor,
)
def _export_drop_layers(self) -> None:
@@ -871,6 +1250,7 @@ class _DynamicMCoreLanguageModel(DynamicModule):
all_pp_layer_scores, layer_scores, group=get_pipeline_model_parallel_group()
)
layer_scores = {k: v for d in all_pp_layer_scores for k, v in d.items()} # type: ignore[attr-defined]
print_rank_0(f"Layerwise scores for depth pruning: {layer_scores}")
assert sorted(layer_scores.keys()) == list(range(1, num_layers_hp.max + 1)) # type: ignore[arg-type]
# sort layers by scores and drop the lowest ones
modelopt/torch/nas/search_space.py
+2 -1
View File
@@ -153,7 +153,8 @@ class SearchSpace(DynamicSpace):
if importance is None:
continue
# compute order from importance and enforce it
order = torch.argsort(importance, descending=True)
# NOTE: use .argsort() instead of torch.argsort() so hp can overwrite the behavior
order = importance.argsort(descending=True)
hp.enforce_order(order)
if verbose:
print(f"Sorted {name} for rank {rank()} with {importance=}")
modelopt/torch/opt/_hooks.py
+2 -11
View File
@@ -20,11 +20,10 @@ from contextlib import ExitStack, nullcontext
import torch
import torch.distributed.checkpoint.state_dict as distributed_state_dict
import torch.nn as nn
from packaging.version import Version
# 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 torch.distributed.fsdp._fully_shard import _fsdp_param
from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
from .dynamic import DynamicModule
@@ -67,14 +66,6 @@ FlatParamHandle._writeback_orig_params_original = FlatParamHandle._writeback_ori
FlatParamHandle._writeback_orig_params = _writeback_orig_param
if Version(torch.__version__) >= Version("2.6"):
from torch.distributed.fsdp._fully_shard import _fsdp_param
from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
else:
from torch.distributed._composable.fsdp import _fsdp_param
from torch.distributed._composable.fsdp._fsdp_param import FSDPParam
def _unsafe_setattr_param_with_dm_check(module: nn.Module, param_name: str, param: nn.Parameter):
"""A batched version of unsafe_setattr_param ensuring compatibility with DMs."""
with module.reset_dynamic_attributes() if isinstance(module, DynamicModule) else nullcontext():
modelopt/torch/opt/plugins/huggingface.py
+2 -3
View File
@@ -96,10 +96,9 @@ def _patch_model_init_for_modelopt(cls, model_path, extra_context=None):
def _new_save_pretrained(self, save_directory, *args, **kwargs):
"""Patch for `cls.save_pretrained` method to save ModelOpt state."""
save_modelopt_state = kwargs.pop("save_modelopt_state", True)
outputs = self._modelopt_cache["save_pretrained"](self, save_directory, *args, **kwargs)
if ModeloptStateManager.is_converted(self) and not getattr(
self, "_disable_modelopt_save", False
):
if save_modelopt_state and ModeloptStateManager.is_converted(self):
path = _get_modelopt_state_path(save_directory)
torch.save(modelopt_state(self), path)
print_rank_0(f"Saved ModelOpt state to {path}")
modelopt/torch/opt/plugins/mcore_dist_checkpointing.py
+4 -8
View File
@@ -42,7 +42,7 @@ def remove_per_module_state(
"""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
and expert parallelism. As a result, the metadata must be stored as several ShardedObject with
global and local layer offset mapping.
Args:
@@ -57,6 +57,8 @@ def remove_per_module_state(
if metadata is not None:
_ = metadata.pop("quantizer_state", None)
_ = metadata.pop("subnet_config", None)
_ = metadata.pop("real_quantizer_state", None)
_ = metadata.pop("q_tensor_state", None)
else:
config["metadata"] = {}
@@ -206,10 +208,4 @@ def restore_sharded_modelopt_state(
#
model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state)
try:
_load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix)
except: # noqa: E722
# [WAR]: nemo2 is calling this function with an empty prefix.
# The prefix however should be `module.` instead. This should be fixed
# from the NeMo side. This is just a WAR.
_load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, "module.")
_load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix)
modelopt/torch/opt/utils.py
+2 -7
View File
@@ -20,6 +20,8 @@ from collections.abc import Generator
from contextlib import contextmanager
import torch.nn as nn
from torch.distributed._composable_state import _get_module_state
from torch.distributed.fsdp import FSDPModule
from modelopt.torch.utils import unwrap_model
@@ -92,13 +94,6 @@ def forward_with_reshard(model: nn.Module):
2) use this context manager to reshard FSDPParam in the root module after forward
passes.
"""
try:
from torch.distributed._composable_state import _get_module_state
from torch.distributed.fsdp import FSDPModule
except Exception:
# If FSDP imports fail, act as a null context manager
yield
return
def _lazy_init_retain_mesh_info(self):
if self._fsdp_param_group and not hasattr(self, "_post_forward_mesh_info_before"):
modelopt/torch/prune/plugins/__init__.py
-1
View File
@@ -19,7 +19,6 @@ from modelopt.torch.utils import import_plugin
with import_plugin("mcore_gpt_minitron"):
from .mcore_gpt_minitron import *
from .megatron import *
with import_plugin("transformers"):
from .transformers import *
modelopt/torch/prune/plugins/mcore_gpt_minitron.py
+27 -7
View File
@@ -15,9 +15,12 @@
"""Module implementing top-level ``mcore_gpt_minitron`` pruning handler for NVIDIA Megatron-Core / NeMo models.
Minitron pruning algorithm uses activation magnitudes to estimate importance of neurons / attention heads in the model.
Minitron pruning algorithm uses activation magnitudes to estimate importance of neurons / attention heads / mamba heads
in the model.
More details on Minitron pruning algorithm can be found here: https://arxiv.org/pdf/2407.14679
Supports both GPT (attention-based) and Mamba (state-space) models, as well as hybrid models with both types of layers.
Actual dynamic module implementations are at :mod:`modelopt.torch.nas.plugins.megatron`.
"""
@@ -27,6 +30,7 @@ from pydantic import create_model
# isort: off
# import nas plugin to check if it is enabled else raises an Exception
from modelopt.torch.nas.plugins.megatron import * # noqa: F403
from modelopt.torch.nas.plugins.megatron import HAS_MAMBA
# isort: on
from modelopt.torch.nas.conversion import NASModeRegistry
@@ -46,6 +50,9 @@ SUPPORTED_HPARAMS = {
"num_attention_heads",
"num_query_groups",
"hidden_size",
# TODO: enable mamba head pruning after debugging
# "mamba_num_heads",
# "mamba_head_dim",
# Depth pruning
"num_layers",
}
@@ -83,6 +90,7 @@ def get_supported_model_config_map() -> dict[type, str]:
return supported_model_config_map
# TODO: Update mode name
class MCoreGPTMinitronSearcher(BaseSearcher):
"""Searcher for Minitron pruning algorithm."""
@@ -130,6 +138,10 @@ class MCoreGPTMinitronSearcher(BaseSearcher):
# 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:
assert export_config["num_attention_heads"] % export_config["num_query_groups"] == 0, (
f"num_attention_heads ({export_config['num_attention_heads']}) must be divisible by"
f" num_query_groups ({export_config['num_query_groups']})!"
)
export_config["num_heads_per_group"] = (
export_config["num_attention_heads"] // export_config["num_query_groups"]
)
@@ -197,12 +209,20 @@ MCoreGPTMinitronConfig: type[ModeloptBaseConfig] = create_model(
"num_query_groups_divisor": 1,
"ffn_hidden_size_divisor": 64,
},
"megatron.core.models.mamba.MambaModel": {
"hidden_size_divisor": 64,
"num_heads_per_group_divisor": 1,
"num_query_groups_divisor": 1,
"ffn_hidden_size_divisor": 64,
},
**(
{
"megatron.core.models.mamba.MambaModel": {
"hidden_size_divisor": 64,
"num_heads_per_group_divisor": 1,
"num_query_groups_divisor": 1,
"ffn_hidden_size_divisor": 64,
"mamba_num_heads_divisor": 4,
"mamba_head_dim_divisor": 4,
}
}
if HAS_MAMBA
else {}
),
},
doc='Configuration for the ``"mcore_gpt_minitron"`` mode.',
),
modelopt/torch/prune/plugins/megatron.py
-18
View File
@@ -1,18 +0,0 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""[Deprecated]. Please use :mod:`modelopt.torch.prune.plugins.mcore_gpt_minitron` instead."""
from .mcore_gpt_minitron import * # noqa: F403
modelopt/torch/quantization/backends/fp8_per_tensor_gemm.py
+13 -2
View File
@@ -113,7 +113,9 @@ class Fp8PerTensorLinear(Function):
"""Linear layer with FP8 per tensor quantization."""
@staticmethod
def forward(ctx, quant_module, input_tensor, weight, bias=None):
def forward(
ctx, quant_module, input_tensor, weight, bias=None, allreduce_dgrad=False, tp_group=None
):
"""Forward method."""
ctx.save_for_backward(
input_tensor if weight.requires_grad else None,
@@ -122,6 +124,10 @@ class Fp8PerTensorLinear(Function):
getattr(quant_module.weight_quantizer, "_scale", None),
)
ctx.block_sizes = getattr(quant_module.weight_quantizer, "_block_sizes", None)
ctx.allreduce_dgrad = allreduce_dgrad
ctx.tp_group = tp_group
ret = fp8_per_tensor_gemm(quant_module, input_tensor, bias)
return ret
@@ -147,7 +153,12 @@ class Fp8PerTensorLinear(Function):
if compute_bias_grad is not None:
# Sum all dimensions except the last one
grad_bias = grad_outputs.sum(dim=list(range(grad_outputs.dim() - 1)))
return None, grad_input, grad_weight, grad_bias
if ctx.allreduce_dgrad:
# All-reduce. Note: here async and sync are effectively the same.
torch.distributed.all_reduce(grad_input, group=ctx.tp_group)
return None, grad_input, grad_weight, grad_bias, None, None
@classmethod
def apply(cls, *args, **kwargs):
modelopt/torch/quantization/backends/nvfp4_gemm.py
+12 -2
View File
@@ -76,7 +76,9 @@ class Nvfp4Linear(Function):
"""Linear layer with FP4 quantization."""
@staticmethod
def forward(ctx, quant_module, input_tensor, weight, bias=None):
def forward(
ctx, quant_module, input_tensor, weight, bias=None, allreduce_dgrad=False, tp_group=None
):
"""Forward method."""
ctx.save_for_backward(
input_tensor if weight.requires_grad else None,
@@ -85,6 +87,9 @@ class Nvfp4Linear(Function):
getattr(quant_module.weight_quantizer, "_scale", None),
getattr(quant_module.weight_quantizer, "_double_scale", None),
)
ctx.allreduce_dgrad = allreduce_dgrad
ctx.tp_group = tp_group
ret = nvfp4_gemm(quant_module, input_tensor, bias)
return ret
@@ -113,7 +118,12 @@ class Nvfp4Linear(Function):
if compute_bias_grad is not None:
# Sum all dimensions except the last one
grad_bias = grad_outputs.sum(dim=list(range(grad_outputs.dim() - 1)))
return None, grad_input, grad_weight, grad_bias
if ctx.allreduce_dgrad:
# All-reduce. Note: here async and sync are effectively the same.
torch.distributed.all_reduce(grad_input, group=ctx.tp_group)
return None, grad_input, grad_weight, grad_bias, None, None
@classmethod
def apply(cls, *args, **kwargs):
modelopt/torch/quantization/compress.py
+31 -4
View File
@@ -51,9 +51,26 @@ RealQuantModuleRegistry = _DMRegistryCls("RealQuant")
def compress_convert(
model, config: CompressConfig, use_real_quant_gemm: bool = True
model,
config: CompressConfig,
use_real_quant_gemm: bool = True,
skip_real_quantize_weight: bool = False,
) -> ConvertReturnType:
"""Compress entry point."""
"""Compress entry point.
This function converts the model to a real quantized model.
Args:
model: The model to compress.
config: The compression configuration.
use_real_quant_gemm: Whether to use real quantize GEMM implementation.
skip_real_quantize_weight: Whether to skip the real quantize step. Currently, it is
only set to True in the Megatron restore path to unify the restore behavior regardless
of whether the model is initialized on meta device or not.
Returns:
The compressed model.
"""
for _, module in model.named_modules():
if is_quantized_linear(module) and type(module) not in RealQuantModuleRegistry:
class_to_register = RealQuantLinear
@@ -90,7 +107,8 @@ def compress_convert(
f"Invalid compression configuration: {to_compress}, expected a boolean as value."
)
# If real quant quantizer is present, real quantize the weights.
pack_real_quantize_weight(model)
if not skip_real_quantize_weight:
pack_real_quantize_weight(model)
def _has_qtensorwrapper(module):
if hasattr(module, "weight") and isinstance(module.weight, QTensorWrapper):
@@ -118,12 +136,21 @@ def compress_convert(
def compress_restore(
model: ModelLikeModule, config: CompressConfig, metadata: MetadataDict
) -> nn.Module:
"""Restore the model from the compressed state."""
"""Restore the model from the compressed state.
Note:
When restoring Megatron distributed checkpoint, real_quantizer_state and q_tensor_state
have been removed from metadata and stored as a part of QuantModule.extra_state.
Restoring happends in set_extra_state when load_state_dict is called. We also skip real
quantize weight (skip_real_quantize_weight). All these steps are
delayed. For details, see plugins.megatron.quant_module_set_extra_state.
"""
# Compress with dummy weights
model, _ = compress_convert(
model,
config,
use_real_quant_gemm=metadata.get("use_real_quant_gemm", False),
skip_real_quantize_weight=("q_tensor_state" not in metadata),
)
# restore scale state in weight quantizer
if "real_quantizer_state" in metadata:
modelopt/torch/quantization/conversion.py
+5 -2
View File
@@ -174,7 +174,10 @@ def quantizer_state(model: nn.Module) -> dict[str, Any]:
def replace_quant_module(model: nn.Module, version=None, registry=QuantModuleRegistry):
"""Recursively replace the module with quantized module."""
from .plugins.custom import register_custom_model_plugins_on_the_fly
from .plugins.custom import (
register_custom_model_plugins_on_the_fly,
register_custom_post_conversion_plugins,
)
assert not is_quantized(model), "Model must not be quantized!"
register_custom_model_plugins_on_the_fly(model)
@@ -183,7 +186,7 @@ def replace_quant_module(model: nn.Module, version=None, registry=QuantModuleReg
model = registry.convert(model)
_replace_quant_module(model, version=version, registry=registry)
register_custom_post_conversion_plugins(model)
replaced_modules = sum(isinstance(m, TensorQuantizer) for _, m in model.named_modules())
print(f"Inserted {replaced_modules} quantizers")
modelopt/torch/quantization/model_calib.py
+34 -17
View File
@@ -33,14 +33,30 @@ from .nn import QuantModule, SequentialQuantizer, TensorQuantizer
from .utils import (
enable_weight_access_and_writeback,
is_quantized_column_parallel_linear,
is_quantized_layer_with_weight,
is_quantized_linear,
is_quantized_row_parallel_linear,
quantizer_attr_names,
weight_attr_names,
)
__all__ = ["awq", "max_calibrate", "smoothquant", "svdquant"]
def weight_only_quantize(model: nn.Module):
"""Just quantize the weights of the model."""
seen_modules = set()
for name, module in model.named_modules():
if module in seen_modules:
continue
for weight_name in weight_attr_names(module):
with enable_weight_access_and_writeback(module, model):
weight_quantizer = getattr(
module, quantizer_attr_names(weight_name).weight_quantizer
)
weight_quantizer(getattr(module, weight_name))
seen_modules.add(module)
@torch.no_grad()
def max_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None, distributed_sync=True):
"""Calibrate the model using max.
@@ -55,18 +71,9 @@ def max_calibrate(model: nn.Module, forward_loop: ForwardLoop | None = None, dis
"""
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"):
with enable_weight_access_and_writeback(module, model):
module.weight_quantizer(module.weight)
seen_modules.add(module)
forward_loop(model)
weight_only_quantize(model)
else:
forward_loop(model)
finish_stats_collection(model)
if not distributed_sync:
@@ -629,7 +636,8 @@ def awq_lite(
AWQLiteHelper.cache_mode = False
print("Searching awq_lite parameters...")
forward_loop(model)
with torch.no_grad():
forward_loop(model)
def postprocess(module):
update_best_params(module)
@@ -711,7 +719,8 @@ def awq_clip(
else:
self.loss = {
k: torch.zeros(
(co, math.ceil(ci / self.block_size)), device=module.weight.device
(co, math.ceil(ci / self.block_size)),
device=module.weight.device,
)
for k in clip_ratios
}
@@ -791,7 +800,10 @@ def awq_clip(
]
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,
(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
@@ -818,7 +830,12 @@ def awq_clip(
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)
_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(
modelopt/torch/quantization/nn/modules/quant_linear.py
+55 -19
View File
@@ -139,6 +139,26 @@ class RealQuantLinear(QuantModule):
list_of_scale_tensors = ["_scale", "double_scale", "_scale_zeros"]
allow_real_quant_gemm = True
@property
def _should_run_real_quant_gemm(self):
return (
hasattr(self, "_use_real_quant_gemm")
and self._use_real_quant_gemm
and not (self.input_quantizer.is_enabled and self.input_quantizer._if_calib)
and self.allow_real_quant_gemm
)
def get_real_quant_gemm_impl(self, input, *args, **kwargs) -> bool:
"""Get the real quant GEMM implmenetation base on input arguments."""
if not hasattr(self, "_real_quant_gemm_impl"):
self._real_quant_gemm_impl = backends.gemm_registry.find_match(
self, input, *args, **kwargs
)
if self._real_quant_gemm_impl is None:
warnings.warn(f"RealQuantLinear: No real-quant GEMM found: {self}.")
return self._real_quant_gemm_impl is not None
def forward(self, input, *args, **kwargs):
"""RealQuant layer forward function."""
# For torch.export, we use the default fake quant
@@ -146,27 +166,16 @@ class RealQuantLinear(QuantModule):
return super().forward(input, *args, **kwargs)
# Check if real-quant GEMM is available
if (
hasattr(self, "_use_real_quant_gemm")
and self._use_real_quant_gemm
and input.numel() > 1
# If we need to calibrate the input, we fallback to fake quant
and not (self.input_quantizer.is_enabled and self.input_quantizer._if_calib)
# Our forward might not work for every implementation, so we allow user to disable it
and self.allow_real_quant_gemm
):
if self._should_run_real_quant_gemm and input.numel() > 1:
# If the input is not quantized, we use the default GEMM.
real_quant_gemm = (
self._real_quant_gemm_cache
if hasattr(self, "_real_quant_gemm_cache")
else backends.gemm_registry.find_match(self, input, args, kwargs)
)
self.get_real_quant_gemm_impl(input, *args, **kwargs)
# Note: We cache the real-quant GEMM function to avoid matching overhead.
# This assumes that the function will not change after the first call.
if real_quant_gemm:
self._real_quant_gemm_cache = real_quant_gemm
output = real_quant_gemm(self, input, self.weight, self.bias, *args, **kwargs)
if self._real_quant_gemm_impl:
output = self._real_quant_gemm_impl(
self, input, self.weight, self.bias, *args, **kwargs
)
return (
self.output_quantizer(output) if hasattr(self, "output_quantizer") else output
)
@@ -210,12 +219,39 @@ class RealQuantLinear(QuantModule):
# Function to dynamically override load_state_dict
dynamically_update_state_methods(self)
def _apply(self, fn):
def _apply(self, fn, recurse=True):
"""Override the _apply method to ensure that the weight is real-quantized."""
# Check if fn is a tensor_cast_fun and print warning if so
if hasattr(fn, "__name__") and "tensor_cast" in fn.__name__.lower():
warnings.warn("RealQuantLinear does not support tensor_cast_fun.")
return self
elif "to_empty" in str(fn):
# Handle meta device materialization using to_empty(). to_empty() calls _apply()
# with a lambda function over torch.empty_like. The function's name is <lambda>;
# hence we can only detect to_empty keyword in the __repr__. We take care
# recursive _apply over all suubmodules (e.g. input and weight quantizers are
# submodules). Parameters and buffer are all taken care.
#
# Since the parameter is reassigned, the QTensorWrapper will be gone entirely.
# Hence we custom the behavior such that the QTensorWrapper is reapplied afterward.
if recurse:
for module in self.children():
module._apply(fn, recurse=recurse)
for key, param in self._parameters.items():
if param is None:
continue
with torch.no_grad():
if "weight" in key and isinstance(param, QTensorWrapper):
self._parameters[key] = QTensorWrapper(fn(param), metadata=param.metadata)
else:
self._parameters[key] = torch.nn.Parameter(fn(param), requires_grad=False)
for key, buf in self._buffers.items():
if buf is not None:
self._buffers[key] = fn(buf)
return self
else:
# Process the function normally
return super()._apply(fn)
return super()._apply(fn, recurse=recurse)
modelopt/torch/quantization/nn/modules/quant_module.py
+7 -7
View File
@@ -59,8 +59,8 @@ class QuantModule(DynamicModule):
"""Post-restore to correctly configure the TensorQuantizer states.
TensorQuantizer states are restored to their shape before saving. Now we need to further configure them.
1. For non-sharded modules this simply involves moving the TensorQuantizer states to the right device and
dtype. This applies for regular Pytorch models and HuggingFace models.
1. For non-sharded modules this simply involves moving the TensorQuantizer states to the right device.
This applies for regular Pytorch models and HuggingFace models.
2. For sharded modules the restored states of TensorQuantizer could be incorrect. This is because
parallelism such as TP might have been changed between saving and resoring. So we need to re-calculate
the state shapes. Hence such modules should override this and implement their own logic.
@@ -75,16 +75,16 @@ class QuantModule(DynamicModule):
if non_tq_param_or_buffer is None:
warnings.warn(
f"Could not identify the device and dtype for TensorQuantizer states of {prefix}. "
"Please move the model to the right device and dtype now. This can be done by calling "
"`model.to(device, dtype)`."
f"Could not identify the device for TensorQuantizer states of {prefix}. "
"Please move the model to the right device now. This can be done by calling "
"`model.to(device)`."
)
return
# Move the TensorQuantizer states to the right device and dtype
# Move the TensorQuantizer states to the right device (dtype should have been restored).
for module in self.modules():
if isinstance(module, TensorQuantizer):
module.to(non_tq_param_or_buffer.device, non_tq_param_or_buffer.dtype)
module.to(non_tq_param_or_buffer.device)
def fold_weight(self):
"""Fold the weight for faster eval."""
modelopt/torch/quantization/nn/modules/tensor_quantizer.py
+7 -4
View File
@@ -581,6 +581,7 @@ class TensorQuantizer(nn.Module):
outputs, scales = NF4QTensor.quantize(
inputs, self._block_sizes[-1], self._block_sizes["scale_block_sizes"][-1]
)
_scale, _double_scale, _scale_zeros = NF4QTensor.double_quantization(
scales,
self._block_sizes["scale_block_sizes"][-1],
@@ -1095,11 +1096,11 @@ class TensorQuantizer(nn.Module):
"""
metadata = {"params": {}, "buffers": {}}
for k, v in self._parameters.items():
metadata["params"][k] = {"shape": v.shape}
metadata["params"][k] = {"shape": v.shape, "dtype": v.dtype}
for k, v in self._buffers.items():
if k in self._non_persistent_buffers_set:
continue
metadata["buffers"][k] = {"shape": v.shape}
metadata["buffers"][k] = {"shape": v.shape, "dtype": v.dtype}
return metadata
def _del_pytorch_state(self):
@@ -1112,9 +1113,11 @@ class TensorQuantizer(nn.Module):
# Lets delete existing parameters and buffers and create fresh ones
self._del_pytorch_state()
for k, v in metadata.get("params", {}).items():
self.register_parameter(k, nn.Parameter(torch.empty(v["shape"])))
dtype = v.get("dtype", None)
self.register_parameter(k, nn.Parameter(torch.empty(v["shape"], dtype=dtype)))
for k, v in metadata.get("buffers", {}).items():
self.register_buffer(k, torch.empty(v["shape"]))
dtype = v.get("dtype", None)
self.register_buffer(k, torch.empty(v["shape"], dtype=dtype))
def get_modelopt_state(self, properties_only: bool = False) -> dict[str, Any]:
"""Get meta state to be saved in checkpoint.
modelopt/torch/quantization/plugins/__init__.py
+3
View File
@@ -64,5 +64,8 @@ with import_plugin("transformer_engine"):
with import_plugin("transformers trainer"):
from .transformers_trainer import *
with import_plugin("transformers"):
from .transformers import *
with import_plugin("vllm"):
from .vllm import *
modelopt/torch/quantization/plugins/custom.py
+18 -8
View File
@@ -31,6 +31,7 @@ from ..nn.modules.quant_linear import _QuantLinear
from ..utils import multi_context, replace_function
CUSTOM_MODEL_PLUGINS = set()
CUSTOM_POST_CONVERSION_PLUGINS = set()
# TODO: This is a temporary solution
@@ -41,6 +42,12 @@ def register_custom_model_plugins_on_the_fly(model):
callback(model)
def register_custom_post_conversion_plugins(model):
"""Registers custom modules as QUANT_MODULE after conversion."""
for callback in CUSTOM_POST_CONVERSION_PLUGINS:
callback(model)
class _QuantFunctionalMixin(QuantModule):
"""Mixin class for quantized functionals.
@@ -105,6 +112,10 @@ class _ParallelLinear(_QuantFunctionalMixin, QuantModule):
self.output_quantizer = TensorQuantizer(_QuantLinear.default_quant_desc_output)
self.output_quantizer.disable()
# Memorize the original weight.dtype for modelopt_post_restore given that
# the dtype can change later.
self.original_weight_dtype = None if self.weight is None else self.weight.dtype
def modelopt_post_restore(self, prefix: str = ""):
"""Post restore to correctly configure the TensorQuantizer states for MCore/distributed frameworks.
@@ -126,11 +137,6 @@ class _ParallelLinear(_QuantFunctionalMixin, QuantModule):
def _has_state(quantizer, name):
# Handling for SequentialQuantizer
quantizer = quantizer[0] if isinstance(quantizer, SequentialQuantizer) else quantizer
if self.is_version_less_than("0.29") and "dev" not in __version__:
# For backward compatibility, previously we used to save a boolean attribute "_has_amax"
# to indicate if the quantizer has amax.
return hasattr(quantizer, "_has" + name)
return hasattr(quantizer, name)
if self.weight is None:
@@ -147,19 +153,23 @@ class _ParallelLinear(_QuantFunctionalMixin, QuantModule):
if hasattr(self.input_quantizer, "_pre_quant_scale"):
delattr(self.input_quantizer, "_pre_quant_scale")
pqs = torch.zeros(
(self.weight.shape[1]), device=self.weight.device, dtype=self.weight.dtype
(self.weight.shape[1]), device=self.weight.device, dtype=self.original_weight_dtype
)
self.input_quantizer.register_buffer("_pre_quant_scale", pqs)
if _has_state(self.input_quantizer, "_amax"):
self.input_quantizer.reset_amax()
dummy_input = torch.ones(
(1, 1, self.weight.shape[1]), device=self.weight.device, dtype=self.weight.dtype
(1, 1, self.weight.shape[1]),
device=self.weight.device,
dtype=self.original_weight_dtype,
)
max_calibrate(self.input_quantizer, lambda iq: iq(dummy_input), distributed_sync=False)
if _has_state(self.output_quantizer, "_amax"):
self.output_quantizer.reset_amax()
dummy_input = torch.ones(
(1, 1, self.weight.shape[0]), device=self.weight.device, dtype=self.weight.dtype
(1, 1, self.weight.shape[0]),
device=self.weight.device,
dtype=self.original_weight_dtype,
)
max_calibrate(self.output_quantizer, lambda oq: oq(dummy_input), distributed_sync=False)
# If there are any other states, lets move them to the correct device
modelopt/torch/quantization/plugins/huggingface.py
+103 -19
View File
@@ -41,7 +41,7 @@ from ..nn import QuantInputBase, QuantModule, QuantModuleRegistry, TensorQuantiz
from ..nn.modules.quant_linear import _QuantLinear
from ..utils import replace_function
from .attention import register_attention_for_kv_quant
from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear
from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear, _QuantFunctionalMixin
if TYPE_CHECKING:
from types import ModuleType
@@ -69,6 +69,12 @@ class _QuantAttention(QuantModule):
)
def forward(self, *args, **kwargs):
"""Forward method for KV cache quantization compatible with new_attention_interface in transformers >= 4.48.0.
The forward method is used to patch the attention interface with _quantized_attention.
Once output tensors are generated, it restores the original attention interface.
"""
def _is_eager_attention():
if self.config._attn_implementation == "eager":
return True
@@ -80,6 +86,14 @@ class _QuantAttention(QuantModule):
# Get the original transformers module before wrapped in any ModelOpt DynamicModule
module: ModuleType = inspect.getmodule(self.get_attn_type(self))
# Preprocessing logic to patch attention interface
original_attention_interface = (
module.eager_attention_forward
if _is_eager_attention()
else module.ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
)
patch_fn = partial(self._quantized_attention, original_attention_interface)
if _is_eager_attention():
if not hasattr(module, "eager_attention_forward"):
raise AssertionError(
@@ -87,26 +101,20 @@ class _QuantAttention(QuantModule):
"Please use a different attention implementation such as `sdpa` by setting "
"`model.config._attn_implementation = 'sdpa'` before quantization."
)
original_attention_interface = module.eager_attention_forward
module.eager_attention_forward = partial( # type: ignore[attr-defined]
self._quantized_attention, original_attention_interface
)
module.eager_attention_forward = patch_fn # type: ignore[attr-defined]
else:
original_attention_interface = module.ALL_ATTENTION_FUNCTIONS[
self.config._attn_implementation
]
module.ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] = partial(
self._quantized_attention, original_attention_interface
)
module.ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] = patch_fn
outputs = super().forward(*args, **kwargs)
if _is_eager_attention():
module.eager_attention_forward = original_attention_interface # type: ignore[attr-defined]
else:
module.ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] = (
original_attention_interface
)
try:
outputs = super().forward(*args, **kwargs)
finally:
# Cleanup logic to restore the original attention interface
if _is_eager_attention():
module.eager_attention_forward = original_attention_interface # type: ignore[attr-defined]
else:
module.ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] = (
original_attention_interface
)
return outputs
@@ -521,6 +529,82 @@ except ImportError:
pass
class _QuantGptOssExperts(_QuantFunctionalMixin):
"""Quantized wrapper for `transformers.GptOssExperts`.
Quantizes `gate_up_proj` and `down_proj` weights via dynamic attributes inside `quantize_weight()`.
Activations into `gate_up_proj` are quantized by `gate_up_proj_input_quantizer`. For `down_proj`
activation quantiation, we intercept `torch.Tensor.__matmul__`/`torch.bmm` and quantize inputs
on every second call (since the first call computes `gate_up_proj` outputs and second call
computes `down_proj` outputs).
"""
def _setup(self):
def _get_quantized_weight(quantizer, module, weight):
if module._enable_weight_quantization:
return quantizer(weight)
return weight
assert not hasattr(self, "kernel_layer_name"), (
"ModelOpt quantization does not support patched forward for kernel_hub"
)
self.gate_up_proj_input_quantizer = TensorQuantizer()
self.gate_up_proj_weight_quantizer = TensorQuantizer()
self.down_proj_input_quantizer = TensorQuantizer()
self.down_proj_weight_quantizer = TensorQuantizer()
self._register_temp_attribute("_enable_weight_quantization", False)
self._register_dynamic_attribute(
"gate_up_proj", partial(_get_quantized_weight, self.gate_up_proj_weight_quantizer)
)
self._register_dynamic_attribute(
"down_proj", partial(_get_quantized_weight, self.down_proj_weight_quantizer)
)
self._register_temp_attribute("_down_proj_mul", False)
@property
def functionals_to_replace(self):
def _quantized_bmm(batch1, batch2):
batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1
self._down_proj_mul = not self._down_proj_mul # toggle the flag
return torch._bmm(batch1, batch2)
def _tensor_matmul(self_t, other):
self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t
self._down_proj_mul = not self._down_proj_mul
return torch.matmul(self_t, other)
return [
(torch, "bmm", _quantized_bmm),
(torch.Tensor, "__matmul__", _tensor_matmul),
]
@contextmanager
def quantize_weight(self):
"""Context in which weight is quantized."""
self._enable_weight_quantization = True
yield
self._enable_weight_quantization = False
def forward(
self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None
) -> torch.Tensor:
"""Forward method to add quantization."""
hidden_states = self.gate_up_proj_input_quantizer(hidden_states)
with self.quantize_weight():
return super().forward(hidden_states, router_indices, routing_weights)
try:
from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts
if GptOssExperts not in QuantModuleRegistry:
QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts)
except ImportError:
pass
def register_dbrx_moe_on_the_fly(model):
"""Register DBRX MoE modules as QUANT_MODULE.
modelopt/torch/quantization/plugins/megatron.py
+148 -71
View File
@@ -16,14 +16,12 @@
"""Support quantization for megatron linear layers."""
import warnings
from contextlib import contextmanager
from typing import Any
import megatron.core.parallel_state as mcore_parallel
import megatron.core.tensor_parallel.layers as megatron_parallel
import megatron.core.transformer.mlp as megatron_mlp
import torch
import torch.nn as nn
from megatron.core.transformer import MegatronModule
from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
@@ -41,11 +39,31 @@ from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear
__all__ = []
def real_quant_module_get_extra_state(self) -> dict:
"""Populating real_quantizer_state and q_tensor_state."""
extra_state = {}
if isinstance(self, RealQuantLinear) and isinstance(self.weight, QTensorWrapper):
real_quantizer_state = self.weight_quantizer.get_modelopt_state()
q_tensor_state = self.weight.get_state()
elif isinstance(self, RealQuantLinear):
real_quantizer_state = self.weight_quantizer.get_modelopt_state()
q_tensor_state = {}
else:
real_quantizer_state = None
q_tensor_state = None
extra_state["modelopt_real_quantizer_state"] = real_quantizer_state
extra_state["modelopt_q_tensor_state"] = q_tensor_state
return extra_state
def quant_module_get_extra_state(self) -> dict:
"""Populating the extra_state when state_dict() is called.
quantizer_state is usually stored with in the modelopt_state
metadata where the keys are the full module name. The issue
quantizer_state, real_quantizer_state, and q_tensor_state are usually stored
with in the modelopt_state metadata where the keys are the full module name. The issue
is that NeMo-MCore model's full module name can change
if pipeline-parallelism (PP) and expert-parallelism (EP)
are changing. Alternatively, we store quantizer_state in
@@ -60,16 +78,63 @@ def quant_module_get_extra_state(self) -> dict:
return extra_state
quantizer_state = {}
for name, module in self.named_modules():
if isinstance(module, TensorQuantizer):
quantizer_state[name] = module.get_modelopt_state()
extra_state["modelopt_quantizer_state"] = quantizer_state
# Handle real_quantizer_state and q_tensor_state
extra_state.update(real_quant_module_get_extra_state(self))
return extra_state
def real_quant_module_set_extra_state(self, state: Any):
"""Restore q_tensor_state when load_state_dict() is called.
We skip restoring real_quantizer_state (if exists), since it is the same as
the weight_quantizer fake quantizer_state.
Finally, q_tensor_state is restored if meta device initialization is used. During
meta-device initialization, real_quantize is not called.
QTensorWrapper should replace the original weight parameter. Due to TP, we also need
to adjust q_tensor_data_shape and its metadata shape attribute to use the local weight shape.
When not using meta device initialization, real_quantize is called during compress mode
restore where the QTensor will be recomputed based on the local weights. Hence we don't
need to restore q_tensor_state.
Note:
The entire restore process can happen on meta device and be materialized later
with to_empty(). However, to_empty() will reassign the parameter and the
QTensorWrapper will be removed. We patch RealQuantLinear._apply to preserve
QTensorWarpper when to_empty() is applied.
"""
q_tensor_state = state.get("modelopt_q_tensor_state", None)
if q_tensor_state is not None:
q_tensor_metadata = q_tensor_state["metadata"]
q_tensor_metadata["shape"] = self.weight.shape
q_tensor_data_dtype = q_tensor_state["quantized_data.dtype"]
q_tensor_shape = self.weight.shape
# If q_tensor_data_type is uint8, then it is compressed format of 2 elements.
if q_tensor_data_dtype == torch.uint8:
q_tensor_shape = list(q_tensor_shape)
q_tensor_shape[-1] = q_tensor_shape[-1] // 2
q_tensor_shape = torch.Size(q_tensor_shape)
self._parameters["weight"] = QTensorWrapper(
qtensor=torch.empty(
q_tensor_shape, # Use the local shape directly (TP-aware)
dtype=q_tensor_data_dtype,
device=self.weight.device,
),
metadata=q_tensor_metadata,
)
def quant_module_set_extra_state(self, state: Any):
"""Restore quantizer_state when load_state_dict() is called.
@@ -84,18 +149,29 @@ def quant_module_set_extra_state(self, state: Any):
The 2nd load_state_dict() is loading all states including amax and
scalars. We disable QuantModule.modelopt_post_restore() to avoid
reinitialization since set_extra_state() is called at the end.
We first restore all fake quantizer_state. Per QuantModule can have
weight_quantizer, input_quantizer, and output_quantizer.
Once all quantizer_state are resumed, modelopt_post_restore() is called
to adjust the shape of all buffers (amax, pre_qunat_scale, _scale, ...) since
the local shape can be different from the shape in the state due to change
in tensor parallelism (TP).
"""
if state is None:
if state is None or not self.allow_post_restore:
return
quantizer_state = state.get("modelopt_quantizer_state", None)
if quantizer_state is not None and self.allow_post_restore:
if quantizer_state is not None:
for name, module in self.named_modules():
if isinstance(module, TensorQuantizer):
module.set_from_modelopt_state(quantizer_state[name])
module.set_from_modelopt_state(quantizer_state[name], properties_only=False)
self.modelopt_post_restore()
# Handle real_quantizer_state and q_tensor_state
real_quant_module_set_extra_state(self, state)
self.allow_post_restore = False
@@ -286,8 +362,9 @@ class _QuantMegatronMLP(_MegatronMLP):
]
class _RealQuantMegatronColumnParallelLinear(RealQuantLinear, _MegatronColumnParallelLinear):
allow_real_quant_gemm = False # We don't support real quant gemm for ColumnParallelLinear
class _RealQuantMegatronParallelLinear(RealQuantLinear):
allow_real_quant_gemm = True
_scale_tensor_shard_axis = None
def _parameter_to_keep_in_quantizer_state_dict(self, key):
return any(k in key for k in self.list_of_scale_tensors)
@@ -299,74 +376,74 @@ class _RealQuantMegatronColumnParallelLinear(RealQuantLinear, _MegatronColumnPar
any(k.endswith(suffix) for suffix in self.list_of_scale_tensors)
and state_dict[k].dim() > 1
):
shard_axis_dict[k] = 0
assert self._scale_tensor_shard_axis is not None, (
"scale_tensor_shard_axis is not set, please set it in the subclass"
)
shard_axis_dict[k] = self._scale_tensor_shard_axis
return shard_axis_dict
def modelopt_post_restore(self, prefix: str = ""):
# First follow the fake quant behavior to initialize tensor_quantizers
with _view_as_fake_quant_module(self):
super().modelopt_post_restore(prefix=prefix)
"""Post restore to correctly configure the realquant scales.
# Restore dtype of real quant parameters in tensor_quanitzer
_restore_real_quant_parameters(self)
ModelOpt restores the TensorQuantizer states such as `_amax` and `_pre_quant_scale` to their
shape before saving. However this is not enough for MCore/distributed frameworks since the tensor parallelism
could change between saving and restoring. If the tensor parallelism changes, the shape of the quantizer
states also changes. So we need to re-calculate the quantizer states.
Note:
During real quantization, weight_quantizer._fake_quant is set to False which trigger the real quant
forward path and lead to error. We enable the weight_quantizer fake_quant forward path while recompute
the correct shape.
"""
self.weight_quantizer._fake_quant = True
super().modelopt_post_restore(prefix=prefix)
self.weight_quantizer._fake_quant = False
if hasattr(self.weight_quantizer, "_scale"):
# Recompute all real quantization buffer shapes
self.weight_quantizer._real_quantize(self.weight)
def _forward_impl(self, input, *args, **kwargs):
"""Use real quant gemm if available.
Here the forward is patched such that real quant gemm can be called if available. Both conditions
below must be satisfied (static and dynamic check based on input args) to use the kernel.
Otherwise, we fallback.
Note:
RealQuantLinear.forward() is doing the same check inside and will fall back to use the super
class forward(). This is not desired since _forward_impl introduces much more args and kwargs
while the original forward only takes 1 positional argument. We must above the fallback path
in RealQuantLinear.forward().
"""
if self._should_run_real_quant_gemm and self.get_real_quant_gemm_impl(
input, *args, **kwargs
):
allreduce_dgrad = kwargs.get("allreduce_dgrad", False)
tp_group = kwargs.get("tp_group")
return RealQuantLinear.forward(
self,
input,
allreduce_dgrad=allreduce_dgrad,
tp_group=tp_group,
)
else:
return super()._forward_impl(input, *args, **kwargs)
class _RealQuantMegatronRowParallelLinear(RealQuantLinear, _MegatronRowParallelLinear):
allow_real_quant_gemm = False # We don't support real quant gemm for RowParallelLinear
class _RealQuantMegatronColumnParallelLinear(
_RealQuantMegatronParallelLinear, _MegatronColumnParallelLinear
):
_scale_tensor_shard_axis = 0
def _parameter_to_keep_in_quantizer_state_dict(self, key):
return any(k in key for k in self.list_of_scale_tensors)
def _get_shard_axis_dict(self, state_dict):
shard_axis_dict = super()._get_shard_axis_dict(state_dict)
for k in state_dict:
if (
any(k.endswith(suffix) for suffix in self.list_of_scale_tensors)
and state_dict[k].dim() > 1
):
shard_axis_dict[k] = 1
return shard_axis_dict
def modelopt_post_restore(self, prefix: str = ""):
# Fisrt follow the fake quant behavior to initialize tensor_quantizers
with _view_as_fake_quant_module(self):
super().modelopt_post_restore(prefix=prefix)
# Restore dtype of real quant parameters in tensor_quanitzer
_restore_real_quant_parameters(self)
def forward(self, input, *args, **kwargs):
return _MegatronColumnParallelLinear.forward(self, input, *args, **kwargs)
@contextmanager
def _view_as_fake_quant_module(module: RealQuantLinear):
"""View the module as a fake quantized module."""
# skip if the module is not a RealQuantLinear or QTensorWrapper
if not isinstance(module, RealQuantLinear):
yield
return
assert isinstance(module.weight, QTensorWrapper), "module.weight is not a QTensorWrapper"
try:
quantized_weight = module.weight
dummy_dequantized_weight = torch.rand(
module.weight.metadata["shape"],
dtype=module.weight.metadata["dtype"],
device=module.weight.device,
)
module.weight_quantizer._fake_quant = True
module.weight_quantizer._dequantize = False
module.weight = nn.Parameter(dummy_dequantized_weight)
yield
finally:
module.weight_quantizer._fake_quant = False
module.weight_quantizer._dequantize = True
module.weight = quantized_weight
class _RealQuantMegatronRowParallelLinear(
_RealQuantMegatronParallelLinear, _MegatronRowParallelLinear
):
_scale_tensor_shard_axis = 1
def _restore_real_quant_parameters(module: RealQuantLinear):
"""Restore the real quant parameters in the tensor_quanitzer by performing real weight quantization again."""
dequantized_weight = module.weight_quantizer(module.weight)
module.weight_quantizer._fake_quant = False
module.weight_quantizer._dequantize = False
for k in ["_scale", "double_scale", "_scale_zeros"]:
if hasattr(module.weight_quantizer, k):
delattr(module.weight_quantizer, k)
module.weight = QTensorWrapper(module.weight_quantizer(dequantized_weight))
def forward(self, input, *args, **kwargs):
return _MegatronRowParallelLinear.forward(self, input, *args, **kwargs)
modelopt/torch/quantization/plugins/peft.py
+1 -1
View File
@@ -35,7 +35,7 @@ class _QuantLoraLinear(QuantModule):
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)
return super().forward(x, *args, **kwargs)
x = self.input_quantizer(x)
weight = self.base_layer.weight
modelopt/torch/quantization/plugins/transformers.py
+52
View File
@@ -0,0 +1,52 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Support quantization for Transformers."""
import torch.nn as nn
from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
from .custom import CUSTOM_POST_CONVERSION_PLUGINS
def make_deepspeed_compatible(model: nn.Module):
"""Make the model compatible with DeepSpeed."""
try:
from deepspeed.runtime.zero.parameter_offload import ZeROOrderedDict
except ImportError:
return
is_deepspeed_zero3_enabled = any(
hasattr(module, "_parameters") and isinstance(module._parameters, ZeROOrderedDict)
for module in model.modules()
)
if is_deepspeed_zero3_enabled:
# For zero stage 3, the _parameters is a ZeROOrderedDict, tensor_quantizer._parameters
# is usually a dict, so we need to check if it is a ZeROOrderedDict if the model is wrapped
# by deepspeed.
def _make_deepspeed_compatible(module):
"""Make a module's _parameters DeepSpeed compatible."""
if isinstance(module, TensorQuantizer) and not isinstance(
module._parameters, ZeROOrderedDict
):
module._parameters = ZeROOrderedDict(module._parameters)
# Make all modules DeepSpeed compatible
for module in model.modules():
_make_deepspeed_compatible(module)
CUSTOM_POST_CONVERSION_PLUGINS.add(make_deepspeed_compatible)
modelopt/torch/quantization/plugins/transformers_trainer.py
+44 -6
View File
@@ -15,6 +15,7 @@
"""ModelOpt plugin for transformers Trainer."""
import gc
import os
from contextlib import contextmanager, suppress
@@ -27,7 +28,11 @@ from modelopt.torch.distill import KDLossConfig
from modelopt.torch.distill.mode import _convert_for_kd
from modelopt.torch.distill.plugins.huggingface import KDTrainer
from modelopt.torch.opt.conversion import restore_from_modelopt_state
from modelopt.torch.quantization.utils import is_quantized
from modelopt.torch.quantization.utils import (
calibrate_with_adapters,
disable_lora_quantizers_in_config,
is_quantized,
)
from modelopt.torch.utils import print_rank_0
@@ -35,6 +40,20 @@ class EvalOnlyError(Exception):
"""Exception to raise when evaluation is only needed."""
def check_awq_smoothquant(quant_cfg):
# TODO: Remove this once deepspeed for AWQ and SmoothQuant is added
"""Get the quantization type from the configuration."""
if quant_cfg is None:
return False
algorithm = quant_cfg.get("algorithm", {})
is_awq_smoothquant = False
# Check SmoothQuant and AWQ
if algorithm and ("smoothquant" in algorithm or "awq" in algorithm):
is_awq_smoothquant = True
return is_awq_smoothquant
def get_metrics_with_perplexity(metrics):
"""Add perplexity to the metrics."""
metrics = {"perplexity": float(torch.exp(torch.tensor(metrics["eval_loss"]))), **metrics}
@@ -95,7 +114,7 @@ class QATTrainer(Trainer):
):
"""Initialize the trainer with modelopt states."""
self.quant_args = quant_args
if quant_cfg is None and quant_args.quant_cfg is not None:
if quant_cfg is None and getattr(quant_args, "quant_cfg", None) is not None:
quant_cfg = getattr(mtq, quant_args.quant_cfg)
self.quant_cfg = quant_cfg
self._eval_without_training = False
@@ -105,6 +124,17 @@ class QATTrainer(Trainer):
getattr(self.accelerator.state.fsdp_plugin, "fsdp_version", 1) == 2
)
self._modelopt_state_path = os.path.join(self.args.output_dir, "modelopt_state_train.pth")
# Add lora adapter before quantizing the model
if getattr(self.args, "lora_config", None) is not None:
self.model.add_adapter(self.args.lora_config, adapter_name="adapter")
disable_lora_quantizers_in_config(self.quant_cfg, self.args.lora_config.target_modules)
print_rank_0("Lora adapter added.")
assert self.is_deepspeed_enabled and not check_awq_smoothquant(self.quant_cfg), (
f"QAT DeepSpeed does not currently support AWQ or SmoothQuant: {self.quant_cfg}"
)
# FSDP1 requires pre-restoring the quantized model if the modelopt state exists.
if os.path.exists(self._modelopt_state_path) and not self._is_fsdp2:
self._quantize_model()
@@ -142,10 +172,18 @@ class QATTrainer(Trainer):
)
data_loader = self.get_eval_dataloader(dataset)
forward_loop = self._get_quantize_forward_loop(data_loader, use_eval_loop)
with calibrate_with_adapters(model, self.args):
print_rank_0("Quantizing the model...")
mtq.quantize(model, self.quant_cfg, forward_loop)
print_rank_0("Quantization done!")
if getattr(self.quant_args, "compress", False):
print_rank_0("Compressing model after calibration")
mtq.compress(model)
# Force garbage collection to free up memory
gc.collect()
print_rank_0("Quantizing the model...")
mtq.quantize(model, self.quant_cfg, forward_loop)
print_rank_0("Quantization done!")
print_rank_0(f"Saving modelopt state to {self._modelopt_state_path}")
save_modelopt_state_with_weights(model, self._modelopt_state_path, save_weights=True)
torch.cuda.empty_cache()
@@ -177,7 +215,7 @@ class QATTrainer(Trainer):
self._original_evaluate_on_start = (
self.args.eval_on_start if not self._eval_without_training else True
)
if self.quant_args.quant_cfg is not None and not is_quantized(self.model):
if getattr(self.quant_args, "quant_cfg", None) is not None and not is_quantized(self.model):
self.args.eval_on_start = True
with suppress(EvalOnlyError):
super().train(*args, **kwargs)
modelopt/torch/quantization/qtensor/base_qtensor.py
+195 -11
View File
@@ -16,8 +16,13 @@
"""Base Class for Real Quantized Tensor."""
import enum
import warnings
from contextlib import contextmanager
import torch
from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard
from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
from torch.distributed.tensor import DTensor
class QTensorType(enum.Enum):
@@ -132,6 +137,29 @@ class QTensorWrapper(torch.nn.Parameter):
}
class QFSDPParam(FSDPParam):
"""A Quantized FSDPParam class to make weight updates compatible with BaseQuantizedTensor and QTensorWrapper.
With this class, we can keep track of the quantized tensor's metadata when compressing the weights
and recreate the QTensorWrapper with the correct metadata, when unsharding the FSDPModule.
Args:
qtensor (BaseQuantizedTensor): The quantized tensor to be wrapped.
"""
def __init__(self, *args, **kwargs):
# Store qtensor information
self.metadata = args[0].metadata
super().__init__(*args, **kwargs)
self.init_dtype_attrs(self.mp_policy)
def _setattr_on_modules(self, param: torch.nn.Parameter) -> None:
if not isinstance(param, DTensor):
# Create a QTensorWrapper with the correct metadata during unsharding
param = QTensorWrapper(param, metadata=self.metadata)
super()._setattr_on_modules(param)
# Function to dynamically override load_state_dict
def dynamically_update_state_methods(module):
# Original method
@@ -166,21 +194,177 @@ def dynamically_update_state_methods(module):
module._load_from_state_dict = custom_load_from_state_dict.__get__(module, type(module))
def get_prefixed_param_names(parent_model, target_module):
"""Get parameter names for a target module prefixed with the parent model name.
This function is used to get full parameter name from FSDPParam module_info which stores the
unprefixed parameter name.
"""
target_ids = {id(p) for p in target_module.parameters()}
return next(
(
name.rsplit(".", 1)[0]
for name, param in parent_model.named_parameters()
if id(param) in target_ids
),
None, # default value if no match
)
@contextmanager
def no_requires_grad():
"""Context manager to temporarily set requires_grad to False.
This is used to allow us to call init_sharded_parameter() on the compressed weights. Currently FSDP2 creates
a new parameter with default requires_grad and then update the requires_grad attribute as needed. This
triggers an error when torch.nn.Parameter is called on compressed weights as requires_grad cannot be set to True
for integer tensors.
"""
original_new = torch.nn.Parameter.__new__
def patched_new(cls, data=None, requires_grad=True):
return original_new(cls, data, requires_grad=False)
torch.nn.Parameter.__new__ = patched_new
try:
yield
finally:
torch.nn.Parameter.__new__ = original_new
@contextmanager
def enable_fake_quant(module):
"""Temporarily set the fake_quant attribute of a module to True.
This is used to prevent weight compression from being triggered during an unshard() call.
"""
original_fake_quant = []
for m in module.modules():
if hasattr(m, "weight_quantizer"):
original_fake_quant.append(m.weight_quantizer._fake_quant)
m.weight_quantizer._fake_quant = True
yield
for m in module.modules():
if hasattr(m, "weight_quantizer"):
m.weight_quantizer._fake_quant = original_fake_quant.pop(0)
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."""
# Import SequentialQuantizer here to avoid circular import
from ..nn import SequentialQuantizer
def _compress_and_update_module_weight(module):
"""Compresses and updates module weights if quantizer is enabled. Returns True when compression is applied."""
if hasattr(module, "weight") and (module.weight is None or module.weight.is_meta):
# We dont compress meta tensors or None
return False
if (
hasattr(module, "weight_quantizer")
and module.weight_quantizer.is_enabled
and not module.weight_quantizer._fake_quant
and module.weight.element_size() > 1
):
if force_quantize:
module.weight_quantizer._dequantize = False
real_quant_tensor = module.weight_quantizer(module.weight)
module.weight = QTensorWrapper(real_quant_tensor)
return True
return False
def _create_fsdp_param_mapping(fsdp_param_list, model):
"""Builds a mapping from module name to their corresponding FSDPParam.
Args:
fsdp_param_list (list): List of FSDPParam.
model (nn.Module): FSDP root module.
Returns:
dict: Full parameter name → FSDP parameter.
"""
return {
get_prefixed_param_names(model, param._module_info.module): param
for param in fsdp_param_list
}
def _compress_fsdp_module(fsdp_module):
"""Applies weight compression to an FSDP-wrapped module and updates its sharded parameter group.
This function unshards the FSDP module to access full weights and compresses each eligible submodule’s weights.
A new FSDPParam wrapped with `QFSDPParam` is registered to the FSDPParamGroup for future handling of
sharding and unsharding. The weight_scale buffers registered during compression and the FSDPModule are reharded
once compression is complete.
Args:
fsdp_module (nn.Module): The FSDP-wrapped module to compress.
Returns:
None
"""
# Unshard FSDPmodule by temporarily setting _fake_quant to prevent weight compression from being triggered
with enable_fake_quant(fsdp_module):
fsdp_module.unshard()
# Get the FSDPParamGroup for the FSDPModule
fsdp_param_group = fully_shard.state(fsdp_module)._fsdp_param_group
if getattr(fsdp_param_group, "fsdp_params", None) is None:
warnings.warn(
f"FSDPParamGroup for {fsdp_module} has no fsdp_params, skipping compression"
)
return
# Create FSDPParam mapping dictionary to keep track of FSDPParams to update/delete
fsdp_param_mapping = _create_fsdp_param_mapping(fsdp_param_group.fsdp_params, fsdp_module)
for name, submodule in fsdp_module.named_modules():
# This is to handle case where the root FSDPModule has parameters.
# We skip all the parameters that dont belong to the FSDPParamGroup.
if name not in fsdp_param_mapping:
continue
if _compress_and_update_module_weight(submodule):
old_fsdp_param = fsdp_param_mapping[name]
# Update mp policy to reflect the new dtype
new_mp_policy = MixedPrecisionPolicy(
param_dtype=submodule.weight.dtype,
reduce_dtype=None,
output_dtype=None,
cast_forward_inputs=False,
)
with no_requires_grad():
# Create a new QFSDPParam parameter
new_param = QFSDPParam(
submodule.weight,
old_fsdp_param._module_info,
old_fsdp_param.mesh_info,
old_fsdp_param.post_forward_mesh_info,
old_fsdp_param.device,
None,
new_mp_policy,
None,
)
# Update the FSDPParam mapping to keep track of the new FSDPParam
fsdp_param_mapping[name] = new_param
# Remove the post_load_hook_handle to allow gc to collect the old FSDPParam
old_fsdp_param._post_load_hook_handle.remove()
# Update FSDPParam list with new compressed weights
fsdp_param_group.fsdp_params = list(fsdp_param_mapping.values())
# Reshard FSDP root module
fsdp_module.reshard()
with SequentialQuantizer.convert_to_single_quantizer(module), torch.no_grad():
for _, m in module.named_modules():
if hasattr(m, "weight") and (m.weight is None or m.weight.is_meta):
continue
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
real_quant_tensor = m.weight_quantizer(m.weight)
m.weight = QTensorWrapper(real_quant_tensor)
# If FSDP module, we need to additionally process the FSDPParam list
if isinstance(m, FSDPModule):
_compress_fsdp_module(m)
else:
# Compress weights and update module weight
_compress_and_update_module_weight(m)
modelopt/torch/quantization/triton/fp4_kernel.py
+10 -9
View File
@@ -54,7 +54,7 @@ def fp4_fake_quant_kernel(
pid_n = tl.program_id(axis=1)
# Load global scale from tensor
global_scale = tl.load(global_scale_ptr)
global_scale = tl.load(global_scale_ptr).to(tl.float32)
# Calculate offsets
offs_m = pid_m * TILE_SIZE + tl.arange(0, TILE_SIZE)
@@ -67,12 +67,13 @@ def fp4_fake_quant_kernel(
# Reshape for block processing
x_reshaped = tl.reshape(x, (TILE_SIZE, NUM_FP4_BLOCKS, BLOCK_SIZE))
x_abs = tl.abs(x_reshaped)
# Calculate max values for each FP4 block
block_max = tl.max(tl.abs(x_reshaped), axis=2, keep_dims=True)
block_max = tl.max(x_abs, axis=2, keep_dims=True)
# global_scale = global_amax / (448 * 6)
block_max_quant = (
tl.clamp((block_max / (6.0 * global_scale)), -448.0, 448.0).to(tl.float8e4nv).to(tl.float32)
tl.minimum((block_max / (6.0 * global_scale)), 448.0).to(tl.float8e4nv).to(tl.float32)
* global_scale
)
@@ -80,11 +81,13 @@ def fp4_fake_quant_kernel(
block_max_quant_broadcast = tl.broadcast_to(
block_max_quant, (TILE_SIZE, NUM_FP4_BLOCKS, BLOCK_SIZE)
)
x_scaled = x_reshaped / block_max_quant_broadcast
# Set scale to 1 if block amax is 0
block_max_quant_broadcast = tl.where(
block_max_quant_broadcast < 1e-5, 1.0, block_max_quant_broadcast
)
abs_scaled = x_abs / block_max_quant_broadcast
# Quantize to FP4 values: {0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6}, following round to even
abs_scaled = tl.abs(x_scaled)
q_val = tl.where(
abs_scaled <= 0.25,
0.0,
@@ -108,10 +111,8 @@ def fp4_fake_quant_kernel(
)
# Apply signs and rescale
sign = tl.where(x_scaled >= 0, 1.0, -1.0)
x_rescaled = q_val * block_max_quant_broadcast
x_rescaled = x_rescaled * sign
x_rescaled = tl.where(x_reshaped >= 0, x_rescaled, -x_rescaled)
# Reshape back and store
x_rescaled = tl.reshape(x_rescaled, (TILE_SIZE, TILE_SIZE))
modelopt/torch/quantization/utils.py
+93 -17
View File
@@ -15,10 +15,17 @@
"""Quantization utilities."""
from collections import namedtuple
from collections.abc import Generator
from contextlib import ExitStack, contextmanager, nullcontext
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributed.fsdp import FSDPModule
from torch.distributed.tensor import Replicate
from modelopt.torch.utils import print_rank_0
__all__ = [
"EXPORT_MODE",
@@ -26,11 +33,11 @@ __all__ = [
"export_torch_mode",
"is_quantized",
"is_quantized_column_parallel_linear",
"is_quantized_layer_with_weight",
"is_quantized_linear",
"is_quantized_row_parallel_linear",
"reduce_amax",
"replace_function",
"weight_attr_names",
]
@@ -173,6 +180,62 @@ def reduce_amax(input, axis=None, keepdims=True, squeeze_scalar=True):
return output
def weight_attr_names(module: nn.Module) -> Generator[str, None, None]:
"""Get the weight param attribute names in a converted module, non-recursive.
We consider the following two cases for each weight param attribute:
- The standard weight attribute (e.g. nn.Linear).
- The custom `weight_attr_name`. (e.g. Llama4TextExperts has weight attributes `gate_up_proj` and `down_proj`)
"""
from .nn import SequentialQuantizer, TensorQuantizer
# the standard weight and quantizer case
weight = getattr(module, "weight", None)
weight_quantizer = getattr(module, "weight_quantizer", None)
if isinstance(weight, nn.Parameter) and isinstance(
weight_quantizer, (TensorQuantizer, SequentialQuantizer)
):
yield "weight"
# other weight and quantizer case
for name, _ in module.named_parameters(recurse=False):
weight = getattr(module, name, None)
weight_quantizer = getattr(module, f"{name}_weight_quantizer", None)
if isinstance(weight, nn.Parameter) and isinstance(
weight_quantizer, (TensorQuantizer, SequentialQuantizer)
):
yield name
"""The whole set of quantizer related attribute names for a given weight name."""
QuantizerAttrNames = namedtuple(
"QuantizerAttrNames",
(
"weight_quantizer",
"input_quantizer",
"output_quantizer",
"weight_scale",
"weight_scale_2",
"input_scale",
"output_scale",
),
)
def quantizer_attr_names(weight_name: str = "weight") -> QuantizerAttrNames:
"""Get all the quantizer related attribute names for a given weight name."""
prefix = f"{weight_name}_" if weight_name != "weight" else ""
return QuantizerAttrNames(
weight_quantizer=f"{prefix}weight_quantizer",
input_quantizer=f"{prefix}input_quantizer",
output_quantizer=f"{prefix}output_quantizer",
weight_scale=f"{prefix}weight_scale",
weight_scale_2=f"{prefix}weight_scale_2",
input_scale=f"{prefix}input_scale",
output_scale=f"{prefix}output_scale",
)
def is_quantized(module):
"""Check if a module is quantized."""
from .nn import TensorQuantizer
@@ -180,11 +243,6 @@ def is_quantized(module):
return any(isinstance(_module, TensorQuantizer) for _module in module.modules())
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."""
from .nn import QuantModule, TensorQuantizer
@@ -213,6 +271,31 @@ def is_quantized_parallel_linear(module):
return is_quantized_column_parallel_linear(module) or is_quantized_row_parallel_linear(module)
@contextmanager
def calibrate_with_adapters(model, args):
"""Disables LoRA adapters during calibration, then re-enables them afterward."""
is_lora = getattr(args, "lora", None)
if is_lora:
print_rank_0("Disabling LoRA adapters during calibration...")
model.disable_adapters()
yield
if is_lora:
print_rank_0("Enabling LoRA adapters after calibration...")
model.enable_adapters()
def disable_lora_quantizers_in_config(config, layers):
"""Turns off input, weight, and output quantizers for LoRA weights and LoRALinear layers in config."""
config["quant_cfg"]["*lora*"] = {"enable": False}
for layer in layers:
config["quant_cfg"][f"*{layer}.input_quantizer"] = {"enable": False}
config["quant_cfg"][f"*{layer}.weight_quantizer"] = {"enable": False}
config["quant_cfg"][f"*{layer}.output_quantizer"] = {"enable": False}
return config
@contextmanager
def replace_function(package, name, new_func):
"""Replace a function with a new one within a context."""
@@ -256,7 +339,7 @@ def is_pow2(n):
return (n != 0) and (n & (n - 1) == 0)
def _get_fsdp2_mesh(module: torch.nn.Module):
def _get_fsdp2_mesh(module: nn.Module):
"""Get the mesh info of the model."""
try:
from torch.distributed._composable_state import _get_module_state
@@ -271,13 +354,8 @@ def _get_fsdp2_mesh(module: torch.nn.Module):
return fsdp_state._fsdp_param_group.post_forward_mesh_info.mesh
def _get_enclosing_fsdp_module(module: torch.nn.Module, root_model: torch.nn.Module):
def _get_enclosing_fsdp_module(module: nn.Module, root_model: nn.Module):
"""Get the enclosing FSDP module for a given module."""
try:
from torch.distributed.fsdp import FSDPModule
except ImportError:
return None
if isinstance(module, FSDPModule):
return module
@@ -300,14 +378,12 @@ def _get_enclosing_fsdp_module(module: torch.nn.Module, root_model: torch.nn.Mod
@contextmanager
def fsdp2_weight_access_and_writeback_context(module: torch.nn.Module, root_model: torch.nn.Module):
def fsdp2_weight_access_and_writeback_context(module: nn.Module, root_model: nn.Module):
"""Context manager for FSDP2 weight access and writeback.
Note this context will gather the weight across FSDP/HSDP shards. If TP is implemented with DTensor,
the weight will be a local tensor of the TP DTensor under this context.
"""
from torch.distributed.tensor import Replicate
assert isinstance(root_model, torch.distributed.fsdp.FSDPModule), "We only support FSDP2"
assert not hasattr(module, "_hf_hook"), "We dont support FSDP2 with HF accelerate hooks"
@@ -331,7 +407,7 @@ def fsdp2_weight_access_and_writeback_context(module: torch.nn.Module, root_mode
placements=[Replicate()] * fsdp_dim + list(original_placements[fsdp_dim:]),
device_mesh=original_device_mesh,
)
new_weight = torch.nn.Parameter(weight_collected.to_local())
new_weight = nn.Parameter(weight_collected.to_local())
module._parameters["weight"] = new_weight
yield
modelopt/torch/speculative/plugins/megatron_eagle.py
+16 -14
View File
@@ -434,24 +434,26 @@ class EagleModule(MegatronModule):
# This linear was previously a ColumnParallelLinear. We changed it to a normal linear
# since ColumnParallelLinear will have try to gather the input sequence when sequence
# parallel is used and does not allow gathering the outputs.
self.fc = Linear(
eagle_config.hidden_size * fc_input_size_multiplier,
eagle_config.hidden_size,
config=eagle_config,
init_method=(lambda w: None), # not used
bias=bias,
).to(device)
with torch.device(device):
self.fc = Linear(
eagle_config.hidden_size * fc_input_size_multiplier,
eagle_config.hidden_size,
config=eagle_config,
init_method=(lambda w: None), # not used
bias=bias,
)
self.rotary_pos_emb = rotary_pos_emb
# Eagle does not use the final_layernorm in decoder.
self.decoder = EagleTransformerBlock(
config=eagle_config,
spec=eagle_transformer_layer_spec,
post_layer_norm=use_last_layernorm,
pre_process=True,
post_process=True,
).to(device)
with torch.device(device):
self.decoder = EagleTransformerBlock(
config=eagle_config,
spec=eagle_transformer_layer_spec,
post_layer_norm=use_last_layernorm,
pre_process=True,
post_process=True,
)
if self._num_aux_hidden_states > 0:
layer = self.decoder.layers[0]
modelopt/torch/speculative/plugins/transformers.py
+764 -147
View File
File diff suppressed because it is too large Load Diff
modelopt/torch/speculative/utils.py
+5
View File
@@ -34,6 +34,9 @@ def calibrate_frequent_vocab(tokenizer, text, target_vocab_size, output_file=Non
counter = Counter(conversations)
vocab = counter.most_common(target_vocab_size)
mapping = torch.zeros(target_vocab_size, dtype=torch.int64)
assert len(vocab) == target_vocab_size, (
f"Not enough vocabs to calibrate ({len(vocab)}/{target_vocab_size}). Please increase data size."
)
for i in range(target_vocab_size):
idx = vocab[i][0]
mapping[i] = idx - i
@@ -292,6 +295,8 @@ class AcceptanceRateValidation:
Use rank 0 data as the golden set to broadcast to all ranks.
Each rank will then compare to this data and through error if different.
"""
if not torch.distributed.is_initialized():
return data
if data is None:
return
golden_set = copy.deepcopy(data)
modelopt/torch/utils/dataset_utils.py
+82 -41
View File
@@ -31,30 +31,53 @@ if TYPE_CHECKING:
# Use dict to store the config for each dataset.
# If we want to export more options to user like target languages, we need more standardized approach like dataclass.
SUPPORTED_DATASET_CONFIG: dict[str, Any] = {
"open_code_reasoning": {
"config": {"path": "nvidia/OpenCodeReasoning", "name": "split_0", "split": ["split_0"]},
"preprocess": lambda sample: "\n".join([sample["input"], sample["output"]]),
},
"open_math_reasoning": {
"config": {
"path": "nvidia/OpenMathReasoning",
"split": ["cot", "tir", "genselect"],
},
"preprocess": lambda sample: "\n".join([sample["problem"], sample["generated_solution"]]),
},
"llama-nemotron-post-training-dataset": {
"config": {
"path": "nvidia/Llama-Nemotron-Post-Training-Dataset",
"name": "SFT",
"split": ["code", "math", "science", "chat", "safety"],
},
"preprocess": lambda sample: "\n".join(turn["content"] for turn in sample["input"])
+ "\n"
+ sample["output"],
},
"magpie": {
"config": {"path": "Magpie-Align/Magpie-Pro-MT-300K-v0.1"},
"target": "conversations",
"preprocess": lambda sample: "\n".join(turn["value"] for turn in sample),
"config": {
"path": "Magpie-Align/Magpie-Pro-MT-300K-v0.1",
"split": ["train"],
},
"preprocess": lambda sample: "\n".join(turn["value"] for turn in sample["conversations"]),
},
"cnn_dailymail": {
"config": {"path": "cnn_dailymail", "name": "3.0.0"},
"target": "article",
"config": {"path": "cnn_dailymail", "name": "3.0.0", "split": ["train"]},
"preprocess": lambda sample: sample["article"],
},
"pile": {
"config": {"path": "monology/pile-uncopyrighted"},
"target": "text",
"config": {"path": "monology/pile-uncopyrighted", "name": "v1.0", "split": ["train"]},
"preprocess": lambda sample: sample["text"],
},
"pg19": {
"config": {"path": "pg19"},
"target": "text",
"config": {"path": "pg19", "name": "v1.0", "split": ["train"]},
"preprocess": lambda sample: sample["text"],
},
"wikipedia": {
"config": {"path": "wikipedia", "name": "20220301.en"},
"target": "text",
"config": {"path": "wikipedia", "name": "20220301.en", "split": ["train"]},
"preprocess": lambda sample: sample["text"],
},
"c4": {
"config": {"path": "c4", "name": "en"},
"target": "text",
"config": {"path": "c4", "name": "en", "split": ["train"]},
"preprocess": lambda sample: sample["text"],
},
}
@@ -77,36 +100,41 @@ def _get_dataset_samples(dataset_name: str, num_samples: int) -> list[str]:
Samples: The list of samples.
"""
# Load the dataset
if dataset_name in SUPPORTED_DATASET_CONFIG:
from datasets import load_dataset
dataset_config = SUPPORTED_DATASET_CONFIG[dataset_name]
dataset = load_dataset(
split="train",
streaming=True,
**dataset_config["config"],
)
else:
if dataset_name not in SUPPORTED_DATASET_CONFIG:
raise NotImplementedError(
f"dataset {dataset_name} is not supported. Please use one of the following:"
f" {get_supported_datasets()}."
)
# Access only the required samples
from datasets import load_dataset
dataset_config = SUPPORTED_DATASET_CONFIG[dataset_name]
# It's unfortunate that the load_dataset function does not support split a list while streaming.
# So we need to load the dataset for each split.
config = dataset_config["config"].copy()
splits = config.pop("split", [None])
dataset_splits = [
load_dataset(
streaming=True,
**config,
split=split,
)
for split in splits
]
# Split the samples evenly across the splits
# For streaming datasets, there is no reliable way to get the number of samples in each split
# other than loading the entire dataset. So, we just use the same number of samples for each split.
num_samples_splits = [num_samples // len(dataset_splits) for _ in dataset_splits]
num_samples_splits[-1] += num_samples - sum(num_samples_splits)
samples = []
target_key = dataset_config["target"]
for i, sample in enumerate(dataset):
if i >= num_samples:
break
for dataset, num_samples_split in zip(dataset_splits, num_samples_splits):
for i, sample in enumerate(dataset):
if i >= num_samples_split:
break
# Get raw value
value = sample[target_key]
# Apply preprocessing if defined
if "preprocess" in dataset_config:
value = dataset_config["preprocess"](value)
samples.append(value)
# Apply preprocess function to the sample
samples.append(dataset_config["preprocess"](sample))
return samples
@@ -127,10 +155,10 @@ class _CustomDataset(torch.utils.data.Dataset):
def get_dataset_dataloader(
dataset_name: str = "cnn_dailymail",
dataset_name: str | list[str] = "cnn_dailymail",
tokenizer: "PreTrainedTokenizerBase | None" = None,
batch_size: int = 1,
num_samples: int = 512,
num_samples: int | list[int] = 512,
max_sample_length: int = 512,
device: str | None = None,
include_labels: bool = False,
@@ -158,12 +186,25 @@ def get_dataset_dataloader(
"Tokenizer with the right padding_side may impact calibration accuracy. Recommend set to left"
)
num_samples = math.ceil(num_samples / batch_size) * batch_size
if isinstance(num_samples, int):
num_samples = [num_samples]
dataset = _get_dataset_samples(dataset_name, num_samples=num_samples)
if isinstance(dataset_name, str):
dataset_name = [dataset_name]
num_samples = [math.ceil(num_sample / batch_size) * batch_size for num_sample in num_samples]
assert len(dataset_name) == len(num_samples), (
"dataset_name and num_samples must be the same length"
)
all_samples = []
for ds_name, num_sample in zip(dataset_name, num_samples):
samples = _get_dataset_samples(ds_name, num_sample)
all_samples.extend(samples)
batch_encoded = tokenizer.batch_encode_plus(
dataset,
all_samples,
return_tensors="pt",
padding=True,
truncation=True,
modelopt/torch/utils/speech_dataset_utils.py
+3 -2
View File
@@ -25,7 +25,9 @@ from transformers import WhisperProcessor
# Use dict to store the config for each dataset.
# If we want to export more options to user like target languages, we need more standardized approach like dataclass.
SUPPORTED_SPEECH_DATASET_CONFIG: dict[str, dict[str, Any]] = {
"peoples_speech": {"config": {"path": "MLCommons/peoples_speech", "name": "clean"}},
"peoples_speech": {
"config": {"path": "MLCommons/peoples_speech", "name": "clean", "split": "train"},
},
}
__all__ = ["get_speech_dataset_dataloader", "get_supported_speech_datasets"]
@@ -47,7 +49,6 @@ def _get_speech_dataset(dataset_name: str, num_samples: int):
# Use streaming can reduce the downloading time for large datasets
dataset = load_dataset(
split="train",
**SUPPORTED_SPEECH_DATASET_CONFIG[dataset_name]["config"],
trust_remote_code=True,
streaming=True,
modelopt/torch/utils/tensor.py
+33 -1
View File
@@ -20,7 +20,13 @@ from collections import abc
import numpy as np
import torch
__all__ = ["numpy_to_torch", "torch_detach", "torch_to", "torch_to_numpy"]
__all__ = [
"numpy_to_torch",
"to_empty_if_meta_device",
"torch_detach",
"torch_to",
"torch_to_numpy",
]
def torch_to(data, *args, **kwargs):
@@ -53,3 +59,29 @@ def torch_to_numpy(inputs: list[torch.Tensor]) -> list[np.ndarray]:
def numpy_to_torch(np_outputs: list[np.ndarray]) -> list[torch.Tensor]:
"""Convert numpy arrays to torch tensors."""
return [torch.from_numpy(arr) for arr in np_outputs]
def to_empty_if_meta_device(module: torch.nn.Module, *, device: torch.device, recurse=True):
"""Move tensors to device if not meta device; otherwise materialize with empty_like().
Officially, torch suggests to_empty() for meta device materialization. Under the hood,
torch.empty_like() is applied to all parameters or buffers (see _apply). This may
accidently overwrite buffers with precomputed values during construction. Given the
goal is to only materialize those tensors on meta device, this function checks the
device first and only move the tensor to the destination if it is not on meta device.
Args:
module: The target module to apply this transformation.
device: The desired device of the parameters
and buffers in this module.
recurse: Whether parameters and buffers of submodules should
be recursively moved to the specified device.
"""
def _empty_like_if_meta(tensor: torch.Tensor, *, device: torch.device):
if tensor.device == torch.device("meta"):
return torch.empty_like(tensor, device=device)
else:
return tensor.to(device)
return module._apply(lambda t: _empty_like_if_meta(t, device=device), recurse=recurse)
modelopt/torch/utils/vlm_dataset_utils.py
+1 -2
View File
@@ -25,7 +25,7 @@ from .image_processor import MllamaImageProcessor
# Use dict to store the config for each dataset.
# If we want to export more options to user like target languages, we need more standardized approach like dataclass.
SUPPORTED_VLM_DATASET_CONFIG: dict[str, dict[str, Any]] = {
"scienceqa": {"config": {"path": "derek-thomas/ScienceQA"}},
"scienceqa": {"config": {"path": "derek-thomas/ScienceQA", "split": "train"}},
}
__all__ = ["get_supported_vlm_datasets", "get_vlm_dataset_dataloader"]
@@ -47,7 +47,6 @@ def _get_vlm_dataset(dataset_name: str, num_samples: int):
# Use streaming can reduce the downloading time for large datasets
dataset = load_dataset(
split="train",
**SUPPORTED_VLM_DATASET_CONFIG[dataset_name]["config"],
)
else:
setup.py
+2 -1
View File
@@ -43,7 +43,7 @@ required_deps = [
"pulp",
"regex",
"safetensors",
"torch>=2.5",
"torch>=2.6",
"torchprofile>=0.0.4",
"torchvision",
]
@@ -70,6 +70,7 @@ optional_deps = {
"huggingface_hub>=0.24.0",
"peft>=0.12.0",
"transformers>=4.48,<5.0", # Version match done in modelopt/torch/__init__.py as well
"deepspeed>=0.9.6 ; platform_system != 'Windows'",
],
# linter tools
"dev-lint": [
tests/_test_utils/torch_dist/fsdp_test.py
+1 -6
View File
@@ -19,13 +19,8 @@ from typing import Any
import torch
import torch.nn as nn
from packaging.version import Version
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP # noqa: N817
if Version(torch.__version__) >= Version("2.6"):
from torch.distributed.fsdp import fully_shard
else:
from torch.distributed._composable.fsdp.fully_shard import fully_shard
from torch.distributed.fsdp import fully_shard
from modelopt.torch.opt.dynamic import DynamicModule, _pytorch_managed
tests/_test_utils/torch_dist/plugins/megatron_common.py
+34 -17
View File
@@ -52,6 +52,7 @@ from modelopt.torch.opt.plugins.mcore_dist_checkpointing import (
restore_sharded_modelopt_state,
save_sharded_modelopt_state,
)
from modelopt.torch.utils import to_empty_if_meta_device
try:
from megatron.core.extensions.transformer_engine import TENorm
@@ -134,6 +135,8 @@ def get_mcore_gpt_model(
initialize_megatron: bool = False,
*,
num_layers: int = 2,
num_layers_in_first_pipeline_stage: int | None = None,
num_layers_in_last_pipeline_stage: int | None = None,
hidden_size: int = 64,
num_attention_heads: int = 8,
num_query_groups: int | None = None,
@@ -143,9 +146,8 @@ def get_mcore_gpt_model(
activation_func: str = "swiglu",
normalization: str = "LayerNorm",
transformer_impl: str = "modelopt" if HAS_TE else "local",
# Uneven PP
num_layers_in_first_pipeline_stage: int | None = None,
num_layers_in_last_pipeline_stage: int | None = None,
use_cpu_initialization: bool = False,
bf16: bool = True,
) -> GPTModel:
assert activation_func in ["swiglu", "squared_relu"]
assert normalization in ["LayerNorm", "RMSNorm"]
@@ -163,6 +165,8 @@ def get_mcore_gpt_model(
pipeline_model_parallel_size=pipeline_model_parallel_size,
sequence_parallel=False,
num_layers=num_layers,
num_layers_in_first_pipeline_stage=num_layers_in_first_pipeline_stage,
num_layers_in_last_pipeline_stage=num_layers_in_last_pipeline_stage,
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
num_query_groups=num_query_groups,
@@ -170,10 +174,10 @@ def get_mcore_gpt_model(
activation_func=squared_relu if activation_func == "squared_relu" else F.silu,
normalization=normalization,
gated_linear_unit=(activation_func == "swiglu"),
pipeline_dtype=torch.float32,
add_bias_linear=False,
num_layers_in_first_pipeline_stage=num_layers_in_first_pipeline_stage,
num_layers_in_last_pipeline_stage=num_layers_in_last_pipeline_stage,
use_cpu_initialization=use_cpu_initialization,
pipeline_dtype=torch.bfloat16 if bf16 else torch.float32,
bf16=bf16,
)
if transformer_impl == "local":
@@ -197,6 +201,8 @@ def get_mcore_gpt_model(
share_embeddings_and_output_weights=False,
position_embedding_type="rope",
)
if bf16:
model = model.to(torch.bfloat16)
return model
@@ -207,6 +213,8 @@ def get_mcore_mamba_model(
initialize_megatron: bool = False,
*,
num_layers: int = 3,
num_layers_in_first_pipeline_stage: int | None = None,
num_layers_in_last_pipeline_stage: int | None = None,
hybrid_override_pattern: str | None = None,
hidden_size: int = 64,
num_attention_heads: int = 8,
@@ -214,13 +222,11 @@ def get_mcore_mamba_model(
ffn_hidden_size: int | None = 128,
max_sequence_length: int = 4,
vocab_size: int = 64,
bf16: bool = True,
# Mamba-specific parameters
mamba_state_dim: int = 32,
mamba_head_dim: int = 16,
mamba_num_groups: int = 2,
# Uneven PP
num_layers_in_first_pipeline_stage: int | None = None,
num_layers_in_last_pipeline_stage: int | None = None,
) -> MambaModel:
assert HAS_MAMBA, "Mamba not installed"
@@ -232,16 +238,17 @@ def get_mcore_mamba_model(
pipeline_model_parallel_size=pipeline_model_parallel_size,
sequence_parallel=False,
num_layers=num_layers,
num_layers_in_first_pipeline_stage=num_layers_in_first_pipeline_stage,
num_layers_in_last_pipeline_stage=num_layers_in_last_pipeline_stage,
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
num_query_groups=num_query_groups,
ffn_hidden_size=ffn_hidden_size,
pipeline_dtype=torch.float32,
mamba_state_dim=mamba_state_dim,
mamba_head_dim=mamba_head_dim,
mamba_num_groups=mamba_num_groups,
num_layers_in_first_pipeline_stage=num_layers_in_first_pipeline_stage,
num_layers_in_last_pipeline_stage=num_layers_in_last_pipeline_stage,
pipeline_dtype=torch.bfloat16 if bf16 else torch.float32,
bf16=bf16,
)
if hybrid_override_pattern is None:
@@ -262,8 +269,10 @@ def get_mcore_mamba_model(
pre_process=is_pipeline_first_stage(),
post_process=is_pipeline_last_stage(),
share_embeddings_and_output_weights=False,
position_embedding_type="rope",
position_embedding_type="none",
)
if bf16:
model = model.to(torch.bfloat16)
return model
@@ -298,7 +307,7 @@ def run_mcore_inference(
hidden_size=active_hidden_size,
inference_batch_times_seqlen_threshold=batch_size * model.max_sequence_length,
fp32_residual_connection=False,
params_dtype=torch.float,
params_dtype=torch.bfloat16 if model.config.bf16 else torch.float32,
padded_vocab_size=model.vocab_size,
)
wrapped_model = GPTInferenceWrapper(model, inference_wrapper_config)
@@ -312,7 +321,7 @@ def run_mcore_inference(
logits = wrapped_model.run_one_forward_step(inference_input)
logits = broadcast_from_last_pipeline_stage(
[batch_size, model.max_sequence_length, model.vocab_size],
dtype=torch.float32,
dtype=torch.bfloat16 if model.config.bf16 else torch.float32,
tensor=logits,
)
return logits # shape: (batch_size, max_sequence_length, vocab_size)
@@ -353,7 +362,9 @@ def load_distributed_checkpoint(checkpoint_path, gpt_model):
return gpt_model
def sharded_state_dict_test_helper(tmp_path, model_ref, model_test, forward_fn, version=None):
def sharded_state_dict_test_helper(
tmp_path, model_ref, model_test, forward_fn, meta_device=False, version=None
):
logits_ref = forward_fn(model_ref)
state_dict = copy.deepcopy(model_ref.state_dict())
@@ -363,6 +374,8 @@ def sharded_state_dict_test_helper(tmp_path, model_ref, model_test, forward_fn,
# Restore model_test from `torch-dist`.
restore_sharded_modelopt_state([model_test], tmp_path)
if meta_device:
to_empty_if_meta_device(model_test, device="cuda")
model_test = load_distributed_checkpoint(tmp_path, model_test)
state_dict_test = model_test.state_dict()
@@ -392,4 +405,8 @@ def sharded_state_dict_test_helper(tmp_path, model_ref, model_test, forward_fn,
)
logits_test = forward_fn(model_test)
assert torch.allclose(logits_ref, logits_test), f"ref: {logits_ref}, test: {logits_test}"
logits_diff = (logits_test - logits_ref) / logits_ref
assert torch.allclose(logits_ref, logits_test), (
f"diff: {logits_diff.max()} ref: {logits_ref}, test: {logits_test}"
)
tests/_test_utils/torch_export/export_utils.py
+7 -6
View File
@@ -18,13 +18,14 @@ import torch
# Models
class ToyModel(torch.nn.Module):
def __init__(self):
def __init__(self, dims=[10, 10, 10, 10]):
super().__init__()
self.linears = torch.nn.Sequential(
torch.nn.Linear(10, 10),
torch.nn.Linear(10, 10),
torch.nn.Linear(10, 10),
)
assert len(dims) >= 2
if len(dims) == 2:
self.linears = torch.nn.Linear(dims[0], dims[1])
else:
linears = [torch.nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)]
self.linears = torch.nn.Sequential(*linears)
def forward(self, x):
return self.linears(x)
tests/_test_utils/torch_model/transformers_models.py
+6 -6
View File
@@ -29,7 +29,7 @@ from transformers import (
Qwen3Config,
Qwen3ForCausalLM,
T5Config,
T5Model,
T5ForConditionalGeneration,
T5Tokenizer,
)
@@ -68,7 +68,7 @@ def get_tiny_llama(**config_kwargs) -> LlamaForCausalLM:
return tiny_llama
def get_tiny_t5(**config_kwargs) -> T5Model:
def get_tiny_t5(**config_kwargs) -> T5ForConditionalGeneration:
kwargs = {
"vocab_size": 32,
"d_model": 32,
@@ -81,7 +81,7 @@ def get_tiny_t5(**config_kwargs) -> T5Model:
"decoder_start_token_id": 0,
}
kwargs.update(**config_kwargs)
t5_model = T5Model(T5Config(**kwargs))
t5_model = T5ForConditionalGeneration(T5Config(**kwargs))
return t5_model
@@ -138,10 +138,10 @@ def tf_output_tester(model_ref, model_test):
output_ref = model_ref(**inputs)
output_test = model_test(**inputs)
if hasattr(output_ref, "logits"):
assert torch.allclose(output_ref.logits, output_test.logits)
assert torch.allclose(output_ref.logits, output_test.logits, atol=1e-6)
else:
assert torch.allclose(output_ref.start_logits, output_test.start_logits)
assert torch.allclose(output_ref.end_logits, output_test.end_logits)
assert torch.allclose(output_ref.start_logits, output_test.start_logits, atol=1e-6)
assert torch.allclose(output_ref.end_logits, output_test.end_logits, atol=1e-6)
def tf_modelopt_state_and_output_tester(model_ref, model_test):
tests/_test_utils/torch_quantization/models.py
+17 -4
View File
@@ -44,18 +44,31 @@ class QuantConvLinear(nn.Module):
class SimpleLinear(nn.Module):
"""Test Linear model for ONNX export."""
def __init__(self):
def __init__(self, bias=True, dtype=torch.float32, add_linear=False):
super().__init__()
self.add_linear = add_linear
self.net = nn.Sequential(
nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 16)
nn.Linear(16, 32, bias=bias, dtype=dtype),
nn.ReLU(),
nn.Linear(32, 64, bias=bias, dtype=dtype),
nn.ReLU(),
nn.Linear(64, 16, bias=bias, dtype=dtype),
)
if add_linear:
self.linear1 = nn.Linear(16, 16, bias=bias, dtype=dtype)
def forward(self, x):
return self.net(x)
x = self.net(x)
if self.add_linear:
x = self.linear1(x)
return x
@classmethod
def get_input(cls):
return torch.randn(2, 16)
return torch.randn(
2,
16,
)
class SimpleConv(nn.Module):
tests/_test_utils/torch_quantization/quantize_common.py
+12 -9
View File
@@ -95,15 +95,18 @@ def save_restore_test(model_cls, device, quant_config, compress=False, version=N
# Rest of the tests are not needed for version < 0.29
return
# gpu: test restoring to a model on cpu. If the quantizer states are not initialized correctly,
# the buffers will be created on cuda and this test will fail
model_ref = model_cls().to("cpu")
state_dict = torch_to(state_dict, device="cuda" if torch.cuda.is_available() else "cpu")
mto.restore_from_modelopt_state(model_ref, state_dict)
model_ref.load_state_dict(model_quant.state_dict())
model_ref(calib_data[0].to("cpu")) # make sure all the buffers are created in the right device
model_ref.to(device)
assert torch.allclose(model_quant(calib_data[0]), model_ref(calib_data[0]))
if not compress:
# gpu: test restoring to a model on cpu. If the quantizer states are not initialized correctly,
# the buffers will be created on cuda and this test will fail
model_ref = model_cls().to("cpu")
state_dict = torch_to(state_dict, device="cuda" if torch.cuda.is_available() else "cpu")
mto.restore_from_modelopt_state(model_ref, state_dict)
model_ref.load_state_dict(model_quant.state_dict())
model_ref(
calib_data[0].to("cpu")
) # make sure all the buffers are created in the right device
model_ref.to(device)
assert torch.allclose(model_quant(calib_data[0]), model_ref(calib_data[0]))
# Test that smoothquant is restored correctly
if quant_config == mtq.INT8_SMOOTHQUANT_CFG:
tests/gpu/torch/export/test_export_weight.py
+123
View File
@@ -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.
import math
import torch
import torch.nn as nn
from _test_utils.torch_export.export_utils import ToyModel, partial_w4a8_config
from torch.nn import functional as F
from torch.nn import init
import modelopt.torch.quantization as mtq
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
from modelopt.torch.quantization.nn.modules.quant_module import QuantModule, QuantModuleRegistry
from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
from modelopt.torch.quantization.tensor_quant import QUANT_DESC_8BIT_PER_TENSOR
from modelopt.torch.quantization.utils import quantizer_attr_names
class ToyLinear(nn.Module):
in_features: int
out_features: int
toyweight: torch.Tensor # intentionally not named weight
def __init__(
self,
in_features: int,
out_features: int,
device=None,
dtype=None,
) -> None:
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.toyweight = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs))
self.reset_parameters()
def reset_parameters(self) -> None:
init.kaiming_uniform_(self.toyweight, a=math.sqrt(5))
def forward(self, input: torch.Tensor) -> torch.Tensor:
return F.linear(input, self.toyweight)
def extra_repr(self) -> str:
return f"in_features={self.in_features}, out_features={self.out_features}"
class ToyModelLinear(torch.nn.Module):
def __init__(self, dims=[10, 10, 10, 10]):
super().__init__()
assert len(dims) >= 2
if len(dims) == 2:
self.linears = ToyLinear(dims[0], dims[1])
else:
linears = [ToyLinear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)]
self.linears = torch.nn.Sequential(*linears)
def forward(self, x):
return self.linears(x)
@QuantModuleRegistry.register({ToyLinear: "ToyLinear"})
class _ToyLinearQuant(QuantModule):
"""Base class for modules where the input is quantized."""
toyweight_input_quantizer: TensorQuantizer
toyweight_weight_quantizer: TensorQuantizer
toyweight_output_quantizer: TensorQuantizer
default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR
default_quant_desc_output = QUANT_DESC_8BIT_PER_TENSOR
default_quant_desc_weight = QUANT_DESC_8BIT_PER_TENSOR
def forward(self, input, *args, **kwargs):
"""Quantize the input before calling the original forward method."""
input = self.toyweight_input_quantizer(input)
weight = self.toyweight_weight_quantizer(self.toyweight)
output = F.linear(input, weight)
return self.toyweight_output_quantizer(output)
def _setup(self):
"""Patch the module's forward method to quantize the input."""
self._register_temp_attribute(
"toyweight_weight_quantizer", TensorQuantizer(self.default_quant_desc_weight)
)
self._register_temp_attribute(
"toyweight_input_quantizer", TensorQuantizer(self.default_quant_desc_input)
)
self._register_temp_attribute(
"toyweight_output_quantizer", TensorQuantizer(self.default_quant_desc_output)
)
self.toyweight_output_quantizer.disable()
def test_export_per_block_quantized_weight():
model = ToyModel(dims=[32, 256, 256, 32])
mtq.quantize(model, partial_w4a8_config, lambda x: x(torch.randn(1, 4, 32)))
quantizer_attrs = quantizer_attr_names("weight")
_export_quantized_weight(model.linears[2], torch.float32, "weight")
assert model.linears[2].weight.dtype == torch.uint8
assert hasattr(model.linears[2], quantizer_attrs.weight_quantizer)
assert hasattr(model.linears[2], quantizer_attrs.weight_scale)
assert hasattr(model.linears[2], quantizer_attrs.weight_scale_2)
assert hasattr(model.linears[2], quantizer_attrs.input_scale)
assert hasattr(model.linears[2], quantizer_attrs.input_quantizer)
assert hasattr(model.linears[2], quantizer_attrs.output_quantizer)
assert not getattr(model.linears[2], quantizer_attrs.output_quantizer).is_enabled
assert not hasattr(model.linears[2], quantizer_attrs.output_scale)
tests/gpu/torch/nas/plugins/test_megatron_gpt_dynamic_modules.py
+1
View File
@@ -169,6 +169,7 @@ def _test_gpt_parameter_sorting(activation_func, rank, size):
max_sequence_length=max_sequence_length,
vocab_size=vocab_size,
activation_func=activation_func,
bf16=False,
)
# Randomize layernorm weights instead of all zeros or ones
tests/gpu/torch/nas/plugins/test_megatron_mamba_dynamic_modules.py
+258
View File
@@ -0,0 +1,258 @@
# 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 pytest
import torch
from _test_utils.import_helper import skip_if_no_megatron
skip_if_no_megatron(apex_or_te_required=True, mamba_required=True)
from _test_utils.torch_dist.dist_utils import spawn_multiprocess_job
from _test_utils.torch_dist.plugins.megatron_common import (
get_mcore_mamba_model,
run_mcore_inference,
run_mcore_inference_with_dummy_input,
)
from _test_utils.torch_misc import set_seed
from megatron.core.parallel_state import is_pipeline_first_stage, is_pipeline_last_stage
from megatron.core.transformer.identity_op import IdentityOp
import modelopt.torch.nas as mtn
from modelopt.torch.nas.modules.conv import _DynamicConvNd
from modelopt.torch.nas.plugins.megatron import (
MambaDInnerHp,
MambaNumHeadsHp,
_DynamicColumnParallelLinear,
_DynamicExtendedRMSNorm,
_DynamicLayerNorm,
_DynamicMambaLayer,
_DynamicMambaMixer,
_DynamicMCoreLanguageModel,
_DynamicRowParallelLinear,
_DynamicVocabParallelEmbedding,
)
from modelopt.torch.nas.search_space import generate_search_space
from modelopt.torch.nas.traced_hp import TracedHp
from modelopt.torch.opt.utils import named_dynamic_modules, search_space_size
from modelopt.torch.utils import flatten_tree
from modelopt.torch.utils.random import centroid
SEED = 1234
def _test_mamba_search_space(rank, size):
channel_divisor = 64
mamba_num_heads_divisor = 4
mamba_head_dim_divisor = 4
num_layers = size
hybrid_override_pattern = "M" * size
hidden_size = 256
mamba_state_dim = 64
mamba_head_dim = 16
mamba_num_groups = 2
max_sequence_length = 16
vocab_size = 32
batch_size = 2
model = get_mcore_mamba_model(
tensor_model_parallel_size=1,
pipeline_model_parallel_size=size,
initialize_megatron=True,
num_layers=num_layers,
hybrid_override_pattern=hybrid_override_pattern,
hidden_size=hidden_size,
mamba_state_dim=mamba_state_dim,
mamba_head_dim=mamba_head_dim,
mamba_num_groups=mamba_num_groups,
max_sequence_length=max_sequence_length,
vocab_size=vocab_size,
)
mamba_num_heads = model.decoder.layers[0].mixer.nheads
model = mtn.convert(model, "mcore_gpt_minitron")
assert isinstance(model, _DynamicMCoreLanguageModel)
if is_pipeline_first_stage():
assert isinstance(model.embedding.word_embeddings, _DynamicVocabParallelEmbedding)
for layer in model.decoder.layers:
assert isinstance(layer, _DynamicMambaLayer)
assert isinstance(layer.mixer, _DynamicMambaMixer)
assert isinstance(layer.mixer.in_proj, _DynamicColumnParallelLinear)
assert isinstance(layer.mixer.out_proj, _DynamicRowParallelLinear)
assert isinstance(layer.mixer.conv1d, _DynamicConvNd)
if layer.mixer.rmsnorm:
assert isinstance(layer.mixer.norm, _DynamicExtendedRMSNorm)
if is_pipeline_last_stage():
assert isinstance(model.decoder.final_norm, _DynamicLayerNorm)
assert isinstance(model.output_layer, _DynamicColumnParallelLinear)
# NOTE: `search_space_size` does not reduce across TP/PP groups
ss_size_per_pp = search_space_size(model)
num_heads_choices = mamba_num_heads // mamba_num_heads_divisor
head_dim_choices = mamba_head_dim // mamba_head_dim_divisor
hidden_size_choices = hidden_size // channel_divisor
num_layers_per_pp = num_layers // size
assert (
ss_size_per_pp
== (num_heads_choices * head_dim_choices) ** num_layers_per_pp
* num_layers
* hidden_size_choices
)
# Make sure forward pass works on min and centroid subnets
prompt_tokens = torch.randint(0, vocab_size, (batch_size, max_sequence_length)).cuda()
for sample_func in [min, max, centroid]:
mtn.sample(model, sample_func)
output = run_mcore_inference(model, prompt_tokens)
assert output.shape == (batch_size, max_sequence_length, vocab_size)
# Make sure export and forward pass works on centroid model
mtn.export(model)
_ = run_mcore_inference(model, prompt_tokens, model.hidden_size)
assert not any(named_dynamic_modules(model))
def test_mamba_search_space():
spawn_multiprocess_job(
size=torch.cuda.device_count(), job=_test_mamba_search_space, backend="nccl"
)
def _test_mamba_parameter_sorting(rank, size):
num_layers = size
hybrid_override_pattern = "M" * size
hidden_size = 256
mamba_state_dim = 64
mamba_head_dim = 16
mamba_num_groups = 2
max_sequence_length = 32
vocab_size = 64
batch_size = 2
model = get_mcore_mamba_model(
tensor_model_parallel_size=1,
pipeline_model_parallel_size=size,
initialize_megatron=True,
num_layers=num_layers,
hybrid_override_pattern=hybrid_override_pattern,
hidden_size=hidden_size,
mamba_state_dim=mamba_state_dim,
mamba_head_dim=mamba_head_dim,
mamba_num_groups=mamba_num_groups,
max_sequence_length=max_sequence_length,
vocab_size=vocab_size,
bf16=False,
)
# Randomize norm weights instead of all zeros or ones
for n, m in model.named_modules():
if "norm" in n and not isinstance(m, IdentityOp):
m.weight.data = torch.randn_like(m.weight)
model.eval()
search_space = generate_search_space(model)
# Compute activations for sorting
for _ in range(5):
run_mcore_inference_with_dummy_input(model, batch_size)
# Get the output of the original model
prompt_tokens = torch.randint(0, vocab_size, (batch_size, max_sequence_length)).cuda()
y1 = run_mcore_inference(model, prompt_tokens)
search_space.sort_parameters()
# check if all mamba_num_heads, mamba_head_dim, hidden_size have been sorted
sortable_per_pp = [
n for n, hp in search_space.named_hparams(configurable=True) if hp.importance is not None
]
# 2 mamba hps per layer + 1 for hidden_size (num_layers is not sorted!)
assert len(sortable_per_pp) == 2 * num_layers // size + 1
# Export since sorting force reassigns SelfAttention weights which we dont want to re-sort!
# TODO: ideally we shouldn't need this
search_space.export()
# sanity check if the model functionality is preserved after sorting
y2 = run_mcore_inference(model, prompt_tokens)
# # check if the inference results after sorting is the same
if rank == 0:
for i, (t1, t2) in enumerate(zip(flatten_tree(y1)[0], flatten_tree(y2)[0])):
if not torch.allclose(t1, t2, rtol=1e-5, atol=1e-2):
print(f"Mismatch at index {i}")
print(f"{t1=}")
print(f"{t2=}")
diff = (t1 - t2).abs()
print(f"{diff=}")
print(f"{diff.max()=}")
print(f"{diff.min()=}")
print(f"{diff.mean()=}")
print(f"{diff.std()=}")
print(f"{diff.median()=}")
print(f"{diff.quantile(0.25)=}")
print(f"{diff.quantile(0.75)=}")
else:
print(f"Match at index {i}")
@pytest.mark.skip("Need to fix")
def test_mamba_parameter_sorting(need_2_gpus):
set_seed(SEED)
spawn_multiprocess_job(
size=torch.cuda.device_count(),
job=_test_mamba_parameter_sorting,
backend="nccl",
)
def test_mamba_num_heads_hp():
num_heads = MambaNumHeadsHp([2, 4, 6, 8], ngroups=2) # 4 heads per group
assert num_heads.choices == [2, 4, 6, 8]
assert num_heads.active_slice == slice(8)
num_heads.active = 4 # 2 heads per group
assert num_heads.active_slice.tolist() == [0, 1, 4, 5]
num_heads_ranking = torch.tensor([1, 0, 3, 2, 4, 7, 6, 5])
num_heads_ranking.argsort = lambda *args, **kwargs: num_heads_ranking
num_heads._get_importance = lambda: num_heads_ranking
num_heads.enforce_order(num_heads.importance.argsort(descending=True))
assert num_heads.active_slice.tolist() == [1, 0, 4, 7]
def test_mamba_d_inner_hp():
num_heads = TracedHp([2, 4, 6, 8])
head_dim = TracedHp([1, 2, 3])
d_inner = MambaDInnerHp(num_heads, head_dim)
assert d_inner.choices == [2, 4, 6, 8, 12, 16, 18, 24]
assert d_inner.active_slice == slice(24)
# Set importance and slice order
num_heads._get_importance = lambda: torch.tensor([2.2, 0.1, 1.1, 2.1, 3.0, 2.0, 0.0, 1.0])
head_dim._get_importance = lambda: torch.tensor([2.0, 3.0, 1.0])
num_heads.enforce_order(torch.argsort(num_heads.importance, descending=True))
head_dim.enforce_order(torch.argsort(head_dim.importance, descending=True))
assert num_heads.active_slice.tolist() == [4, 0, 3, 5, 2, 7, 1, 6]
assert head_dim.active_slice.tolist() == [1, 0, 2]
# check if we get correct selection of sorted + pruned heads after setting active values
num_heads.active = 6 # top 6 heads
head_dim.active = 2 # top 2 dims per head
assert d_inner.active == 12 # (6 * 2)
assert d_inner.active_slice.tolist() == [13, 12, 1, 0, 10, 9, 16, 15, 7, 6, 22, 21]
tests/gpu/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py
+114
View File
@@ -0,0 +1,114 @@
# 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 torch
from _test_utils.import_helper import skip_if_no_megatron
skip_if_no_megatron(apex_or_te_required=True, mamba_required=True)
from _test_utils.torch_dist.dist_utils import spawn_multiprocess_job
from _test_utils.torch_dist.plugins.megatron_common import (
get_mcore_mamba_model,
run_mcore_inference_with_dummy_input,
)
from megatron.core.ssm.mamba_layer import MambaLayer
import modelopt.torch.prune as mtp
def _test_mcore_mamba_pruning(rank, size):
num_layers = min(size * 2, 8)
hidden_size = 256
ffn_hidden_size = 128
num_attention_heads = 8
num_query_groups = 4
mamba_state_dim = 64
mamba_head_dim = 16
mamba_num_groups = 2
batch_size = 2
model = get_mcore_mamba_model(
tensor_model_parallel_size=1,
pipeline_model_parallel_size=size,
initialize_megatron=True,
num_layers=num_layers,
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
num_query_groups=num_query_groups,
mamba_state_dim=mamba_state_dim,
mamba_head_dim=mamba_head_dim,
mamba_num_groups=mamba_num_groups,
)
mamba_num_heads = torch.tensor(0, device=torch.cuda.current_device())
if rank == 0:
assert isinstance(model.decoder.layers[0], MambaLayer)
mamba_num_heads += model.decoder.layers[0].mixer.nheads
torch.distributed.broadcast(mamba_num_heads, 0, async_op=True)
mamba_num_heads = mamba_num_heads.item()
assert mamba_num_heads > 0, "No MambaLayer found in the model rank 0!"
def forward_loop(m):
for _ in range(5):
run_mcore_inference_with_dummy_input(m, batch_size, hidden_size)
# Traditional GPT pruning parameters
pruned_ffn_hidden_size = ffn_hidden_size // 2
pruned_num_attention_heads = num_attention_heads // 2
pruned_num_query_groups = num_query_groups // 2
pruned_hidden_size = hidden_size // 2
pruned_num_layers = num_layers // 2
# Mamba-specific pruning parameters
# pruned_mamba_num_heads = mamba_num_heads // 2
# pruned_mamba_head_dim = mamba_head_dim // 2
# Base export config with GPT/Attention parameters
# TODO: enable mamba head pruning after debugging
export_config = {
"ffn_hidden_size": pruned_ffn_hidden_size,
"num_attention_heads": pruned_num_attention_heads,
"num_query_groups": pruned_num_query_groups,
"hidden_size": pruned_hidden_size,
# "mamba_num_heads": pruned_mamba_num_heads,
# "mamba_head_dim": pruned_mamba_head_dim,
"num_layers": pruned_num_layers,
}
model, _ = mtp.prune(
model,
mode="mcore_gpt_minitron",
constraints={"export_config": export_config},
dummy_input=None, # Not used
config={"forward_loop": forward_loop},
)
# Assert forward pass works on the pruned model
run_mcore_inference_with_dummy_input(model, batch_size, pruned_hidden_size)
# Assert model.config is updated for correct save/restoring
assert model.config.ffn_hidden_size == pruned_ffn_hidden_size
assert model.config.num_attention_heads == pruned_num_attention_heads
assert model.config.num_query_groups == pruned_num_query_groups
assert model.config.hidden_size == pruned_hidden_size
assert model.config.num_layers == pruned_num_layers
# assert model.config.mamba_num_heads == pruned_mamba_num_heads
# assert model.config.mamba_head_dim == pruned_mamba_head_dim
def test_mcore_mamba_pruning():
spawn_multiprocess_job(
size=torch.cuda.device_count(), job=_test_mcore_mamba_pruning, backend="nccl"
)
tests/gpu/torch/quantization/backends/test_fp8_per_tensor_gemm.py
+28 -1
View File
@@ -20,9 +20,36 @@ from _test_utils.torch_quantization.quantize_common import compute_backward_grad
import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.backends import gemm_registry
from modelopt.torch.quantization.backends.fp8_per_tensor_gemm import Fp8PerTensorLinear
from modelopt.torch.quantization.backends.utils import fp8_compatible
@pytest.mark.skipif(not fp8_compatible(), reason="FP8 is not supported on this GPU")
@pytest.mark.parametrize("model_cls", [SimpleLinear])
@pytest.mark.parametrize("config", [mtq.FP8_DEFAULT_CFG])
def test_fp8_per_tensor_gemm_available(model_cls, config):
"""Test for fp8_per_tensor_gemm function with hardware-friendly dimensions."""
model = model_cls().cuda()
calib_data = [model.get_input().cuda() for _ in range(8)]
def forward_loop(model, run_backward=False):
for batch in calib_data:
output = model(batch)
if run_backward:
output.sum().backward()
mtq.quantize(model, config, forward_loop)
mtq.compress(model)
# Take the first module in the net
module = model.net[0]
input_tensor = calib_data[0].clone()
# Find the matching GEMM implementation
gemm_forward = gemm_registry.find_match(module, input_tensor, [], {})
assert gemm_forward == Fp8PerTensorLinear.apply
@pytest.mark.skipif(not fp8_compatible(), reason="FP8 is not supported on this GPU")
@pytest.mark.parametrize("model_cls", [SimpleLinear])
@pytest.mark.parametrize("config", [mtq.FP8_DEFAULT_CFG])
@@ -45,7 +72,7 @@ def test_fp8_per_tensor_gemm(model_cls, config):
expected = torch.nn.functional.linear(input_tensor, module.weight, bias=None)
# Find the matching GEMM implementation
gemm_forward = gemm_registry.find_match(module, input_tensor, [], {})
gemm_forward = Fp8PerTensorLinear.apply
assert gemm_forward is not None
# Test without bias
tests/gpu/torch/quantization/plugins/test_megatron.py
+45 -19
View File
@@ -118,22 +118,39 @@ def test_tensor_parallel(need_2_gpus, config):
)
def _gpt_model_provider(tp_size: int, hidden_size=256, vocab_size=64):
def _gpt_model_provider(tp_size: int, hidden_size=256, vocab_size=64, meta_device=False):
"""Build the model."""
gpt_model = get_mcore_gpt_model(
tensor_model_parallel_size=tp_size,
num_layers=4,
ffn_hidden_size=None,
num_attention_heads=4,
activation_func="squared_relu",
transformer_impl="local",
hidden_size=hidden_size,
vocab_size=vocab_size,
)
return gpt_model.cuda().eval()
if meta_device:
with torch.device("meta"):
gpt_model = get_mcore_gpt_model(
tensor_model_parallel_size=tp_size,
num_layers=4,
ffn_hidden_size=None,
num_attention_heads=4,
activation_func="squared_relu",
transformer_impl="local",
hidden_size=hidden_size,
vocab_size=vocab_size,
use_cpu_initialization=meta_device,
)
else:
gpt_model = get_mcore_gpt_model(
tensor_model_parallel_size=tp_size,
num_layers=4,
ffn_hidden_size=None,
num_attention_heads=4,
activation_func="squared_relu",
transformer_impl="local",
hidden_size=hidden_size,
vocab_size=vocab_size,
).cuda()
return gpt_model.eval()
def _test_sharded_state_dict(tmp_path, config, hidden_size, modelopt_version, compress, rank, size):
def _test_sharded_state_dict(
tmp_path, config, hidden_size, modelopt_version, compress, meta_device, rank, size
):
# Must disable output_layer quantization since output_layer amax cannot be restore via
# sharded_state_dict. All output_layer quantizers state are removed.
config["quant_cfg"]["*output_layer*"] = {"enable": False}
@@ -145,7 +162,8 @@ def _test_sharded_state_dict(tmp_path, config, hidden_size, modelopt_version, co
initialize_for_megatron(tensor_model_parallel_size=size, seed=SEED)
model_ref = _gpt_model_provider(size, hidden_size, vocab_size=256)
model_test = _gpt_model_provider(size, hidden_size, vocab_size=256)
model_test = _gpt_model_provider(size, hidden_size, vocab_size=256, meta_device=meta_device)
prompt_tokens = torch.randint(
0, model_ref.vocab_size, (2, model_ref.max_sequence_length)
).cuda()
@@ -162,7 +180,12 @@ def _test_sharded_state_dict(tmp_path, config, hidden_size, modelopt_version, co
delattr(module, "_amax_for_smoothing")
sharded_state_dict_test_helper(
tmp_path, model_ref, model_test, forward_fn, version=modelopt_version
tmp_path,
model_ref,
model_test,
forward_fn,
meta_device=meta_device,
version=modelopt_version,
)
if modelopt_version is not None:
@@ -211,13 +234,14 @@ mixed_block_size_config["quant_cfg"].update(
],
)
@pytest.mark.parametrize("compress", [False, True])
def test_homogeneous_sharded_state_dict(need_2_gpus, tmp_path, config, compress):
@pytest.mark.parametrize("meta_device", [False, True])
def test_homogeneous_sharded_state_dict(need_2_gpus, tmp_path, config, compress, meta_device):
if compress and config is mtq.W4A8_AWQ_BETA_CFG:
pytest.skip("W4A8_AWQ_BETA_CFG is not supported for compress")
spawn_multiprocess_job(
size=2,
job=partial(_test_sharded_state_dict, tmp_path, config, 256, None, compress),
job=partial(_test_sharded_state_dict, tmp_path, config, 256, None, compress, meta_device),
backend="nccl",
)
@@ -236,7 +260,7 @@ def test_homogeneous_sharded_state_dict(need_2_gpus, tmp_path, config, compress)
def test_heterogenous_sharded_state_dict(need_2_gpus, tmp_path, config):
spawn_multiprocess_job(
size=2,
job=partial(_test_sharded_state_dict, tmp_path, config, 256, None, False),
job=partial(_test_sharded_state_dict, tmp_path, config, 256, None, False, False),
backend="nccl",
)
@@ -256,7 +280,9 @@ def test_heterogenous_sharded_state_dict(need_2_gpus, tmp_path, config):
def test_sharded_state_dict_old_checkpoints(need_2_gpus, tmp_path, config, modelopt_version):
spawn_multiprocess_job(
size=2,
job=partial(_test_sharded_state_dict, tmp_path, config, 256, modelopt_version, False),
job=partial(
_test_sharded_state_dict, tmp_path, config, 256, modelopt_version, False, False
),
backend="nccl",
)
tests/gpu/torch/quantization/test_deepspeed.py
+193
View File
@@ -0,0 +1,193 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Test of quantization with DeepSpeed."""
import argparse
import copy
import os
from functools import partial
import pytest
pytest.importorskip("deepspeed")
pytest.importorskip("accelerate")
import deepspeed
import torch
import torch.nn as nn
from _test_utils.torch_dist.dist_utils import spawn_multiprocess_job, synchronize_state_dict
from accelerate import Accelerator
from accelerate.utils import DeepSpeedPlugin
import modelopt.torch.quantization as mtq
from modelopt.torch.opt.dynamic import _pytorch_managed
def get_ds_config(zero_stage: int = 3):
return {
"train_micro_batch_size_per_gpu": 1,
"gradient_accumulation_steps": 1,
"zero_optimization": {"stage": zero_stage}, # Restore Stage 3
"fp16": {"enabled": False},
"bf16": {"enabled": False},
}
def _test_deepspeed_simple_linear(zero_stage, rank, size):
deepspeed.init_distributed()
os.environ["LOCAL_RANK"] = str(rank)
os.environ["RANK"] = str(rank)
os.environ["WORLD_SIZE"] = str(size)
dim = 32
model = nn.Linear(dim, dim).cuda(rank)
inputs = torch.randn(2, 2, dim).cuda(rank)
synchronize_state_dict(model)
deepspeed_model_after = copy.deepcopy(model)
model = mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda model: model(inputs))
manager = model._get_dm_attribute_manager()
assert "weight" in manager.da_keys()
assert model._get_dm_attribute_manager().get_da_value("weight") is _pytorch_managed
out_ref = model(inputs)
# Create cmd_args namespace for DeepSpeed initialization
cmd_args = argparse.Namespace()
cmd_args.deepspeed_config = get_ds_config(zero_stage)
cmd_args.local_rank = rank
cmd_args.world_size = size
# Initialize DeepSpeed for the test model
optimizer_test = torch.optim.Adam(model.parameters(), lr=0.1)
deepspeed_model, _, _, _ = deepspeed.initialize(
args=cmd_args, model=model, optimizer=optimizer_test
)
assert "weight" in manager.da_keys()
out_test = deepspeed_model(inputs)
assert torch.allclose(out_ref, out_test)
# Test quantization after DeepSpeed initialization
optimizer_after_test = torch.optim.Adam(deepspeed_model_after.parameters(), lr=0.1)
accelerator = Accelerator(
deepspeed_plugin=DeepSpeedPlugin(hf_ds_config=get_ds_config(zero_stage))
)
deepspeed_model_after, _ = accelerator.prepare(deepspeed_model_after, optimizer_after_test)
deepspeed_unwrapped = accelerator.unwrap_model(deepspeed_model_after)
mtq.quantize(deepspeed_unwrapped, mtq.INT8_DEFAULT_CFG, lambda model: model(inputs))
out_deepspeed_model_after = deepspeed_model_after(inputs)
assert torch.allclose(out_ref, out_deepspeed_model_after)
def _test_nested_deepspeed_backward(zero_stage, rank, size, quant_cfg):
# Set required environment variables for DeepSpeed
os.environ["LOCAL_RANK"] = str(rank)
os.environ["RANK"] = str(rank)
os.environ["WORLD_SIZE"] = str(size)
dim = 32
torch.manual_seed(1)
model = nn.Sequential(
nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, dim)),
nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, dim)),
nn.Linear(dim, dim),
).cuda(rank)
inputs = torch.randn(2, 2, dim).cuda(rank)
inputss = inputs.detach().clone()
# test for quantization after DeepSpeed
deepspeed_model_quant_after = copy.deepcopy(model)
model = mtq.quantize(model, quant_cfg, lambda model: model(inputs))
deepspeed_model = copy.deepcopy(model)
optimizer_ref = torch.optim.Adam(model.parameters(), lr=0.1)
out_ref = model(inputs)
out_ref.sum().backward()
# Initialize DeepSpeed for the test model
cmd_args = argparse.Namespace()
cmd_args.deepspeed_config = get_ds_config(zero_stage)
cmd_args.local_rank = rank
cmd_args.world_size = size
# Create optimizer for DeepSpeed
optimizer_test = torch.optim.Adam(deepspeed_model.parameters(), lr=0.1)
deepspeed_model, optimizer_test, _, _ = deepspeed.initialize(
args=cmd_args, model=deepspeed_model, optimizer=optimizer_test
)
out_test = deepspeed_model(inputs)
deepspeed_model.backward(out_test.sum())
assert torch.allclose(out_ref, out_test)
optimizer_ref.step()
optimizer_ref.zero_grad()
optimizer_test.step()
optimizer_test.zero_grad()
out_ref_1 = model(inputss)
out_test_1 = deepspeed_model(inputss)
assert torch.allclose(out_ref_1, out_test_1, rtol=1e-4)
# Initialize DeepSpeed for quantization after DeepSpeed
optimizer_quant_after = torch.optim.Adam(deepspeed_model_quant_after.parameters(), lr=0.1)
accelerator = Accelerator(
deepspeed_plugin=DeepSpeedPlugin(hf_ds_config=get_ds_config(zero_stage))
)
deepspeed_model_quant_after, optimizer_quant_after = accelerator.prepare(
deepspeed_model_quant_after, optimizer_quant_after
)
deepspeed_unwrapped = accelerator.unwrap_model(deepspeed_model_quant_after)
mtq.quantize(deepspeed_unwrapped, quant_cfg, lambda model: model(inputs))
out_quant_after = deepspeed_model_quant_after(inputs)
accelerator.backward(out_quant_after.sum())
assert torch.allclose(out_ref, out_quant_after)
out_quant_after_1 = deepspeed_model_quant_after(inputss)
assert torch.allclose(out_ref_1, out_quant_after_1, rtol=1e-4)
@pytest.mark.parametrize("zero_stage", [1, 2, 3])
def test_deepspeed_simple_linear(zero_stage):
spawn_multiprocess_job(
size=torch.cuda.device_count(),
job=partial(_test_deepspeed_simple_linear, zero_stage),
backend="nccl",
)
@pytest.mark.parametrize("quant_cfg", [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG])
@pytest.mark.parametrize("zero_stage", [1, 2, 3])
def test_nested_deepspeed_backward(quant_cfg, zero_stage):
spawn_multiprocess_job(
size=torch.cuda.device_count(),
job=partial(_test_nested_deepspeed_backward, zero_stage, quant_cfg=quant_cfg),
backend="nccl",
)
tests/gpu/torch/quantization/test_hadamard.py
+8 -22
View File
@@ -39,7 +39,8 @@ def test_hadamard_transform(dim):
xxt = x @ x.T
x_h = normalized_hadamard_transform(x)
xxt_h = x_h @ x_h.T
assert torch.allclose(xxt_h, xxt, atol=1e-3)
# The numerical error can be large, especially for 16-bit floats.
assert torch.allclose(xxt_h, xxt, atol=0.05)
def test_kv_rotate():
@@ -59,33 +60,18 @@ def test_kv_rotate():
},
):
output_test = model(dummy_input)
assert torch.allclose(output_ref, output_test, atol=1e-3)
assert torch.allclose(output_ref, output_test, atol=0.05)
set_quantizer_by_cfg(
# Test the rotation is actually applied by turning on only one of the query, key quantizers
with set_quantizer_by_cfg_context(
model,
{
"*q_bmm_quantizer": {
"enable": False,
"rotate": False,
},
"*k_bmm_quantizer": {
"num_bits": 4,
"axis": -1,
"enable": True,
"rotate": False,
},
},
)
output_ref1 = model(dummy_input)
set_quantizer_by_cfg(
model,
{
"*[qk]_bmm_quantizer": {
"rotate": True,
},
},
)
output_test1 = model(dummy_input)
torch.not_equal(output_ref1, output_test1)
):
output_test1 = model(dummy_input)
assert not torch.allclose(output_ref, output_test1, atol=0.05)
mtq.unregister(SDPAAttention)
tests/gpu/torch/quantization/test_real_quantize_cuda.py
+65
View File
@@ -18,10 +18,14 @@
import fnmatch
import pytest
import torch
from _test_utils.torch_dist.dist_utils import get_device_counts, spawn_multiprocess_job
from _test_utils.torch_model.transformers_models import create_tiny_llama_dir
from _test_utils.torch_quantization.models import SimpleConv, SimpleConvLinear, SimpleLinear
from _test_utils.torch_quantization.quant_utils import get_model_size
from _test_utils.torch_quantization.quantize_common import save_restore_test
from torch.distributed.fsdp import FSDPModule, fully_shard
from torch.distributed.tensor import DTensor
import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights
@@ -202,3 +206,64 @@ def test_real_quantize_linear(quant_config, tmp_path):
and not module.weight_quantizer.fake_quant
):
assert isinstance(module.weight, QTensorWrapper)
def _test_mtq_compress_fsdp_module(
rank, size, model_cls=SimpleLinear, quant_config=mtq.NVFP4_DEFAULT_CFG
):
# Load model and shard it
model = model_cls(bias=False, dtype=torch.bfloat16, add_linear=True).cuda()
# Shard model
for n, m in model.named_modules():
if isinstance(m, torch.nn.Sequential):
fully_shard(m)
fully_shard(model)
# Create calib data
calib_data = [model.get_input().to(torch.bfloat16).cuda() for _ in range(8)]
# Forward loop
def forward_loop(model, run_backward=False):
for batch in calib_data:
output = model(batch)
if run_backward:
output.sum().backward()
# Calibrate model
mtq.quantize(model, quant_config, forward_loop)
# Compress model
mtq.compress(model)
# Verify that model is in sharded state after compression
for n, m in model.named_parameters():
assert isinstance(m, DTensor), f"Parameter {n} is not in sharded state after compression"
# Verify model unshard, module parameters must be torch.nn.Parameter or QTensorWrapper after unsharding
for n, m in model.named_modules():
if isinstance(m, FSDPModule):
m.unshard()
for n, m in model.named_parameters():
assert not isinstance(m, DTensor), (
f"Parameter {n} is not in unsharded state after unsharding"
)
# Verify model reshard, module parameters must be DTensors after reshard
for n, m in model.named_modules():
if isinstance(m, FSDPModule):
m.reshard()
for n, m in model.named_parameters():
assert isinstance(m, DTensor), (
f"Parameter {n} {m} is not in sharded state after calling reshard"
)
# Verify forward pass after compressing model
model(model.get_input().to(torch.bfloat16).cuda())
@pytest.mark.parametrize("device_count", get_device_counts())
def test_compress_fsdp_module(device_count):
spawn_multiprocess_job(size=device_count, job=_test_mtq_compress_fsdp_module, backend="nccl")
tests/gpu/torch/quantization/test_tensor_quant_cuda.py
+5 -3
View File
@@ -219,7 +219,7 @@ class Testfp4:
(test_out,) * (block_size // 8), dim=-1
)
def _test_fp4_kernel(test_in, test_out):
def _test_fp4_kernel(test_in, test_out, skip_triton=False):
inputs, expected_outputs = _get_test_inputs_outputs(test_in, test_out)
quantized_outputs = cuda_ext_mx.fused_amax_convert(
inputs,
@@ -229,7 +229,7 @@ class Testfp4:
inputs.abs().amax(),
)
assert torch.allclose(quantized_outputs, expected_outputs)
if triton_kernel.IS_AVAILABLE:
if triton_kernel.IS_AVAILABLE and not skip_triton:
quantized_outputs_triton = triton_kernel.fp4_fake_quant_block(
inputs, inputs.abs().amax()
)
@@ -242,7 +242,9 @@ class Testfp4:
# Test with e2m1 boundary values. The even indexes are rounded down and odd indexes are rounded up.
test_in = torch.tensor([[0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5, 6]]).cuda() * sign
test_out = torch.tensor([[0.0, 1, 1, 2, 2, 4, 4, 6]]).cuda() * sign
_test_fp4_kernel(test_in, test_out)
# The triton kernel has a numerical issue, the values are not exactly at the boundary after scaling,
# e.g. 0.25 -> 0.250061, this won't cause visible error for real-world quantizations.
_test_fp4_kernel(test_in, test_out, skip_triton=True)
# Test slightly below the e2m1 boundary values.
# Numbers should be quantized down to the corresponding e2m1 value.
tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py
+6
View File
@@ -72,6 +72,9 @@ def _test_speculative_gpt_model(
else:
raise ValueError("Only algo={eagle, medusa} are supported!")
# Bfloat16
model = model.to(torch.bfloat16)
# Prepare inputs for forward.
prompt_tokens = torch.randint(0, vocab_size, (batch_size, max_sequence_length)).cuda()
attention_mask = torch.tril(torch.ones((1, 1, max_sequence_length, max_sequence_length))).cuda()
@@ -189,6 +192,9 @@ def _test_tree_decode(tree_paths, greedy_steps, rank, size):
model = mtsp.convert(model, [("eagle", config)])
# Bfloat16
model = model.to(torch.bfloat16)
# Prepare inputs for forward.
prompt_tokens = torch.randint(0, vocab_size, (batch_size, max_sequence_length)).cuda()
attention_mask = torch.tril(torch.ones((1, 1, max_sequence_length, max_sequence_length))).cuda()
tests/unit/torch/export/test_export_weight.py
+98
View File
@@ -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.
import pytest
import torch
from _test_utils.torch_export.export_utils import ToyModel, partial_fp8_config, partial_w4a8_config
import modelopt.torch.quantization as mtq
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
from modelopt.torch.quantization.utils import quantizer_attr_names
@pytest.mark.parametrize(
"weight_name",
["weight", "weight_2", "some_other_w"],
)
def test_quantizer_attr_names(weight_name):
quantizer_attrs = quantizer_attr_names(weight_name)
if weight_name == "weight":
assert quantizer_attrs.weight_scale == "weight_scale"
assert quantizer_attrs.input_scale == "input_scale"
assert quantizer_attrs.weight_scale_2 == "weight_scale_2"
assert quantizer_attrs.weight_quantizer == "weight_quantizer"
assert quantizer_attrs.input_quantizer == "input_quantizer"
assert quantizer_attrs.output_quantizer == "output_quantizer"
assert quantizer_attrs.output_scale == "output_scale"
else:
assert quantizer_attrs.weight_scale == f"{weight_name}_weight_scale"
assert quantizer_attrs.input_scale == f"{weight_name}_input_scale"
assert quantizer_attrs.weight_scale_2 == f"{weight_name}_weight_scale_2"
assert quantizer_attrs.weight_quantizer == f"{weight_name}_weight_quantizer"
assert quantizer_attrs.input_quantizer == f"{weight_name}_input_quantizer"
assert quantizer_attrs.output_quantizer == f"{weight_name}_output_quantizer"
assert quantizer_attrs.output_scale == f"{weight_name}_output_scale"
def test_export_per_tensor_quantized_weight():
model = ToyModel(dims=[32, 256, 32, 128])
mtq.quantize(model, partial_fp8_config, lambda x: x(torch.randn(1, 4, 32)))
orig_dtype = model.linears[0].weight.dtype
quantizer_attrs = quantizer_attr_names("weight")
_export_quantized_weight(model.linears[0], torch.float32, "weight")
assert model.linears[0].weight.dtype == orig_dtype
assert hasattr(model.linears[0], quantizer_attrs.weight_quantizer)
assert not getattr(model.linears[0], quantizer_attrs.weight_quantizer).is_enabled
assert not hasattr(model.linears[0], quantizer_attrs.weight_scale)
assert not hasattr(model.linears[0], quantizer_attrs.weight_scale_2)
assert not hasattr(model.linears[0], quantizer_attrs.input_scale)
assert hasattr(model.linears[0], quantizer_attrs.input_quantizer)
assert not getattr(model.linears[0], quantizer_attrs.input_quantizer).is_enabled
assert hasattr(model.linears[0], quantizer_attrs.output_quantizer)
assert not getattr(model.linears[0], quantizer_attrs.output_quantizer).is_enabled
assert not hasattr(model.linears[0], quantizer_attrs.output_scale)
_export_quantized_weight(model.linears[1], torch.float32, "weight")
assert model.linears[1].weight.dtype == torch.float8_e4m3fn
assert hasattr(model.linears[1], quantizer_attrs.weight_quantizer)
assert hasattr(model.linears[1], quantizer_attrs.weight_scale)
assert not hasattr(model.linears[1], quantizer_attrs.weight_scale_2)
assert hasattr(model.linears[1], quantizer_attrs.input_quantizer)
assert hasattr(model.linears[1], quantizer_attrs.input_scale)
assert hasattr(model.linears[1], quantizer_attrs.output_quantizer)
assert not getattr(model.linears[1], quantizer_attrs.output_quantizer).is_enabled
assert not hasattr(model.linears[1], quantizer_attrs.output_scale)
def test_export_per_block_quantized_weight():
model = ToyModel(dims=[32, 256, 256, 32])
mtq.quantize(model, partial_w4a8_config, lambda x: x(torch.randn(1, 4, 32)))
quantizer_attrs = quantizer_attr_names("weight")
_export_quantized_weight(model.linears[2], torch.float32, "weight")
assert model.linears[2].weight.dtype == torch.uint8
assert hasattr(model.linears[2], quantizer_attrs.weight_quantizer)
assert hasattr(model.linears[2], quantizer_attrs.weight_scale)
assert hasattr(model.linears[2], quantizer_attrs.weight_scale_2)
assert hasattr(model.linears[2], quantizer_attrs.input_scale)
assert hasattr(model.linears[2], quantizer_attrs.input_quantizer)
assert hasattr(model.linears[2], quantizer_attrs.output_quantizer)
assert not getattr(model.linears[2], quantizer_attrs.output_quantizer).is_enabled
assert not hasattr(model.linears[2], quantizer_attrs.output_scale)
tests/unit/torch/export/test_get_quantization.py
+1 -35
View File
@@ -15,47 +15,13 @@
import pytest
import torch
from _test_utils.torch_export.export_utils import ToyModel, partial_fp8_config, partial_w4a8_config
import modelopt.torch.quantization as mtq
from modelopt.torch.export.layer_utils import get_quantization_format
from modelopt.torch.export.model_config import QUANTIZATION_FP8, QUANTIZATION_W4A8_AWQ
class ToyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linears = torch.nn.Sequential(
torch.nn.Linear(10, 10),
torch.nn.Linear(10, 10),
torch.nn.Linear(10, 10),
)
def forward(self, x):
return self.linears(x)
partial_fp8_config = {
"quant_cfg": {
"*.1.weight_quantizer": {"num_bits": (4, 3), "axis": None},
"*.1.input_quantizer": {"num_bits": (4, 3), "axis": None},
"default": {"num_bits": 8, "enable": False},
},
"algorithm": "max",
}
partial_w4a8_config = {
"quant_cfg": {
"*.2.weight_quantizer": [
{"num_bits": 4, "block_sizes": {-1: 128, "type": "static"}, "enable": True},
{"num_bits": (4, 3), "axis": None, "enable": True},
],
"*.2.input_quantizer": {"num_bits": (4, 3), "axis": None, "enable": True},
"default": {"num_bits": 8, "enable": False},
},
"algorithm": "awq_lite",
}
@pytest.mark.parametrize(
("config", "expected"),
[(partial_fp8_config, QUANTIZATION_FP8), (partial_w4a8_config, QUANTIZATION_W4A8_AWQ)],
tests/unit/torch/speculative/plugins/test_hf_speculative.py
+139
View File
@@ -15,8 +15,11 @@
import os
import pytest
import torch
from _test_utils.torch_model.transformers_models import (
create_tiny_llama_dir,
get_tiny_llama,
tf_modelopt_state_and_output_tester,
)
from transformers import AutoModelForCausalLM, LlamaForCausalLM
@@ -41,3 +44,139 @@ def test_medusa_model_convert_save_and_restore(tmp_path):
model_test = AutoModelForCausalLM.from_pretrained(tiny_llama_dir / "modelopt_model")
assert isinstance(model_test, mtsp.plugins.HFMedusaModel)
tf_modelopt_state_and_output_tester(model_ref, model_test)
def test_eagle_model_convert_save_and_restore(tmp_path):
model_ref = get_tiny_llama(num_hidden_layers=8)
config = {
"eagle_num_layers": 1,
"use_aux_hidden_state": True,
}
mtsp.convert(model_ref, mode=[("eagle", config)])
assert isinstance(model_ref, mtsp.plugins.HFEagleModel)
model_ref.save_pretrained(tmp_path / "modelopt_model")
assert os.path.exists(tmp_path / "modelopt_model/modelopt_state.pth")
model_test = AutoModelForCausalLM.from_pretrained(tmp_path / "modelopt_model")
assert isinstance(model_test, mtsp.plugins.HFEagleModel)
tf_modelopt_state_and_output_tester(model_ref, model_test)
# fmt: off
@pytest.mark.parametrize("dtype", [torch.bfloat16])
def test_eagle_model_prepare_eagle_inputs(dtype):
dummy_model = get_tiny_llama(num_hidden_layers=4)
config = {
"eagle_num_layers": 1,
"use_aux_hidden_state": True,
}
mtsp.convert(dummy_model, mode=[("eagle", config)])
eagle_input_ids_0 = torch.tensor([[10, 20, 30, 40]], dtype=torch.long)
position_ids_0 = torch.tensor([[0, 1, 2, 3]], dtype=torch.long)
#This is concatenated from 3 intermediate base model layers
cat_aux_hidden_states = torch.randn(1, 4, 32, dtype=dtype)
#This is eagle output from previous eagle forward pass
dummy_eagle_output_hidden_states = torch.randn(1, 4, 32, dtype=dtype)
#This is the causal mask for the 0th eagle step
m = torch.finfo(dtype).min
attention_mask_0 = torch.tensor([[0, m, m, m], # input tok 10-> predicting token 20
[0, 0, m, m], # 20 -> 30
[0, 0, 0, m], # 30 -> 40
[0, 0, 0, 0]] # 40 -> tok after 40
, dtype=dtype).view(1, 1, 4, 4)
# 2nd eagle step
eagle_input_h_1, eagle_input_ids_1, attention_mask_1, position_ids_1 = dummy_model._concat_eagle_inputs(
eagle_input_ids_0,
cat_aux_hidden_states,
attention_mask_0,
position_ids_0,
dummy_eagle_output_hidden_states,
)
assert eagle_input_ids_1.equal(torch.tensor([[10, 20, 30, 40, 10, 20, 30, 40]], dtype=torch.long))
assert position_ids_1.equal(torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3]], dtype=torch.long))
assert attention_mask_1.equal(torch.tensor([[0, m, m, m, m, m, m, m], # (x) output discarded
[0, 0, m, m, m, m, m, m], # (x)
[0, 0, 0, m, m, m, m, m], # (x)
[0, 0, 0, 0, m, m, m, m], # (x)
[m, m, m, m, m, m, m, m], # (x) input tok 10-> predicting token 20
[0, m, m, m, m, 0, m, m], # 20 -> 30
[0, 0, m, m, m, m, 0, m], # 30 -> 40
[0, 0, 0, 0, m, m, m, m], # (x) 40 -> tok after 40
], dtype=dtype).view(1, 1, 8, 8))
# 3rd eagle step
eagle_input_hidden_states_2, eagle_input_ids_2, attention_mask_2, position_ids_2 = dummy_model._concat_eagle_inputs(
eagle_input_ids_0,
cat_aux_hidden_states,
attention_mask_0,
position_ids_0,
torch.cat([dummy_eagle_output_hidden_states, dummy_eagle_output_hidden_states], dim=1),
)
assert eagle_input_ids_2.equal(torch.tensor([[10, 20, 30, 40, 10, 20, 30, 40, 10, 20, 30, 40]], dtype=torch.long))
assert position_ids_2.equal(torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]], dtype=torch.long))
assert attention_mask_2.equal(torch.tensor([[0, m, m, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, m, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, 0, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, 0, 0, m, m, m, m, m, m, m, m], # (x)
[m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[0, m, m, m, m, 0, m, m, m, m, m, m], # (x)
[0, 0, m, m, m, m, 0, m, m, m, m, m], # (x)
[0, 0, 0, 0, m, m, m, m, m, m, m, m], # (x)
[m, m, m, m, m, m, m, m, m, m, m, m], # (x)10 -> 20
[m, m, m, m, m, m, m, m, m, m, m, m], # (x)20 -> 30
[0, m, m, m, m, 0, m, m, m, m, 0, m], # 30 -> 40
[0, 0, 0, 0, m, m, m, m, m, m, m, m], # (x) 40 -> tok after 40
], dtype=dtype).view(1, 1, 12, 12))
# 4th eagle step
eagle_input_hidden_states_3, eagle_input_ids_3, attention_mask_3, position_ids_3 = dummy_model._concat_eagle_inputs(
eagle_input_ids_0,
cat_aux_hidden_states,
attention_mask_0,
position_ids_0,
torch.cat([dummy_eagle_output_hidden_states, dummy_eagle_output_hidden_states,
dummy_eagle_output_hidden_states],dim=1),
)
assert eagle_input_ids_3.equal(torch.tensor([[10, 20, 30, 40, 10, 20, 30, 40,
10, 20, 30, 40, 10, 20, 30, 40]], dtype=torch.long))
assert position_ids_3.equal(torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]], dtype=torch.long))
assert attention_mask_3.equal(torch.tensor([[0, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, m, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, 0, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, 0, 0, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[0, m, m, m, m, 0, m, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, m, m, m, m, 0, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, 0, 0, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[0, m, m, m, m, 0, m, m, m, m, 0, m, m, m, m, m], # (x)
[0, 0, 0, 0, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)10 -> 20
[m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)20 -> 30
[m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
[0, 0, 0, 0, m, m, m, m, m, m, m, m, m, m, m, m], # (x)
], dtype=dtype).view(1, 1, 16, 16))