# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team import numpy as np import torch import torch.nn as nn from unit.common import DistributedTest from deepspeed.accelerator import get_accelerator from deepspeed.inference.quantization.quantization import _init_group_wise_weight_quantization from deepspeed.inference.quantization.utils import Quantizer, DeQuantizer from deepspeed.inference.quantization.layers import QuantizedLinear from deepspeed.utils.torch import required_torch_version from transformers.models.opt.modeling_opt import OPTDecoderLayer from transformers import AutoConfig, OPTConfig, AutoModel import pytest from collections import OrderedDict from typing import Dict from deepspeed.ops.aio import AsyncIOBuilder device = get_accelerator().device_name() if get_accelerator().is_available() else 'cpu' if not required_torch_version(min_version=1.11): pytest.skip("torch.Tensor.bitwise_left_shift in INT4 quantizer needs torch 1.11 or above.", allow_module_level=True) def reset_random(seed=1234): np.random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) get_accelerator().manual_seed_all(seed) def quantization_test_helper(pre_quant_type: torch.dtype, num_bits: int): reset_random() num_group = 1024 * 32 group_size = 64 quantization_config = {'num_bits': num_bits, 'group_size': group_size, 'group_dim': 1, 'symmetric': False} quantizer = Quantizer(config=quantization_config) dequantizer = DeQuantizer(config=quantization_config, dtype=pre_quant_type) data = torch.randn(num_group, group_size, dtype=pre_quant_type, device=device) quantized_data, scale_buf, min_vals = quantizer.quantize(data) dequantized_data = dequantizer.dequantize(quantized_data, scale_buf, min_vals) max_diff = torch.max(torch.abs(data - dequantized_data)) mean_diff = torch.mean(torch.abs(data - dequantized_data)) # This threshold value is emperically selected. assert mean_diff < 0.15 and max_diff < 0.5, f'Numeric error exceed threshold, mean diff {mean_diff} (threshold 0.15), max diff {max_diff} (threshold 0.5)' def zero3_post_init_quantization_test_helper(cpu_offload: bool, nvme_offload: bool, bits: int, nvme_path=None): import deepspeed from transformers.integrations.deepspeed import HfDeepSpeedConfig if nvme_offload and not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]: pytest.skip('Skip tests since async-io is not compatible') def get_zero3_ds_config(hf_config: OPTConfig, cpu_offload: bool, nvme_offload: bool, bits: int) -> Dict: GB = 1 << 30 ds_config = { "fp16": { "enabled": True, }, "zero_optimization": { "stage": 3, "stage3_prefetch_bucket_size": 2 * hf_config.hidden_size * hf_config.hidden_size, "stage3_param_persistence_threshold": hf_config.hidden_size, "stage3_max_live_parameters": 2 * hf_config.hidden_size * hf_config.hidden_size }, "steps_per_print": 2000, "train_micro_batch_size_per_gpu": 1, "wall_clock_breakdown": False, 'weight_quantization': { 'post_init_quant': { 'fc': { 'num_bits': bits, 'group_size': 32, 'group_dim': 1, 'symmetric': False }, 'self_attn.q_proj': { 'num_bits': bits, 'group_size': 32, 'group_dim': 1, 'symmetric': False }, 'self_attn.k_proj': { 'num_bits': bits, 'group_size': 32, 'group_dim': 1, 'symmetric': False }, 'self_attn.v_proj': { 'num_bits': bits, 'group_size': 32, 'group_dim': 1, 'symmetric': False }, 'self_attn.out_proj': { 'num_bits': bits, 'group_size': 32, 'group_dim': 1, 'symmetric': False }, 'lm_head': { 'num_bits': bits, 'group_size': 32, 'group_dim': 1, 'symmetric': False }, 'embed_tokens': { 'num_bits': bits, 'group_size': 32, 'group_dim': 1, 'symmetric': False }, } } } if cpu_offload: ds_config["zero_optimization"]["offload_param"] = dict(device="cpu", pin_memory=1) if nvme_offload: ds_config["zero_optimization"]["offload_param"] = dict( device="nvme", pin_memory=True, nvme_path=nvme_path or '~/tmp_offload_dir', buffer_count=5, buffer_size=1 * GB, ) ds_config["aio"] = { "block_size": 1048576, "queue_depth": 8, "thread_count": 1, "single_submit": False, "overlap_events": True, } return ds_config hf_config = AutoConfig.from_pretrained('facebook/opt-125m') ds_config = get_zero3_ds_config(hf_config=hf_config, cpu_offload=cpu_offload, nvme_offload=nvme_offload, bits=bits) input_ids = torch.ones(1, 16, dtype=torch.int32, device=device) attention_mask = torch.ones(1, 16, dtype=torch.float32, device=device) with torch.no_grad(): ref_model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16).to(device) ref_model.eval() ref_output = ref_model(input_ids=input_ids, attention_mask=attention_mask) with torch.no_grad(): dschf = HfDeepSpeedConfig(ds_config) model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16) model = model.eval() model = _init_group_wise_weight_quantization(model, ds_config) ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0] ds_engine.module.eval() model = ds_engine.module output = model(input_ids=input_ids, attention_mask=attention_mask) mean_diff = torch.mean(torch.abs(output.last_hidden_state - ref_output.last_hidden_state)) # This threshold value is emperically selected. assert mean_diff < 0.4, f'Numeric error exceed threshold, relative error {mean_diff} (threshold 0.4)' def zero3_quantized_initialization_test_helper(cpu_offload: bool, nvme_offload: bool, bits: int, nvme_path=None): import deepspeed from transformers.integrations.deepspeed import HfDeepSpeedConfig if nvme_offload and not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]: pytest.skip('Skip tests since async-io is not compatible') def get_zero3_ds_config(hf_config: OPTConfig, cpu_offload: bool, nvme_offload: bool, bits: int) -> Dict: GB = 1 << 30 ds_config = { "fp16": { "enabled": True, }, "zero_optimization": { "stage": 3, "stage3_prefetch_bucket_size": 2 * hf_config.hidden_size * hf_config.hidden_size, "stage3_param_persistence_threshold": hf_config.hidden_size, "stage3_max_live_parameters": 2 * hf_config.hidden_size * hf_config.hidden_size }, "steps_per_print": 2000, "train_micro_batch_size_per_gpu": 1, "wall_clock_breakdown": False, 'weight_quantization': { 'quantized_initialization': { 'num_bits': bits, 'group_size': 32, 'group_dim': 1, 'symmetric': False }, } } if cpu_offload: ds_config["zero_optimization"]["offload_param"] = dict(device="cpu", pin_memory=1) if nvme_offload: ds_config["zero_optimization"]["offload_param"] = dict( device="nvme", pin_memory=True, nvme_path=nvme_path or '~/tmp_offload_dir', buffer_count=5, buffer_size=1 * GB, ) ds_config["aio"] = { "block_size": 1048576, "queue_depth": 8, "thread_count": 1, "single_submit": False, "overlap_events": True, } return ds_config hf_config = AutoConfig.from_pretrained('facebook/opt-125m') ds_config = get_zero3_ds_config(hf_config=hf_config, cpu_offload=cpu_offload, nvme_offload=nvme_offload, bits=bits) input_ids = torch.ones(1, 16, dtype=torch.int32, device=device) attention_mask = torch.ones(1, 16, dtype=torch.float32, device=device) with torch.no_grad(): ref_model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16).to(device) ref_model.eval() ref_output = ref_model(input_ids=input_ids, attention_mask=attention_mask) with torch.no_grad(): dschf = HfDeepSpeedConfig(ds_config) model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16) model = model.eval() ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0] ds_engine.module.eval() model = ds_engine.module output = model(input_ids=input_ids, attention_mask=attention_mask) mean_diff = torch.mean(torch.abs(output.last_hidden_state - ref_output.last_hidden_state)) # This threshold value is emperically selected. assert mean_diff < 0.4, f'Numeric error exceed threshold, relative error {mean_diff} (threshold 0.4)' @pytest.fixture(params=[4, 8], ids=["4bits", "8bits"]) def quantization_bits(request): return request.param @pytest.fixture(params=[0, 1], ids=["0", "1"]) def group_dim(request): return request.param class TestQuantizedInt(DistributedTest): def test_model_quantization(self, quantization_bits): reset_random() config = AutoConfig.from_pretrained('facebook/opt-125m') with torch.no_grad(): model = OPTDecoderLayer(config).half().to(device) bits = quantization_bits ds_config = { 'weight_quantization': { 'post_init_quant': { 'fc': { 'num_bits': bits, 'group_size': 64, 'group_dim': 0, 'symmetric': False }, 'self_attn.q_proj': { 'num_bits': bits, 'group_size': 64, 'group_dim': 0, 'symmetric': False }, 'self_attn.k_proj': { 'num_bits': bits, 'group_size': 64, 'group_dim': 0, 'symmetric': False }, 'self_attn.v_proj': { 'num_bits': bits, 'group_size': 64, 'group_dim': 0, 'symmetric': False }, 'self_attn.out_proj': { 'num_bits': bits, 'group_size': 64, 'group_dim': 0, 'symmetric': False } } } } model = _init_group_wise_weight_quantization(model, ds_config) assert type(model.fc1) is QuantizedLinear assert type(model.fc2) is QuantizedLinear assert type(model.self_attn.q_proj) is QuantizedLinear assert type(model.self_attn.k_proj) is QuantizedLinear assert type(model.self_attn.v_proj) is QuantizedLinear assert type(model.self_attn.out_proj) is QuantizedLinear @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM') def test_quantized_linear(self, quantization_bits, group_dim): reset_random() layers = [] for idx in range(5): layers.append( (f'layer_{idx}', nn.Linear(in_features=128, out_features=128, dtype=torch.float16, device=device))) input_tensor = torch.randn(32, 128, dtype=torch.float16, device=device) with torch.no_grad(): model = nn.Sequential(OrderedDict(layers)) ref_output = model(input_tensor) ds_config = { 'weight_quantization': { 'post_init_quant': { 'layer': { 'num_bits': quantization_bits, 'group_size': 64, 'group_dim': group_dim, 'symmetric': False } } } } model = _init_group_wise_weight_quantization(model, ds_config) assert type(model.layer_0) is QuantizedLinear assert type(model.layer_1) is QuantizedLinear assert type(model.layer_2) is QuantizedLinear assert type(model.layer_3) is QuantizedLinear assert type(model.layer_4) is QuantizedLinear output = model(input_tensor) mean_diff = torch.mean(torch.abs(ref_output - output)) # This threshold value is emperically selected. assert mean_diff < 0.15, f'Numeric error exceed threshold, mean diff {mean_diff}' def test_float_int4_quantization(self): reset_random() quantization_test_helper(torch.float32, 4) @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM') def test_half_int4_quantization(self): reset_random() quantization_test_helper(torch.float16, 4) def test_float_int8_quantization(self): reset_random() quantization_test_helper(torch.float32, 8) def test_half_int8_quantization(self): reset_random() quantization_test_helper(torch.float16, 8) @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM') def test_zero3_int4_post_init_quant(self, quantization_bits): reset_random() zero3_post_init_quantization_test_helper(cpu_offload=False, nvme_offload=False, bits=quantization_bits) @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM') def test_zero3_int4_post_init_quant_cpu_offload(self, quantization_bits): reset_random() zero3_post_init_quantization_test_helper(cpu_offload=True, nvme_offload=False, bits=quantization_bits) @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM') def test_zero3_int4_post_init_quant_nvme_offload(self, tmpdir): reset_random() zero3_post_init_quantization_test_helper(cpu_offload=False, nvme_offload=True, bits=4, nvme_path=str(tmpdir.join("nvme_offload"))) @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM') def test_zero3_int4_quantized_initialization(self, quantization_bits): reset_random() zero3_quantized_initialization_test_helper(cpu_offload=False, nvme_offload=False, bits=quantization_bits) @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM') def test_zero3_int4_quantized_initialization_cpu_offload(self, quantization_bits): reset_random() zero3_quantized_initialization_test_helper(cpu_offload=True, nvme_offload=False, bits=quantization_bits) @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM') def test_zero3_int4_quantized_initialization_nvme_offload(self, tmpdir): reset_random() zero3_quantized_initialization_test_helper(cpu_offload=False, nvme_offload=True, bits=4, nvme_path=str(tmpdir.join("nvme_offload")))