|
|
import fnmatch |
|
|
|
|
|
from .._utils import get_init_params |
|
|
from ..layers import (MLP, Attention, ColumnLinear, Embedding, GatedMLP, |
|
|
LayerNorm, RmsNorm, RowLinear) |
|
|
from ..layers.moe import MixtureOfExperts |
|
|
from ..models.modeling_utils import QuantConfig |
|
|
from ..parameter import Parameter |
|
|
from .layers import (FP8Linear, FP8RowLinear, Fp8RowwiseGatedMLP, Fp8RowwiseMLP, |
|
|
Fp8RowwiseRmsNorm, Int8SmoothQuantLinear, |
|
|
Int8SmoothQuantRowLinear, SmoothQuantAttention, |
|
|
SmoothQuantGatedMLP, SmoothQuantLayerNorm, SmoothQuantMLP, |
|
|
SmoothQuantRmsNorm, WeightOnlyGroupwiseQuantColumnLinear, |
|
|
WeightOnlyGroupwiseQuantRowLinear, |
|
|
WeightOnlyQuantColumnLinear, WeightOnlyQuantEmbedding, |
|
|
WeightOnlyQuantRowLinear) |
|
|
from .mode import W8A8_SQ_PLUGIN_LIST, QuantAlgo, QuantMode |
|
|
|
|
|
|
|
|
def quantize_layers( |
|
|
model, |
|
|
quant_config: QuantConfig, |
|
|
quant_map, |
|
|
preprocess_init_params=None, |
|
|
): |
|
|
exclude_modules = quant_config.exclude_modules or [ |
|
|
'*lm_head', |
|
|
'*router', |
|
|
'*vocab_embedding', |
|
|
'*position_embedding', |
|
|
'*block_embedding', |
|
|
'*shared_expert_gate', |
|
|
] |
|
|
|
|
|
for name, module, parent in model.named_modules_with_parent(): |
|
|
module_name = name.rsplit('.', 1)[-1] |
|
|
is_excluded = False |
|
|
for exclude_module in exclude_modules: |
|
|
if fnmatch.fnmatchcase(name, exclude_module): |
|
|
is_excluded = True |
|
|
|
|
|
|
|
|
if isinstance(module, MixtureOfExperts): |
|
|
init_params = get_init_params(module, MixtureOfExperts) |
|
|
init_params["quant_mode"] = QuantMode(0) |
|
|
original_layer = MixtureOfExperts(**init_params) |
|
|
if parent is not None: |
|
|
setattr(parent, module_name, original_layer) |
|
|
else: |
|
|
model = original_layer |
|
|
break |
|
|
if not is_excluded: |
|
|
quant_cls = None |
|
|
for cls in quant_map: |
|
|
if isinstance(module, cls): |
|
|
quant_cls = quant_map[cls] |
|
|
break |
|
|
|
|
|
if quant_cls is None: |
|
|
continue |
|
|
|
|
|
init_params = get_init_params(module, quant_cls) |
|
|
if "bias" in init_params: |
|
|
init_params["bias"] = init_params["bias"] is not None |
|
|
if isinstance(module, ColumnLinear): |
|
|
init_params[ |
|
|
"out_features"] = module.out_features * module.tp_size |
|
|
elif isinstance(module, RowLinear): |
|
|
init_params["in_features"] = module.in_features * module.tp_size |
|
|
if preprocess_init_params is not None: |
|
|
preprocess_init_params(init_params, name, module) |
|
|
quant_layer = quant_cls(**init_params) |
|
|
if parent is not None: |
|
|
setattr(parent, module_name, quant_layer) |
|
|
else: |
|
|
model = quant_layer |
|
|
|
|
|
setattr(model, 'quant_mode', quant_config.quant_mode) |
|
|
return model |
|
|
|
|
|
|
|
|
def weight_only_quantize(model, quant_config: QuantConfig): |
|
|
assert quant_config.quant_mode.is_weight_only() |
|
|
|
|
|
quant_map = { |
|
|
ColumnLinear: WeightOnlyQuantColumnLinear, |
|
|
RowLinear: WeightOnlyQuantRowLinear, |
|
|
Embedding: WeightOnlyQuantEmbedding, |
|
|
} |
|
|
|
|
|
def preprocess_init_params(init_params, name, module): |
|
|
init_params["quant_mode"] = quant_config.quant_mode |
|
|
if isinstance(module, ColumnLinear): |
|
|
module_name = name.rsplit('.', 1)[-1] |
|
|
init_params["transb"] = module_name == "lm_head" |
|
|
|
|
|
model = quantize_layers( |
|
|
model, |
|
|
quant_config, |
|
|
quant_map, |
|
|
preprocess_init_params, |
|
|
) |
|
|
return model |
|
|
|
|
|
|
|
|
def weight_only_groupwise_quantize(model, quant_config: QuantConfig): |
|
|
assert quant_config.quant_mode.is_weight_only() |
|
|
|
|
|
quant_map = { |
|
|
ColumnLinear: WeightOnlyGroupwiseQuantColumnLinear, |
|
|
RowLinear: WeightOnlyGroupwiseQuantRowLinear, |
|
|
} |
|
|
|
|
|
def preprocess_init_params(init_params, name, module): |
|
|
init_params["group_size"] = quant_config.group_size |
|
|
init_params["pre_quant_scale"] = quant_config.pre_quant_scale |
|
|
init_params["zero"] = quant_config.has_zero_point |
|
|
init_params[ |
|
|
"use_w4a8_awq"] = quant_config.quant_algo == QuantAlgo.W4A8_AWQ |
|
|
|
|
|
model = quantize_layers( |
|
|
model, |
|
|
quant_config, |
|
|
quant_map, |
|
|
preprocess_init_params, |
|
|
) |
|
|
return model |
|
|
|
|
|
|
|
|
def smooth_quantize_ootb( |
|
|
model, |
|
|
quant_config: QuantConfig, |
|
|
): |
|
|
quant_map = { |
|
|
ColumnLinear: Int8SmoothQuantLinear, |
|
|
RowLinear: Int8SmoothQuantRowLinear, |
|
|
} |
|
|
|
|
|
model = quantize_layers( |
|
|
model, |
|
|
quant_config, |
|
|
quant_map, |
|
|
) |
|
|
return model |
|
|
|
|
|
|
|
|
def smooth_quantize_plugin(model, quant_mode): |
|
|
quant_map = { |
|
|
RmsNorm: SmoothQuantRmsNorm, |
|
|
LayerNorm: SmoothQuantLayerNorm, |
|
|
GatedMLP: SmoothQuantGatedMLP, |
|
|
MLP: SmoothQuantMLP, |
|
|
Attention: SmoothQuantAttention, |
|
|
} |
|
|
for name, layer, parent in model.named_modules_with_parent(): |
|
|
layer_name = name.rsplit('.', 1)[-1] |
|
|
if layer_name in ['ln_f', 'ln_embed']: |
|
|
continue |
|
|
|
|
|
quant_cls = None |
|
|
for cls in quant_map: |
|
|
if isinstance(layer, cls): |
|
|
quant_cls = quant_map[cls] |
|
|
break |
|
|
|
|
|
if quant_cls is None: |
|
|
continue |
|
|
|
|
|
init_params = get_init_params(layer, quant_cls) |
|
|
init_params["quant_mode"] = quant_mode |
|
|
if isinstance(layer, Attention): |
|
|
init_params[ |
|
|
"num_attention_heads"] = layer.num_attention_heads * layer.tp_size |
|
|
quant_layer = quant_cls(**init_params) |
|
|
if parent is not None: |
|
|
setattr(parent, layer_name, quant_layer) |
|
|
else: |
|
|
model = quant_layer |
|
|
|
|
|
setattr(model, 'quant_mode', quant_mode) |
|
|
return model |
|
|
|
|
|
|
|
|
def smooth_quantize(model, quant_config: QuantConfig): |
|
|
assert quant_config.quant_mode.has_act_and_weight_quant() |
|
|
if quant_config.quant_algo in W8A8_SQ_PLUGIN_LIST: |
|
|
return smooth_quantize_plugin(model, quant_config.quant_mode) |
|
|
else: |
|
|
return smooth_quantize_ootb(model, quant_config) |
|
|
|
|
|
|
|
|
def fp8_quantize(model, quant_config: QuantConfig): |
|
|
assert quant_config.quant_mode.has_fp8_qdq() |
|
|
|
|
|
quant_map = { |
|
|
ColumnLinear: FP8Linear, |
|
|
RowLinear: FP8RowLinear, |
|
|
} |
|
|
|
|
|
model = quantize_layers( |
|
|
model, |
|
|
quant_config, |
|
|
quant_map, |
|
|
) |
|
|
return model |
|
|
|
|
|
|
|
|
def fp8_rowwise_quantize(model, quant_config: QuantConfig): |
|
|
assert quant_config.quant_mode.has_fp8_rowwise() |
|
|
|
|
|
quant_map = { |
|
|
RmsNorm: Fp8RowwiseRmsNorm, |
|
|
GatedMLP: Fp8RowwiseGatedMLP, |
|
|
MLP: Fp8RowwiseMLP, |
|
|
} |
|
|
for name, layer, parent in model.named_modules_with_parent(): |
|
|
layer_name = name.rsplit('.', 1)[-1] |
|
|
if layer_name in ['ln_f', 'ln_embed'] or "input_layernorm" in name: |
|
|
continue |
|
|
|
|
|
quant_cls = None |
|
|
for cls in quant_map: |
|
|
if isinstance(layer, cls): |
|
|
quant_cls = quant_map[cls] |
|
|
break |
|
|
|
|
|
if quant_cls is None: |
|
|
continue |
|
|
|
|
|
init_params = get_init_params(layer, quant_cls) |
|
|
init_params["quant_mode"] = quant_config.quant_mode |
|
|
quant_layer = quant_cls(**init_params, clamp_val=quant_config.clamp_val) |
|
|
if parent is not None: |
|
|
setattr(parent, layer_name, quant_layer) |
|
|
else: |
|
|
model = quant_layer |
|
|
|
|
|
setattr(model, 'quant_mode', quant_config.quant_mode) |
|
|
return model |
|
|
|
|
|
|
|
|
def kv_cache_quantize(model, quant_config: QuantConfig): |
|
|
assert quant_config.quant_mode.has_kv_cache_quant() |
|
|
for name, module in model.named_modules(): |
|
|
if isinstance(module, (Attention, SmoothQuantAttention)): |
|
|
module.kv_cache_scaling_factor = Parameter(shape=(1, ), |
|
|
dtype='float32') |
|
|
|
|
|
|
|
|
def quantize(model, quant_config: QuantConfig): |
|
|
quant_mode = quant_config.quant_mode |
|
|
|
|
|
if quant_mode.has_fp8_qdq(): |
|
|
model = fp8_quantize(model, quant_config) |
|
|
elif quant_mode.has_fp8_rowwise(): |
|
|
model = fp8_rowwise_quantize(model, quant_config) |
|
|
elif quant_mode.has_act_and_weight_quant(): |
|
|
model = smooth_quantize(model, quant_config) |
|
|
elif quant_mode.is_weight_only(): |
|
|
if quant_mode.has_per_group_scaling(): |
|
|
model = weight_only_groupwise_quantize(model, quant_config) |
|
|
else: |
|
|
model = weight_only_quantize(model, quant_config) |
|
|
|
|
|
if quant_mode.has_kv_cache_quant(): |
|
|
model = kv_cache_quantize(model, quant_config) |
|
|
|
|
|
return model |
|
|
|
