Upload tests/test_quantize.py with huggingface_hub

tests/test_quantize.py

@@ -0,0 +1,476 @@
+"""
+量化模块单元测试
+
+测试量化模块的所有功能，包括：
+- GGUF 量化
+- AWQ 量化
+- GPTQ 量化
+- PPL 评估
+- 格式转换
+- VRAM 检查和优化
+
+使用 mock 避免实际模型加载和量化过程。
+"""
+
+import os
+import sys
+import pytest
+from pathlib import Path
+from unittest.mock import MagicMock, patch, Mock
+import subprocess
+
+# 添加项目路径
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from hos_optimizer.quantize import (
+    QuantizationError,
+    check_vram_availability,
+    optimize_for_low_vram,
+    quantize_gguf,
+    quantize_awq,
+    quantize_gptq,
+    evaluate_perplexity,
+    convert_format,
+    get_model_size,
+    VRAM_8GB_CONFIG,
+)
+
+
+class TestVRAMCheck:
+    """VRAM 检查相关测试"""
+
+    def test_check_vram_no_cuda(self):
+        """测试无 CUDA 支持时的 VRAM 检查"""
+        with patch("torch.cuda.is_available", return_value=False):
+            result = check_vram_availability()
+            
+            assert result["available"] is False
+            assert result["total_vram_gb"] == 0
+            assert result["free_vram_gb"] == 0
+            assert result["device"] == "cpu"
+
+    def test_check_vram_with_cuda(self):
+        """测试有 CUDA 支持时的 VRAM 检查"""
+        mock_device = MagicMock()
+        mock_device.total_memory = 8 * 1024 ** 3  # 8GB
+        
+        with patch("torch.cuda.is_available", return_value=True), \
+             patch("torch.cuda.current_device", return_value=0), \
+             patch("torch.cuda.get_device_properties", return_value=mock_device), \
+             patch("torch.cuda.memory_allocated", return_value=2 * 1024 ** 3), \
+             patch("torch.cuda.get_device_name", return_value="Test GPU"):
+            
+            result = check_vram_availability()
+            
+            assert result["available"] is True
+            assert result["total_vram_gb"] == 8.0
+            assert result["free_vram_gb"] == 6.0
+            assert result["device"] == "Test GPU"
+
+    def test_optimize_for_low_vram_enabled(self):
+        """测试低 VRAM 优化配置启用用的情况"""
+        config = {"batch_size": 4, "seq_length": 1024}
+        
+        with patch("hos_optimizer.quantize.check_vram_availability") as mock_check:
+            mock_check.return_value = {
+                "available": True,
+                "free_vram_gb": 6.0,
+                "total_vram_gb": 8.0,
+                "device": "Test GPU"
+            }
+            
+            optimized = optimize_for_low_vram(config)
+            
+            # 应该应用 8GB 优化配置
+            assert optimized["max_batch_size"] == VRAM_8GB_CONFIG["max_batch_size"]
+            assert optimized["max_seq_length"] == VRAM_8GB_CONFIG["max_seq_length"]
+            assert optimized["gradient_checkpointing"] is True
+            assert optimized["offload_to_cpu"] is True
+            # 原有配置应该保留
+            assert optimized["batch_size"] == 4
+            assert optimized["seq_length"] == 1024
+
+    def test_optimize_for_low_vram_disabled(self):
+        """测试低 VRAM 优化配置不适用的情况"""
+        config = {"batch_size": 4, "seq_length": 1024}
+        
+        with patch("hos_optimizer.quantize.check_vram_availability") as mock_check:
+            mock_check.return_value = {
+                "available": True,
+                "free_vram_gb": 12.0,  # 超过 8GB
+                "total_vram_gb": 16.0,
+                "device": "Test GPU"
+            }
+            
+            optimized = optimize_for_low_vram(config)
+            
+            # 不应该应用优化配置
+            assert optimized == config
+
+
+class TestGGUFQuantization:
+    """GGUF 量化测试"""
+
+    def test_quantize_gguf_success(self, tmp_dir):
+        """测试 GGUF 量化成功场景"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model.gguf")
+        llama_cpp_path = "/path/to/llama.cpp"
+        
+        # Mock 所有依赖
+        with patch("subprocess.run") as mock_run, \
+             patch("hos_optimizer.quantize.AutoModelForCausalLM") as mock_model_cls, \
+             patch("hos_optimizer.quantize.AutoTokenizer") as mock_tokenizer_cls, \
+             patch("tempfile.TemporaryDirectory") as mock_tmpdir:
+            
+            # Mock subprocess 调用
+            mock_run.return_value = MagicMock(returncode=0)
+            
+            # Mock 临时目录
+            mock_tmpdir.return_value.__enter__.return_value = tmp_dir
+            
+            # Mock 模型和分词器
+            mock_model = MagicMock()
+            mock_tokenizer = MagicMock()
+            mock_model_cls.from_pretrained.return_value = mock_model
+            mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer
+            
+            result = quantize_gguf(
+                model_path=model_path,
+                output_path=output_path,
+                quant_type="Q4_K_M",
+                llama_cpp_path=llama_cpp_path
+            )
+            
+            assert result == output_path
+            # 验证调用了转换和量化命令
+            assert mock_run.call_count >= 2
+
+    def test_quantize_gguf_tool_not_found(self, tmp_dir):
+        """测试 GGUF 量化工具不存在的情况"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model.gguf")
+        
+        with patch("subprocess.run") as mock_run:
+            mock_run.side_effect = FileNotFoundError()
+            
+            with pytest.raises(QuantizationError) as exc_info:
+                quantize_gguf(model_path, output_path)
+            
+            assert "找不到 llama-quantize 工具" in str(exc_info.value)
+
+    def test_quantize_gguf_conversion_failed(self, tmp_dir):
+        """测试 GGUF 量化转换失败的情况"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model.gguf")
+        
+        with patch("subprocess.run") as mock_run, \
+             patch("hos_optimizer.quantize.AutoModelForCausalLM") as mock_model_cls, \
+             patch("hos_optimizer.quantize.AutoTokenizer") as mock_tokenizer_cls, \
+             patch("tempfile.TemporaryDirectory") as mock_tmpdir:
+            
+            # 第一次调用成功（检查工具），第二次调用失败（转换）
+            mock_run.side_effect = [
+                MagicMock(returncode=0),  # 检查工具
+                subprocess.CalledProcessError(1, "convert", stderr="Conversion failed")
+            ]
+            
+            mock_tmpdir.return_value.__enter__.return_value = tmp_dir
+            mock_model_cls.from_pretrained.return_value = MagicMock()
+            mock_tokenizer_cls.from_pretrained.return_value = MagicMock()
+            
+            with pytest.raises(QuantizationError) as exc_info:
+                quantize_gguf(model_path, output_path)
+            
+            assert "GGUF 量化失败" in str(exc_info.value)
+
+
+class TestAWQQuantization:
+    """AWQ 量化测试"""
+
+    def test_quantize_awq_success(self, tmp_dir):
+        """测试 AWQ 量化成功场景"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model-awq")
+        
+        with patch("hos_optimizer.quantize.AutoAWQForCausalLM") as mock_awq_cls, \
+             patch("hos_optimizer.quantize.AutoTokenizer") as mock_tokenizer_cls, \
+             patch("hos_optimizer.quantize.optimize_for_low_vram") as mock_optimize:
+            
+            mock_model = MagicMock()
+            mock_tokenizer = MagicMock()
+            mock_awq_cls.from_pretrained.return_value = mock_model
+            mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer
+            mock_optimize.return_value = {
+                "zero_point": True,
+                "q_group_size": 128,
+                "w_bit": 4,
+                "version": "GEMM"
+            }
+            
+            result = quantize_awq(
+                model_path=model_path,
+                output_path=output_path,
+                bits=4,
+                group_size=128
+            )
+            
+            assert result == output_path
+            mock_model.quantize.assert_called_once()
+            mock_model.save_quantized.assert_called_once_with(output_path)
+            mock_tokenizer.save_pretrained.assert_called_once_with(output_path)
+
+    def test_quantize_awq_missing_dependency(self, tmp_dir):
+        """测试 AWQ 量化缺少依赖的情况"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model-awq")
+        
+        with patch("hos_optimizer.quantize.AutoAWQForCausalLM") as mock_awq_cls:
+            mock_awq_cls.from_pretrained.side_effect = ImportError("autoawq")
+            
+            with pytest.raises(QuantizationError) as exc_info:
+                quantize_awq(model_path, output_path)
+            
+            assert "缺少依赖" in str(exc_info.value)
+            assert "autoawq" in str(exc_info.value)
+
+    def test_quantize_awq_quantization_failed(self, tmp_dir):
+        """测试 AWQ 量化过程失败的情况"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model-awq")
+        
+        with patch("hos_optimizer.quantize.AutoAWQForCausalLM") as mock_awq_cls, \
+             patch("hos_optimizer.quantize.AutoTokenizer") as mock_tokenizer_cls, \
+             patch("hos_optimizer.quantize.optimize_for_low_vram") as mock_optimize:
+            
+            mock_model = MagicMock()
+            mock_tokenizer = MagicMock()
+            mock_awq_cls.from_pretrained.return_value = mock_model
+            mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer
+            mock_optimize.return_value = {"zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM"}
+            
+            # 量化过程抛出异常
+            mock_model.quantize.side_effect = Exception("Quantization failed")
+            
+            with pytest.raises(QuantizationError) as exc_info:
+                quantize_awq(model_path, output_path)
+            
+            assert "AWQ 量化失败" in str(exc_info.value)
+
+
+class TestGPTQQuantization:
+    """GPTQ 量化测试"""
+
+    def test_quantize_gptq_success(self, tmp_dir):
+        """测试 GPTQ 量化成功场景"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model-gptq")
+        
+        with patch("hos_optimizer.quantize.AutoGPTQForCausalLM") as mock_gptq_cls, \
+             patch("hos_optimizer.quantize.AutoTokenizer") as mock_tokenizer_cls, \
+             patch("hos_optimizer.quantize.BaseQuantizeConfig") as mock_config_cls:
+            
+            mock_model = MagicMock()
+            mock_tokenizer = MagicMock()
+            mock_config = MagicMock()
+            
+            mock_gptq_cls.from_pretrained.return_value = mock_model
+            mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer
+            mock_config_cls.return_value = mock_config
+            
+            # Mock tokenizer 调用
+            mock_tokenizer.return_value = {"input_ids": MagicMock()}
+            
+            result = quantize_gptq(
+                model_path=model_path,
+                output_path=output_path,
+                bits=4,
+                group_size=128,
+                desc_act=False
+            )
+            
+            assert result == output_path
+            mock_model.quantize.assert_called_once()
+            mock_model.save_quantized.assert_called_once_with(output_path)
+
+    def test_quantize_gptq_invalid_bits(self, tmp_dir):
+        """测试 GPTQ 量化使用无效位数的情况"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model-gptq")
+        
+        with pytest.raises(ValueError) as exc_info:
+            quantize_gptq(model_path, output_path, bits=3)
+        
+        assert "仅支持 4-bit 或 8-bit" in str(exc_info.value)
+
+    def test_quantize_gptq_missing_dependency(self, tmp_dir):
+        """测试 GPTQ 量化缺少依赖的情况"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model-gptq")
+        
+        with patch("hos_optimizer.quantize.AutoGPTQForCausalLM") as mock_gptq_cls:
+            mock_gptq_cls.from_pretrained.side_effect = ImportError("auto_gptq")
+            
+            with pytest.raises(QuantizationError) as exc_info:
+                quantize_gptq(model_path, output_path, bits=4)
+            
+            assert "缺少依赖" in str(exc_info.value)
+            assert "auto-gptq" in str(exc_info.value)
+
+
+class TestPerplexityEvaluation:
+    """PPL 评估测试"""
+
+    def test_evaluate_perplexity_success(self, tmp_dir):
+        """测试 PPL 评估成功场景"""
+        model_path = os.path.join(tmp_dir, "model")
+        
+        with patch("hos_optimizer.quantize.AutoModelForCausalLM") as mock_model_cls, \
+             patch("hos_optimizer.quantize.AutoTokenizer") as mock_tokenizer_cls, \
+             patch("hos_optimizer.quantize.load_dataset") as mock_load_dataset:
+            
+            mock_model = MagicMock()
+            mock_tokenizer = MagicMock()
+            mock_dataset = MagicMock()
+            
+            mock_model_cls.from_pretrained.return_value = mock_model
+            mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer
+            mock_load_dataset.return_value = mock_dataset
+            
+            # Mock 数据集
+            mock_dataset.__getitem__.return_value = ["text1", "text2"]
+            
+            # Mock tokenizer 调用
+            mock_encodings = MagicMock()
+            mock_encodings.input_ids.size.return_value = (1, 100)
+            mock_tokenizer.return_value = mock_encodings
+            
+            # Mock 模型推理
+            mock_model.return_value = MagicMock(loss=MagicMock(item=MagicMock(return_value=2.5)))
+            mock_model.device = "cpu"
+            
+            result = evaluate_perplexity(
+                model_path=model_path,
+                dataset="wikitext",
+                max_samples=10,
+                stride=512
+            )
+            
+            assert isinstance(result, float)
+            assert result > 0
+
+    def test_evaluate_perplexity_missing_dataset(self, tmp_dir):
+        """测试 PPL 评估缺少 datasets 库的情况"""
+        model_path = os.path.join(tmp_dir, "model")
+        
+        with patch("hos_optimizer.quantize.AutoModelForCausalLM") as mock_model_cls, \
+             patch("hos_optimizer.quantize.AutoTokenizer") as mock_tokenizer_cls:
+            
+            mock_model_cls.from_pretrained.return_value = MagicMock()
+            mock_tokenizer_cls.from_pretrained.return_value = MagicMock()
+            
+            with patch("hos_optimizer.quantize.load_dataset", side_effect=ImportError("datasets")):
+                with pytest.raises(QuantizationError) as exc_info:
+                    evaluate_perplexity(model_path)
+                
+                assert "缺少依赖" in str(exc_info.value)
+                assert "datasets" in str(exc_info.value)
+
+
+class TestFormatConversion:
+    """格式转换测试"""
+
+    def test_convert_hf_to_gguf(self, tmp_dir):
+        """测试 HuggingFace 到 GGUF 格式转换"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model.gguf")
+        
+        with patch("hos_optimizer.quantize.quantize_gguf") as mock_quantize:
+            mock_quantize.return_value = output_path
+            
+            result = convert_format(
+                model_path=model_path,
+                output_path=output_path,
+                from_format="hf",
+                to_format="gguf"
+            )
+            
+            assert result == output_path
+            mock_quantize.assert_called_once()
+
+    def test_convert_hf_to_awq(self, tmp_dir):
+        """测试 HuggingFace 到 AWQ 格式转换"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model-awq")
+        
+        with patch("hos_optimizer.quantize.quantize_awq") as mock_quantize:
+            mock_quantize.return_value = output_path
+            
+            result = convert_format(
+                model_path=model_path,
+                output_path=output_path,
+                from_format="hf",
+                to_format="awq"
+            )
+            
+            assert result == output_path
+            mock_quantize.assert_called_once()
+
+    def test_convert_unsupported_path(self, tmp_dir):
+        """测试不支持的转换路径"""
+        model_path = os.path.join(tmp_dir, "model")
+        output_path = os.path.join(tmp_dir, "model-out")
+        
+        with pytest.raises(QuantizationError) as exc_info:
+            convert_format(
+                model_path=model_path,
+                output_path=output_path,
+                from_format="gguf",
+                to_format="awq"
+            )
+        
+        assert "不支持的转换路径" in str(exc_info.value)
+
+    def test_convert_gguf_to_hf_not_implemented(self, tmp_dir):
+        """测试 GGUF 到 HuggingFace 转换未实现"""
+        model_path = os.path.join(tmp_dir, "model.gguf")
+        output_path = os.path.join(tmp_dir, "model")
+        
+        with pytest.raises(QuantizationError) as exc_info:
+            convert_format(
+                model_path=model_path,
+                output_path=output_path,
+                from_format="gguf",
+                to_format="hf"
+            )
+        
+        assert "尚未实现" in str(exc_info.value)
+
+
+class TestModelSize:
+    """模型大小计算测试"""
+
+    def test_get_model_size_empty_dir(self, tmp_dir):
+        """测试空目录的模型大小"""
+        size = get_model_size(tmp_dir)
+        assert size == 0.0
+
+    def test_get_model_size_with_files(self, tmp_dir):
+        """测试包含模型文件的目录大小"""
+        # 创建测试文件
+        test_file = os.path.join(tmp_dir, "model.safetensors")
+        with open(test_file, "wb") as f:
+            f.write(b"0" * (1024 * 1024))  # 1MB
+        
+        size = get_model_size(tmp_dir)
+        assert size > 0
+        assert size < 0.01  # 应该约等于 0.001GB
+
+    def test_get_model_size_nonexistent_dir(self):
+        """测试不存在的目录"""
+        with pytest.raises(Exception):
+            get_model_size("/nonexistent/path")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])