UT refactor for preview

.azure-pipelines/scripts/ut/run_ut.sh

@@ -19,8 +19,7 @@ cd /auto-round && uv pip install .
 echo "##[endgroup]"
 uv pip list
 
-cd /auto-round/test/test_cpu || exit 1
-find . -type f -exec sed -i '/sys\.path\.insert(0, "\.\.")/d' {} +
+cd /auto-round/test || exit 1
 
 export LD_LIBRARY_PATH=${HOME}/.venv/lib/:$LD_LIBRARY_PATH
 export FORCE_BF16=1
@@ -32,7 +31,7 @@ mkdir -p ${LOG_DIR}
 ut_log_name=${LOG_DIR}/ut.log
 
 # Split test files into 5 parts
-find . -name "test*.py" | sort > all_tests.txt
+find ./test_cpu -name "test*.py" | sort > all_tests.txt
 total_lines=$(wc -l < all_tests.txt)
 NUM_CHUNKS=5
 q=$(( total_lines / NUM_CHUNKS ))

.azure-pipelines/scripts/ut/run_ut_cuda.sh

@@ -27,16 +27,14 @@ function create_conda_env() {
 
     # install AutoRound
     cd ${REPO_PATH}
-    pip uninstall auto-round -y
+    uv pip install torch==2.8.0 torchvision
     uv pip install -r requirements.txt
-    sed -i '/^torch==/d;/^transformers==/d;/^lm-eval==/d' requirements.txt
     if [ -d "/proc/driver/nvidia" ]; then
         export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
         export LD_LIBRARY_PATH=$(python -c "import site; print(site.getsitepackages()[0])")/nvidia/nvjitlink/lib:$LD_LIBRARY_PATH
     fi
     uv pip install --no-build-isolation .
     uv pip install pytest-cov pytest-html cmake==4.0.2
-    uv pip install torch==2.8.0 torchvision
 }
 
 function print_test_results_table() {
@@ -92,23 +90,23 @@ function run_unit_test() {
     # install unit test dependencies
     create_conda_env
 
-    cd ${REPO_PATH}/test/test_cuda
+    cd ${REPO_PATH}/test
     rm -rf .coverage* *.xml *.html
 
     uv pip install -v git+https://github.com/casper-hansen/AutoAWQ.git --no-build-isolation
     uv pip install -v git+https://github.com/ModelCloud/GPTQModel.git@v2.2.0 --no-build-isolation
     uv pip install -r https://raw.githubusercontent.com/ModelCloud/GPTQModel/refs/heads/main/requirements.txt
     CMAKE_ARGS="-DGGML_CUDA=on -DLLAVA_BUILD=off" uv pip install llama-cpp-python
     uv pip install 'git+https://github.com/ggml-org/llama.cpp.git#subdirectory=gguf-py'
-    uv pip install -r requirements.txt
-    uv pip install -r requirements_diffusion.txt
+    uv pip install -r test_cuda/requirements.txt
+    uv pip install -r test_cuda/requirements_diffusion.txt
 
     pip list > ${LOG_DIR}/ut_pip_list.txt
     export COVERAGE_RCFILE=${REPO_PATH}/.azure-pipelines/scripts/ut/.coverage
     local auto_round_path=$(python -c 'import auto_round; print(auto_round.__path__[0])')
 
     # run unit tests individually with separate logs
-    for test_file in $(find . -name "test_*.py" ! -name "test_*vlms.py" ! -name "test_llmc*.py" | sort); do
+    for test_file in $(find ./test_cuda -name "test_*.py" ! -name "test_*vlms.py" ! -name "test_llmc*.py" | sort); do
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_cuda_${test_basename}.log
         echo "Running ${test_file}..."
@@ -128,7 +126,7 @@ function run_unit_test() {
 function run_unit_test_vlm() {
     # install unit test dependencies
     create_conda_env
-    cd ${REPO_PATH}/test/test_cuda
+    cd ${REPO_PATH}/test
     rm -rf .coverage* *.xml *.html
 
     uv pip install git+https://github.com/haotian-liu/LLaVA.git@v1.2.2 --no-deps
@@ -138,14 +136,14 @@ function run_unit_test_vlm() {
     uv pip install git+https://github.com/deepseek-ai/DeepSeek-VL2.git timm attrdict --no-deps
     uv pip install -v git+https://github.com/casper-hansen/AutoAWQ.git@v0.2.0 --no-build-isolation
     uv pip install flash-attn==2.7.4.post1 --no-build-isolation
-    uv pip install -r requirements_vlm.txt
+    uv pip install -r test_cuda/requirements_vlm.txt
 
     pip list > ${LOG_DIR}/vlm_ut_pip_list.txt
     export COVERAGE_RCFILE=${REPO_PATH}/.azure-pipelines/scripts/ut/.coverage
     local auto_round_path=$(python -c 'import auto_round; print(auto_round.__path__[0])')
 
     # run VLM unit tests individually with separate logs
-    for test_file in $(find . -name "test*vlms.py"); do
+    for test_file in $(find ./test_cuda -name "test*vlms.py"); do
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_cuda_vlm_${test_basename}.log
         echo "Running ${test_file}..."
@@ -166,17 +164,17 @@ function run_unit_test_llmc() {
     # install unit test dependencies
     create_conda_env
 
-    cd ${REPO_PATH}/test/test_cuda
+    cd ${REPO_PATH}/test
     rm -rf .coverage* *.xml *.html
 
-    uv pip install -r requirements_llmc.txt
+    uv pip install -r test_cuda/requirements_llmc.txt
 
     pip list > ${LOG_DIR}/llmc_ut_pip_list.txt
     export COVERAGE_RCFILE=${REPO_PATH}/.azure-pipelines/scripts/ut/.coverage
     local auto_round_path=$(python -c 'import auto_round; print(auto_round.__path__[0])')
 
     # run unit tests individually with separate logs
-    for test_file in $(find . -name "test_llmc*.py" | sort); do
+    for test_file in $(find ./test_cuda -name "test_llmc*.py" | sort); do
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_cuda_llmc_${test_basename}.log
         echo "Running ${test_file}..."

.azure-pipelines/scripts/ut/run_ut_hpu.sh

@@ -7,8 +7,7 @@ export TQDM_MININTERVAL=60
 pip install pytest-cov pytest-html
 pip list
 
-cd /auto-round/test/test_hpu || exit 1
-find . -type f -exec sed -i '/sys\.path\.insert(0, "\.\.")/d' {} +
+cd /auto-round/test || exit 1
 
 export LD_LIBRARY_PATH=/usr/local/lib/:$LD_LIBRARY_PATH
 export FORCE_BF16=1
@@ -19,8 +18,8 @@ LOG_DIR=/auto-round/log_dir
 mkdir -p ${LOG_DIR}
 ut_log_name=${LOG_DIR}/ut.log
 
-find . -name "test*.py" | sed "s,\.\/,python -m pytest --cov=\"${auto_round_path}\" --cov-report term --html=report.html --self-contained-html  --cov-report xml:coverage.xml --cov-append -vs --disable-warnings ,g" > run_lazy.sh
-find . -name "test*.py" | sed "s,\.\/,python -m pytest --mode compile --cov=\"${auto_round_path}\" --cov-report term --html=report.html --self-contained-html  --cov-report xml:coverage.xml --cov-append -vs --disable-warnings ,g" > run_compile.sh
+find ./test_hpu -name "test*.py" | sed "s,\.\/,python -m pytest --cov=\"${auto_round_path}\" --cov-report term --html=report.html --self-contained-html  --cov-report xml:coverage.xml --cov-append -vs --disable-warnings ,g" > run_lazy.sh
+find ./test_hpu -name "test*.py" | sed "s,\.\/,python -m pytest --mode compile --cov=\"${auto_round_path}\" --cov-report term --html=report.html --self-contained-html  --cov-report xml:coverage.xml --cov-append -vs --disable-warnings ,g" > run_compile.sh
 
 cat run_lazy.sh
 bash run_lazy.sh 2>&1 | tee ${ut_log_name}

auto_round/compressors/mllm/compressor.py

@@ -206,6 +206,9 @@ def __init__(
         if hasattr(model, "name_or_path") and any([name in model.name_or_path for name in MISTRAL_3_2_MODELS]):
             template = "mistral3_2"
         if iters > 0:
+            # TODO: Remove after fixing https://github.com/huggingface/transformers/issues/43005
+            model.config.model_type = model.config.to_dict()["model_type"]
+
             if template is None and model.config.model_type not in TEMPLATES:
                 self.template = None
             else:

auto_round/utils/model.py

@@ -1047,6 +1047,11 @@ def set_module(model, key, new_module):
     setattr(module, name_list[-1], new_module)
 
 
+# For getting and setting attribution, such as 'lm_head.weight'
+get_attr = get_module
+set_attr = set_module
+
+
 def get_layer_features(layer):
     """Extracts input and output feature dimensions for supported layers."""
     from auto_round.utils import deepspeed_exists

test/README.md

@@ -0,0 +1,46 @@
+# Unit Test (UT) Guide
+
+This project uses `pytest` for unit testing. All test cases are under the `test/` directory. Below is a simple guide for new users to write and run UTs:
+
+## 1. Environment Setup
+- Recommended Python 3.8 or above.
+- Install dependencies:
+  ```sh
+  pip install -r ../requirements.txt
+  pip install pytest
+  ```
+
+## 2. Test Structure
+- Place your test files in the `test/` directory, and name them starting with `test_`.
+- You can refer to existing `test_*.py` files.
+- Common fixtures (such as `tiny_opt_model`, `opt_model`, `opt_tokenizer`, `dataloader`) and helper functions (such as `model_infer`) are defined in `confest.py` and `helpers.py` and can be imported directly.
+- Example:
+  ```python
+  # test_example.py
+    from ..helpers import model_infer
+
+    def test_model_infer(tiny_opt_model, opt_tokenizer):
+        result = model_infer(tiny_opt_model, opt_tokenizer, input_text="hello world")
+        assert result is not None
+  ```
+
+## 3. Running Tests
+- In the `test/` directory, run:
+  ```sh
+  pytest
+  ```
+- You can specify a single file or test case:
+  ```sh
+  pytest test_xxx.py
+  pytest -k "test_func_name"
+  ```
+
+## 4. Debugging Tips
+- `confest.py` adds the parent directory to `sys.path`, so you can debug without installing the local package.
+- You can directly import project source code in your test cases.
+
+## 5. Reference
+- Fixtures are defined in `confest.py` and `fixtures.py`
+- Helper functions are in `helpers.py`
+
+If you have any questions, feel free to open an issue.

test/test_hpu/conftest.py → test/conftest.py

@@ -1,9 +1,16 @@
 import os
+import sys
 from typing import Mapping
 
 import pytest
 
+from .fixtures import *
 
+# Easy debugging without installing auto-round.
+sys.path.insert(0, "..")
+
+
+### HPU related configuration, usage: `pytest --mode=compile/lazy``
 def pytest_addoption(parser):
     parser.addoption(
         "--mode",

test/fixtures.py

@@ -0,0 +1,169 @@
+import os
+import shutil
+
+import pytest
+import torch
+import transformers
+
+from .helpers import (
+    DataLoader,
+    deepseek_v2_name_or_path,
+    gemma_name_or_path,
+    get_tiny_model,
+    gptj_name_or_path,
+    lamini_name_or_path,
+    opt_name_or_path,
+    phi2_name_or_path,
+    qwen_2_5_vl_name_or_path,
+    qwen_moe_name_or_path,
+    qwen_name_or_path,
+    qwen_vl_name_or_path,
+    save_tiny_model,
+)
+
+
+# Create tiny model path fixtures for testing
+@pytest.fixture(scope="session")
+def tiny_opt_model_path():
+    model_name_or_path = opt_name_or_path
+    tiny_model_path = "./tmp/tiny_opt_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_lamini_model_path():
+    model_name_or_path = lamini_name_or_path
+    tiny_model_path = "./tmp/tiny_lamini_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_gptj_model_path():
+    model_name_or_path = gptj_name_or_path
+    tiny_model_path = "./tmp/tiny_gptj_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_phi2_model_path():
+    model_name_or_path = phi2_name_or_path
+    tiny_model_path = "./tmp/tiny_phi2_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_deepseek_v2_model_path():
+    model_name_or_path = deepseek_v2_name_or_path
+    tiny_model_path = "./tmp/tiny_deepseek_v2_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path, num_layers=2)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_gemma_model_path():
+    model_name_or_path = gemma_name_or_path
+    tiny_model_path = "./tmp/tiny_gemma_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path, num_layers=2)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_qwen_model_path():
+    model_name_or_path = qwen_name_or_path
+    tiny_model_path = "./tmp/tiny_qwen_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_untied_qwen_model_path():
+    model_name_or_path = qwen_name_or_path
+    tiny_model_path = "./tmp/tiny_untied_qwen_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path, force_untie=True)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_qwen_moe_model_path():
+    model_name_or_path = qwen_moe_name_or_path
+    tiny_model_path = "./tmp/tiny_qwen_moe_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path, num_layers=2)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_qwen_vl_model_path():
+    model_name_or_path = qwen_vl_name_or_path
+    tiny_model_path = "./tmp/tiny_qwen_vl_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path, num_layers=2, is_mllm=True)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(scope="session")
+def tiny_qwen_2_5_vl_model_path():
+    model_name_or_path = qwen_2_5_vl_name_or_path
+    tiny_model_path = "./tmp/tiny_qwen_2_5_vl_model_path"
+    tiny_model_path = save_tiny_model(model_name_or_path, tiny_model_path, num_layers=2, is_mllm=True)
+    yield tiny_model_path
+    shutil.rmtree(tiny_model_path)
+
+
+@pytest.fixture(autouse=True, scope="session")
+def clean_tmp_model_folder():
+    yield
+    shutil.rmtree("./tmp", ignore_errors=True)  # unittest default workspace
+    shutil.rmtree("./tmp_autoround", ignore_errors=True)  # autoround default workspace
+
+
+# Create objective fixtures for testing
+@pytest.fixture(scope="function")
+def tiny_opt_model():
+    model_name_or_path = opt_name_or_path
+    return get_tiny_model(model_name_or_path, num_layers=2)
+
+
+@pytest.fixture(scope="function")
+def opt_model():
+    model_name_or_path = opt_name_or_path
+    model = transformers.AutoModelForCausalLM.from_pretrained(model_name_or_path, dtype="auto", trust_remote_code=True)
+    return model
+
+
+@pytest.fixture(scope="session")
+def opt_tokenizer():
+    model_name_or_path = opt_name_or_path
+    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
+    return tokenizer
+
+
+@pytest.fixture(scope="function")
+def model():
+    model_name_or_path = opt_name_or_path
+    model = transformers.AutoModelForCausalLM.from_pretrained(model_name_or_path, dtype="auto", trust_remote_code=True)
+    return model
+
+
+@pytest.fixture(scope="session")
+def tokenizer():
+    model_name_or_path = opt_name_or_path
+    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
+    return tokenizer
+
+
+@pytest.fixture(scope="session")
+def dataloader():
+    return DataLoader()