""" 量化模块单元测试 测试量化模块的所有功能，包括： - 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"])