test_huggingface.py

# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import warnings
from contextlib import nullcontext

import pytest
import torch
import torch.nn as nn
from _test_utils.torch.misc import set_seed
from _test_utils.torch.transformers_models import (
    create_tiny_llama_dir,
    get_tiny_gpt_oss,
    get_tiny_llama,
    get_tiny_qwen3_5,
    get_tiny_qwen3_5_moe,
    get_tiny_qwen3_moe,
    tf_modelopt_state_and_output_tester,
)
from packaging.version import Version

import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.nn import QuantLinear, QuantModuleRegistry
from modelopt.torch.quantization.plugins.huggingface import (
    get_homogeneous_hf_decoder_layers,
    is_homogeneous_hf_model,
)
from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector

pytest.importorskip("transformers")

import transformers
from transformers import AutoModelForCausalLM, LlamaForCausalLM
from transformers.models.dbrx.configuration_dbrx import DbrxConfig, DbrxFFNConfig
from transformers.models.dbrx.modeling_dbrx import DbrxExpertGLU, DbrxExperts, DbrxFFN


class HFModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            # initialization is (out_features, in_features) instead of (in_features, out_features)
            transformers.pytorch_utils.Conv1D(5, 3),
            nn.ReLU(),
            transformers.pytorch_utils.Conv1D(5, 5),
        )

    def forward(self, x):
        return self.net(x)


class PytorchModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            QuantLinear(3, 5),
            nn.ReLU(),
            QuantLinear(5, 5),
        )

    def forward(self, x):
        return self.net(x)


def test_convert_conv1d():
    set_seed()
    assert transformers.pytorch_utils.Conv1D in QuantModuleRegistry

    model_ref = HFModel()
    model_test = HFModel()
    model_test.load_state_dict(model_ref.state_dict())

    mtq.replace_quant_module(model_test)
    for name, module in model_test.named_modules():
        if isinstance(module, transformers.pytorch_utils.Conv1D):
            assert hasattr(module, "input_quantizer")
            assert hasattr(module, "weight_quantizer")
            assert hasattr(module, "output_quantizer")

    mtq.set_quantizer_attributes_partial(model_test, "*", {"enable": False})

    x = torch.randn(2, 3)
    out_1 = model_ref(x)
    out_2 = model_test(x)

    assert torch.allclose(out_1, out_2)

    mtq.set_quantizer_attributes_partial(model_test, "*input_quantizer", {"enable": True})
    mtq.set_quantizer_attributes_partial(model_test, "*weight_quantizer", {"enable": True})
    model_ref = PytorchModel()
    model_ref.load_state_dict(model_test.state_dict())

    out_1 = model_ref(x)
    out_2 = model_test(x)
    assert torch.allclose(out_1, out_2)


@pytest.mark.skipif(
    Version(transformers.__version__) < Version("5.0"),
    reason="test_dbrx is not supported for transformers<5.0",
)
def test_dbrx():
    assert DbrxExperts in QuantModuleRegistry
    assert DbrxExpertGLU in QuantModuleRegistry

    config = DbrxConfig(
        ffn_config=DbrxFFNConfig(ffn_hidden_size=8, moe_num_experts=2, hidden_size=32),
        hidden_size=32,
    )

    model_ref = DbrxFFN(config)
    model_test = DbrxFFN(config)
    with torch.no_grad():
        model_ref.experts.mlp.w1.copy_(torch.randn(16, 32))
        model_ref.experts.mlp.v1.copy_(torch.randn(16, 32))
        model_ref.experts.mlp.w2.copy_(torch.randn(16, 32))

    model_test.load_state_dict(model_ref.state_dict())

    mtq.replace_quant_module(model_test)

    expertglu_ref = model_ref.experts.mlp
    expertglu_test = model_test.experts.mlp

    assert hasattr(expertglu_test, "w1_linear") and not hasattr(expertglu_test, "w1")
    assert hasattr(expertglu_test, "v1_linear") and not hasattr(expertglu_test, "v1")
    assert hasattr(expertglu_test, "w2_linear") and not hasattr(expertglu_test, "w2")

    # Weights are stored transposed (W = w1[i].T) to match F.linear semantics with
    # transformers 5.0's raw matmul: x @ w1[i] = F.linear(x, w1[i].T)
    assert torch.allclose(
        torch.concat([m.weight.T for m in expertglu_test.w1_linear], dim=0),
        expertglu_ref.w1,
    )

    mtq.set_quantizer_attributes_partial(model_test, "*", {"enable": False})

    # In transformers 5.0, the FFN input dimension is ffn_hidden_size (not hidden_size)
    x = torch.randn(1, 4, 8)
    out_1 = model_ref(x)
    out_2 = model_test(x)
    assert torch.allclose(out_1[0], out_2[0])


@pytest.mark.parametrize("method", ["gradient", "kl_div"])
@pytest.mark.parametrize("model_provider", [get_tiny_llama, get_tiny_qwen3_moe])
def test_autoquantize_huggingface(model_provider, method):
    if model_provider == get_tiny_qwen3_moe and Version(torch.__version__) < Version("2.9"):
        pytest.skip("torch 2.8 grouped_mm is CUDA-only")

    model = model_provider()
    input_ids = model.dummy_inputs["input_ids"]

    def forward_step(model, batch):
        return model(**batch) if method == "gradient" else model(**batch).logits

    warnings.filterwarnings(
        "error", message="AutoQuantize: Error enabling gradient checkpointing for huggingface model"
    )

    # Gradient checkpointing warning should only appear for gradient-based method
    context = (
        pytest.warns(
            UserWarning,
            match="AutoQuantize: Huggingface model detected - Enabling gradient checkpointing. ",
        )
        if method == "gradient"
        else nullcontext()
    )

    with context:
        best_model, search_history = mtq.auto_quantize(
            model,
            constraints={"effective_bits": 11.0},
            quantization_formats=[mtq.INT8_DEFAULT_CFG],
            data_loader=[{"input_ids": input_ids, "labels": input_ids} for _ in range(2)],
            forward_step=forward_step,
            loss_func=lambda output, data: output.loss,
            num_calib_steps=2,
            num_score_steps=2,
            verbose=True,
            method=method,
        )


@pytest.mark.parametrize(
    ("model_cls", "quant_config"),
    [
        (LlamaForCausalLM, mtq.INT4_AWQ_CFG),
        (AutoModelForCausalLM, mtq.INT4_AWQ_CFG),
    ],
)
def test_quantized_transformers_save_restore(tmp_path, model_cls, quant_config):
    tiny_llama_dir = create_tiny_llama_dir(tmp_path, dtype=torch.float32)
    # update config to fit test cases
    if quant_config == mtq.INT4_AWQ_CFG:
        import copy

        quant_config = copy.deepcopy(quant_config)
        for entry in quant_config["quant_cfg"]:
            if entry["quantizer_name"] == "*weight_quantizer":
                entry.setdefault("cfg", {})["block_sizes"] = {-1: 16}
                break
    else:
        raise ValueError(f"Unsupported quant_config: {quant_config}")

    model_ref = model_cls.from_pretrained(tiny_llama_dir)
    mtq.quantize(model_ref, quant_config, lambda model: model(**model.dummy_inputs))
    mtq.compress(model_ref)
    model_ref.save_pretrained(tiny_llama_dir / "modelopt_model")
    assert os.path.exists(tiny_llama_dir / "modelopt_model/modelopt_state.pth")

    model_test = model_cls.from_pretrained(tiny_llama_dir / "modelopt_model")
    tf_modelopt_state_and_output_tester(model_ref, model_test)


def test_is_homogeneous_hf_model_llama():
    model = get_tiny_llama()
    assert is_homogeneous_hf_model(model)


def test_is_homogeneous_hf_model_gpt_oss():
    model = get_tiny_gpt_oss(num_hidden_layers=1)
    assert is_homogeneous_hf_model(model)


def test_hf_decoder_discoverer_registration_path():
    model = get_tiny_llama()
    assert any(
        is_supported is is_homogeneous_hf_model and discoverer is get_homogeneous_hf_decoder_layers
        for is_supported, discoverer in LayerActivationCollector._decoder_layer_support
    )
    assert LayerActivationCollector.get_decoder_layers(model) is get_homogeneous_hf_decoder_layers(
        model
    )


def test_qwen3_5_hybrid_attention_quantize():
    """Verify FP8 quantization disables all linear_attn quantizers while self_attn is quantized."""
    model = get_tiny_qwen3_5()
    mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(**m.dummy_inputs))

    for name, module in model.named_modules():
        if not hasattr(module, "weight_quantizer"):
            continue
        if "linear_attn" in name:
            assert not module.weight_quantizer.is_enabled, (
                f"linear_attn module {name} should have weight_quantizer disabled"
            )
            assert not module.input_quantizer.is_enabled, (
                f"linear_attn module {name} should have input_quantizer disabled"
            )
        elif "self_attn" in name and "layernorm" not in name:
            assert module.weight_quantizer.is_enabled, (
                f"self_attn module {name} should have weight_quantizer enabled"
            )


@pytest.mark.skipif(
    Version(torch.__version__) < Version("2.9"),
    reason="torch 2.8 grouped_mm is CUDA-only",
)
def test_qwen3_5_moe_experts_not_quantized():
    """Verify MoE expert quantizers are disabled when build_quant_cfg rules are applied."""
    model = get_tiny_qwen3_5_moe()

    import copy

    quant_cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
    quant_cfg["quant_cfg"].append({"quantizer_name": "*experts*", "enable": False})

    mtq.quantize(model, quant_cfg, lambda m: m(**m.dummy_inputs))

    for name, module in model.named_modules():
        if not hasattr(module, "weight_quantizer"):
            continue
        if "experts" in name:
            assert not module.weight_quantizer.is_enabled, (
                f"expert module {name} should have weight_quantizer disabled"
            )