diff --git a/angelslim/compressor/quant/core/config.py b/angelslim/compressor/quant/core/config.py index dc1da02c..72428d58 100644 --- a/angelslim/compressor/quant/core/config.py +++ b/angelslim/compressor/quant/core/config.py @@ -174,6 +174,37 @@ def __init__(self, config, global_config=None): "dequant_to_bf16": quantization_args.quant_method.get("dequant_to_bf16", False), "actorder": quantization_args.quant_method.get("actorder", True), } + elif "nvfp4_gptq" in self.quant_algo or "nvfp4_gptaq" in self.quant_algo: + # NVFP4 weight format optimized with the GPTQ error-compensation + # loop (weight-only). Shares the GPTQ runner with int4_gptq; the + # weight_format flag routes the per-format quant primitives. The + # NVFP4 micro-scaling block is fixed at 16. act-order may stay on: + # the GPTQ runner folds the input permutation into the previous + # layer's output channels, so the 16-element block stays physically + # contiguous along K even when columns are reordered. + self.act_observer = None + self.weight_observer = None + self.kv_cache_observer = None + block_size = ( + 16 + if quantization_args.quant_method["group_size"] == -1 + else quantization_args.quant_method["group_size"] + ) + self.quant_algo_info = { + # ``group_size`` is read by the GPTQ runner; ``block_size`` by + # the NVFP4 packing path. Keep them in sync. + "group_size": block_size, + "block_size": block_size, + "ignore_layers": quantization_args.ignore_layers, + "weight_format": "nvfp4", + "checkpoint_format": "nvfp4", + "dequant_to_bf16": quantization_args.quant_method.get("dequant_to_bf16", False), + "actorder": quantization_args.quant_method.get("actorder", True), + "share_gate_up_weight_scale_2": quantization_args.quant_method.get( + "share_gate_up_weight_scale_2", True + ), + "w": f"nvfp4_{weight_quant_method}", + } elif "nvfp4" in self.quant_algo: is_weight_only = "weight_only" in self.quant_algo is_dynamic = "dynamic" if "dynamic" in self.quant_algo else "static" diff --git a/angelslim/compressor/quant/modules/gptq/gptq.py b/angelslim/compressor/quant/modules/gptq/gptq.py index 46108e4a..e822e39b 100644 --- a/angelslim/compressor/quant/modules/gptq/gptq.py +++ b/angelslim/compressor/quant/modules/gptq/gptq.py @@ -24,9 +24,18 @@ from huggingface_hub import save_torch_state_dict from tqdm import tqdm -from .....utils import decide_device_for_distributed, find_layers, print_info +from .....utils import ( + decide_device_for_distributed, + find_layers, + find_parent_layer_and_sub_name, + print_info, +) from ...modules.catcher import Catcher -from ...modules.helper_layer import GPTQQuantLinear +from ...modules.helper_layer import ( + GPTQQuantLinear, + NVFP4QDQModule, + compute_nvfp4_weight_scale_2, +) from .gptaq_module import GPTAQModule from .gptq_module import GPTQModule @@ -56,9 +65,21 @@ def __init__(self, model, seq_length=2048, hidden_size=2560, sym=True, actorder= self.quant_bits = self.model.quant_config.quant_bit self.group_size = self.model.quant_config.quant_algo_info["group_size"] self.ignore_layers = self.model.quant_config.quant_algo_info["ignore_layers"] + # Weight number format backend: "int4" (default) or "nvfp4". + self.weight_format = self.model.quant_config.quant_algo_info.get("weight_format", "int4") + self.block_size = self.model.quant_config.quant_algo_info.get( + "block_size", self.group_size + ) self.dequant_to_bf16 = bool( self.model.quant_config.quant_algo_info.get("dequant_to_bf16", False) ) + # NVFP4 only: share one per-tensor (level-2) weight scale across each + # fused gate/up (and q/k/v) group, matching the single per-tensor scale + # a fused-GEMM deployment applies. Default on; set False to keep the + # legacy per-layer scale_2. + self.share_gate_up_weight_scale_2 = bool( + self.model.quant_config.quant_algo_info.get("share_gate_up_weight_scale_2", True) + ) self.percdamp = 0.01 self.sym = sym self.actorder = actorder @@ -66,6 +87,8 @@ def __init__(self, model, seq_length=2048, hidden_size=2560, sym=True, actorder= self.hidden_size = hidden_size self.dtype = next(iter(self.layers.parameters())).dtype self.quantizers = {} + # NVFP4: per-tensor (level-2) scale per quantized layer. + self.nvfp4_weight_scales_2 = {} self.gptq = {} self.quant_algo = self.model.quant_config.quant_algo self.native_inp_caches = {} @@ -110,12 +133,24 @@ def _move_layer_to_device(self, layer, device): return layer.to(device) def get_actorder_prev_names(self, name, subset): + # Enable act-order for *_proj layers whose input permutation can be + # folded into the output channels of the preceding gate_proj/up_proj. + # This keeps the (NVFP4 / int4) group blocks physically contiguous + # along K even after column reordering, so no g_idx is needed. if not name.endswith("down_proj"): return [] prefix = name[: -len("down_proj")] prev_names = [f"{prefix}gate_proj", f"{prefix}up_proj"] - return [prev_name for prev_name in prev_names if prev_name in subset] + if any(prev_name not in subset for prev_name in prev_names): + return [] + + current_layer = subset[name] + for prev_name in prev_names: + prev_layer = subset[prev_name] + if prev_layer.weight.shape[0] != current_layer.weight.shape[1]: + return [] + return prev_names def reorder_quantizer_output_channels(self, quantizer_name, input_perm): if quantizer_name not in self.quantizers: @@ -123,11 +158,20 @@ def reorder_quantizer_output_channels(self, quantizer_name, input_perm): scale, zero = self.quantizers[quantizer_name] scale_perm = input_perm.to(scale.device) - zero_perm = input_perm.to(zero.device) - self.quantizers[quantizer_name] = ( - scale.index_select(0, scale_perm), - zero.index_select(0, zero_perm), - ) + # Reorder per-output-channel scale rows. For nvfp4 the second element + # is the per-tensor (scalar) weight_scale_2, which is invariant to + # output-channel permutation and must NOT be index_select'd. + if self.weight_format == "nvfp4": + self.quantizers[quantizer_name] = ( + scale.index_select(0, scale_perm), + zero, + ) + else: + zero_perm = input_perm.to(zero.device) + self.quantizers[quantizer_name] = ( + scale.index_select(0, scale_perm), + zero.index_select(0, zero_perm), + ) def fold_input_permutation_into_prev_layers( self, @@ -257,11 +301,70 @@ def _filter_distributed_expert_subset(self, layer, subset): if skipped > 0: print_info( - "Filter non-local expert layers for GPTQ: " - f"local experts [{start}, {end}), skipped {skipped} layer(s)." + f"Filter non-local expert layers for GPTQ: local experts " + f"[{start}, {end}), skipped {skipped} layer(s)." ) return filtered_subset + @staticmethod + def _nvfp4_scale_2_group_key(name): + """Group key for layers that must share one NVFP4 per-tensor scale_2. + + Deployment fuses each expert's gate_proj+up_proj (and attention's + q/k/v) into a single GEMM that can apply only one per-tensor (level-2) + scale across the fused weight. GPTQ must therefore quantize every member + of such a group against the *same* weight_scale_2, otherwise the half of + the fused weight whose scale_2 was overwritten at deploy time is + dequantized with the wrong per-tensor scale -> a systematic multiplicative + bias that occasionally flips a token. + + Returns (group_prefix, member_tag) for fusible layers, else None. + """ + for suffix in ("gate_proj", "up_proj"): + if name.endswith(suffix): + return name[: -len(suffix)], "gate_up" + for suffix in ("q_proj", "k_proj", "v_proj"): + if name.endswith(suffix): + return name[: -len(suffix)], "qkv" + return None + + def _assign_shared_nvfp4_weight_scale_2(self): + """Pre-pass: give fused gate/up (and q/k/v) groups one shared scale_2. + + Runs BEFORE fasterquant so GPTQ compensates against the exact grid that + deployment will use. The shared scale_2 is derived from the group-wide + amax (max of each member's amax). Because the shared amax is >= each + member's own amax, every per-block E4M3 scale only shrinks, so no block + scale overflows the E4M3 range that single-layer scaling would have hit. + """ + groups = {} + for name in self.gptq: + if any(ignore in name for ignore in self.ignore_layers): + continue + key = self._nvfp4_scale_2_group_key(name) + if key is None: + continue + groups.setdefault(key, []).append(name) + + for (prefix, tag), names in groups.items(): + if len(names) < 2: + # Nothing to fuse with; let fasterquant compute scale_2 itself. + continue + group_amax = None + for name in names: + # self.gptq[name].w is the cloned fp32 weight before compensation; + # GPTQ compensation preserves amax magnitude, so this is the same + # amax fasterquant would otherwise use per-layer. + w_amax = self.gptq[name].w.abs().amax().to(torch.float32) + group_amax = w_amax if group_amax is None else torch.maximum(group_amax, w_amax) + shared_scale_2 = compute_nvfp4_weight_scale_2(group_amax) + for name in names: + self.gptq[name].weight_scale_2 = shared_scale_2.to(self.gptq[name].dev) + print_info( + f"NVFP4 shared weight_scale_2 for {tag} group '{prefix}*': " + f"{shared_scale_2.item():.6g} across {len(names)} layers" + ) + @torch.no_grad() def run(self, dataloader): for model_module in self.layers: @@ -329,7 +432,12 @@ def run(self, dataloader): self.native_inp_caches[name] = [] self.gptq[name] = GPTAQModule(subset[name], quant_bits=self.quant_bits) else: - self.gptq[name] = GPTQModule(subset[name], quant_bits=self.quant_bits) + self.gptq[name] = GPTQModule( + subset[name], + quant_bits=self.quant_bits, + weight_format=self.weight_format, + block_size=self.block_size, + ) def pre_process_fwd_hook(layer_name): def tmp(_, inp, out): @@ -380,6 +488,9 @@ def tmp(_, inp, out): for h in handles: h.remove() + if self.weight_format == "nvfp4" and self.share_gate_up_weight_scale_2: + self._assign_shared_nvfp4_weight_scale_2() + for name in self.gptq: if any(ignore in name for ignore in self.ignore_layers): continue @@ -389,8 +500,8 @@ def tmp(_, inp, out): and self.gptq[name].nsamples == 0 ): print_info( - f"Skip {name} because no calibration samples were routed " - "to this local expert layer." + f"Skip {name} because no calibration samples were " + f"routed to this local expert layer." ) self.gptq[name].free() continue @@ -403,10 +514,23 @@ def tmp(_, inp, out): actorder=actorder, sym=self.sym, ) - self.quantizers[f"{self.layers_block_name}.{i}.{name}"] = ( - scale, - zero, - ) + quant_name = f"{self.layers_block_name}.{i}.{name}" + if self.weight_format == "nvfp4": + # NVFP4 uses block_size=16 -> 8x more scale groups than + # int4 (gs=128). Keep the per-layer scales on CPU so they + # do not accumulate on GPU across all 80 layers (the cause + # of the layer-16 OOM). They are only needed at convert(). + self.quantizers[quant_name] = ( + scale.cpu(), + zero.cpu(), + ) + # ``zero`` slot carries the per-tensor level-2 scale. + self.nvfp4_weight_scales_2[quant_name] = zero.cpu() + else: + self.quantizers[quant_name] = ( + scale, + zero, + ) self.fold_input_permutation_into_prev_layers( i, subset, @@ -511,12 +635,56 @@ def _convert_llm(self): force_layer_back_to_cpu=True, ) + def _convert_nvfp4(self): + """Insert NVFP4QDQModule for each quantized layer (real packing). + + GPTQ has already written the compensated, E2M1-snapped weights back to + ``layer.weight`` (bf16). NVFP4QDQModule re-quantizes them into the + packed modelopt NVFP4 format (packed uint4 + E4M3 block scale + FP32 + weight_scale_2) using the scales GPTQ produced, so the deployed grid is + exactly what GPTQ optimized against. + """ + model = self.model.model + model.cpu() + print_info("Packing NVFP4 model...") + layers = find_layers(model, layers=self.model.observer_layer_classes) + + with tctl.threadpool_limits(limits=1): + pbar = tqdm(self.quantizers.keys(), leave=True) + for name in pbar: + pbar.set_description(f"Packing {name}...", refresh=True) + if name not in layers: + continue + sub_layer = layers[name].cpu() + block_scale_e4m3, _zero = self.quantizers[name] + weight_scale_2 = self.nvfp4_weight_scales_2[name] + + qdq_module = NVFP4QDQModule( + weight=sub_layer.weight, + weight_scale=block_scale_e4m3.cpu(), + weight_scale_2=weight_scale_2.cpu(), + bias=sub_layer.bias, + block_size=self.block_size, + input_scale=None, + ) + parent_layer, sub_name = find_parent_layer_and_sub_name(model, name) + setattr(parent_layer, sub_name, qdq_module) + print_info("NVFP4 model packed.") + def convert(self): """ Saves scales and inserts QDQ modules. """ print_info("Start convert model...") - if self.dequant_to_bf16: + if self.weight_format == "nvfp4": + if self.dequant_to_bf16: + print_info( + "dequant_to_bf16=True: skip NVFP4 packing, keep " + "fake-quantized bf16 weights." + ) + else: + self._convert_nvfp4() + elif self.dequant_to_bf16: print_info( "dequant_to_bf16=True: skip int4 packing, keep fake-quantized bf16 weights." ) @@ -600,6 +768,14 @@ def forward(self, x): except Exception: self.model.model.config.quantization_config = None self.model.model.config.torch_dtype = "bfloat16" + elif self.weight_format == "nvfp4": + # Packed modelopt-style NVFP4 checkpoint (two-level scaling). + self.model.model.config.quantization_config = { + "quant_method": "nvfp4", + "kv_cache_scheme": None, + "group_size": self.block_size, + "exclude_modules": self.ignore_layers, + } else: self.model.model.config.quantization_config = { "bits": self.quant_bits, diff --git a/angelslim/compressor/quant/modules/gptq/gptq_module.py b/angelslim/compressor/quant/modules/gptq/gptq_module.py index 3bf18ffe..2ce81bbf 100644 --- a/angelslim/compressor/quant/modules/gptq/gptq_module.py +++ b/angelslim/compressor/quant/modules/gptq/gptq_module.py @@ -19,18 +19,26 @@ from .....utils import get_tensor_item, print_info from ...core import compute_scales_with_zero +from ..helper_layer import ( + compute_nvfp4_block_scale, + compute_nvfp4_weight_scale_2, + nvfp4_quant_dequant, +) __all__ = ["GPTQModule"] class GPTQModule: - def __init__(self, layer, quant_bits=4): + def __init__(self, layer, quant_bits=4, weight_format="int4", block_size=16): """ GPTQ quantization wrapper for neural network layers. Args: layer: Full-precision torch.nn.Module to quantize (Linear) quant_bits: Quantization bitwidth (2-8 bits, default=4) + weight_format: "int4" (default, uniform) or "nvfp4" (E2M1 grid + + two-level scale). Routes compute_quant_params / quant_dequant. + block_size: NVFP4 micro-scaling block size (nvfp4 only). """ super(GPTQModule, self).__init__() self.layer = layer @@ -41,6 +49,10 @@ def __init__(self, layer, quant_bits=4): self.h = torch.zeros((self.columns, self.columns), device=self.dev, dtype=torch.float32) self.nsamples = 0 self.quant_bits = quant_bits + self.weight_format = weight_format + self.block_size = block_size + # Per-tensor (level-2) NVFP4 scale, set at the start of fasterquant. + self.weight_scale_2 = None def add_batch(self, inp, out): # Handle 4D input (e.g., Conv2d or multi-head attention internals) @@ -62,9 +74,16 @@ def add_batch(self, inp, out): self.h += inp.matmul(inp.t()) def compute_quant_params(self, x, bits, sym): + if self.weight_format == "nvfp4": + # Per-block effective scale (block_scale_e4m3 * weight_scale_2). + # weight_scale_2 is computed once in fasterquant before this runs. + eff_scale = compute_nvfp4_block_scale(x, self.weight_scale_2) + return eff_scale, torch.zeros_like(eff_scale) return compute_scales_with_zero(x, bits=bits, sym=sym) def quant_dequant(self, x, weight_scale, weight_zero): + if self.weight_format == "nvfp4": + return nvfp4_quant_dequant(x, weight_scale) maxq = torch.tensor(2**self.quant_bits - 1, device=x.device) q = torch.clamp(torch.round(x / weight_scale) + weight_zero, 0, maxq) return weight_scale * (q - weight_zero) @@ -91,6 +110,18 @@ def fasterquant( hessian[dead, dead] = 1 w_weight[:, dead] = 0 + # NVFP4: per-tensor level-2 scale, computed once from the (dead-zeroed) + # weights before any per-block scale. GPTQ compensation does not change + # the amax magnitude, so this stays valid for the whole layer. + # + # If ``self.weight_scale_2`` + # was already set externally (e.g. a shared + # gate/up or qkv level-2 scale injected by GPTQ.run so fused-GEMM + # deployment uses one per-tensor scale across the group), keep it and do + # NOT recompute from this layer's amax alone. + if self.weight_format == "nvfp4" and self.weight_scale_2 is None: + self.weight_scale_2 = compute_nvfp4_weight_scale_2(w_weight.abs().amax()) + scale = [] zero = [] now_idx = 1 @@ -200,6 +231,27 @@ def fasterquant( zero = torch.zeros_like(weight_scale) scale = torch.cat(scale, dim=1) zero = torch.cat(zero, dim=1) + + if self.weight_format == "nvfp4": + # ``scale`` currently holds the effective scale (block_e4m3 * + # weight_scale_2). Recover the stored E4M3 block scale and hand the + # per-tensor level-2 scale back via the second return value. Move + # both to CPU so the per-layer scales never pile up on GPU across + # all transformer layers. + block_scale_e4m3 = (scale / self.weight_scale_2).to(torch.float8_e4m3fn).cpu() + weight_scale_2 = self.weight_scale_2.detach().clone().cpu() + losses = losses.cpu() + q_weight = q_weight.cpu() + w_weight = w_weight.cpu() + hessian = hessian.cpu() + hinv = hinv.cpu() + del losses, q_weight, w_weight, hessian, hinv + self.w = self.w.cpu() + del self.w + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return block_scale_e4m3, weight_scale_2, input_perm + losses = losses.cpu() q_weight = q_weight.cpu() w_weight = w_weight.cpu() diff --git a/angelslim/compressor/quant/modules/helper_layer.py b/angelslim/compressor/quant/modules/helper_layer.py index c586a1da..c7e06d57 100644 --- a/angelslim/compressor/quant/modules/helper_layer.py +++ b/angelslim/compressor/quant/modules/helper_layer.py @@ -692,6 +692,93 @@ def forward(self, x): return output +# --------------------------------------------------------------------------- +# NVFP4 module-level primitives (E2M1 grid + two-level scale). +# +# These mirror the mxfp4_* helpers but implement NVFP4's two-level scaling: +# per-block scale (FP8 E4M3) * per-tensor scale_2 (FP32). +# The E2M1 weight grid is identical to MXFP4, so the cast/round logic matches +# NVFP4QDQModule._cast_fp4 bit-for-bit. GPTQ's inner loop calls these so that +# the grid it compensates against is exactly the one used at packing time. +# --------------------------------------------------------------------------- +NVFP4_E2M1_BOUNDS = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0]) +NVFP4_E2M1_VALUES = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6]) +NVFP4_E2M1_VALUES_ON_DEVICE = {} +NVFP4_E2M1_MAX = 6.0 +NVFP4_E4M3_MAX = 448.0 + + +def nvfp4_get_e2m1_values(device): + if device not in NVFP4_E2M1_VALUES_ON_DEVICE: + NVFP4_E2M1_VALUES_ON_DEVICE[device] = NVFP4_E2M1_VALUES.to(device) + return NVFP4_E2M1_VALUES_ON_DEVICE[device] + + +def nvfp4_cast_to_e2m1(weight: torch.Tensor): + """Round a (pre-scaled) tensor to nearest E2M1 grid, return uint4 codes. + + Matches NVFP4QDQModule._cast_fp4 exactly (round-half-to-even on ties). + """ + device = weight.device + bounds = NVFP4_E2M1_BOUNDS.to(device) + + mask = torch.tensor([0, 1, 0, 1, 0, 1, 0], dtype=torch.uint8, device=device) + mask = mask.expand([*list(weight.shape), 7]) + + sign_bit = (weight < 0).to(torch.uint8) + weight_abs = weight.abs() + ord_idx = torch.searchsorted(bounds, weight_abs, out_int32=True).to(torch.uint8) + round_up = torch.any((weight_abs.unsqueeze(-1) == bounds) * mask, dim=-1) + return (sign_bit * 0b1000 + ord_idx + round_up).to(torch.uint8) + + +def nvfp4_dequantize_e2m1(code: torch.Tensor, dtype: torch.dtype = torch.float32): + return nvfp4_get_e2m1_values(code.device)[(code.to(torch.long) & 0x0F)].to(dtype) + + +def compute_nvfp4_weight_scale_2(weight_amax: torch.Tensor) -> torch.Tensor: + """Per-tensor (level-2) FP32 scale: amax / 6 / 448. + + Computed once per layer; does not change during GPTQ compensation. + """ + scale2 = weight_amax.float() / NVFP4_E2M1_MAX / NVFP4_E4M3_MAX + return torch.where(scale2 > 0, scale2, torch.ones_like(scale2)) + + +def compute_nvfp4_block_scale( + x_block: torch.Tensor, + weight_scale_2: torch.Tensor, + keep_high_precision: bool = False, +) -> torch.Tensor: + """Per-block (level-1) effective scale, including the E4M3 round-trip. + + Args: + x_block: weight block, last dim is the block (e.g. [rows, block_size]). + weight_scale_2: per-tensor scale from compute_nvfp4_weight_scale_2. + Returns: + Effective scale (block_scale_e4m3 * scale_2), last dim kept as 1, so + that GPTQ compensates against the exact stored scale. + """ + per_block_amax = x_block.abs().amax(dim=-1, keepdim=True).float() + per_block_scale = per_block_amax / NVFP4_E2M1_MAX + q_per_block_scale = per_block_scale / weight_scale_2 + q_per_block_scale = torch.where( + per_block_scale > 0, q_per_block_scale, torch.ones_like(q_per_block_scale) + ) + if not keep_high_precision: + finfo = torch.finfo(torch.float8_e4m3fn) + q_per_block_scale = q_per_block_scale.clamp(min=finfo.min, max=finfo.max) + q_per_block_scale = q_per_block_scale.to(torch.float8_e4m3fn).float() + return q_per_block_scale * weight_scale_2 + + +def nvfp4_quant_dequant(x: torch.Tensor, eff_scale: torch.Tensor): + """Fake-quantize x onto the E2M1 grid and dequantize back (float32).""" + scaled = x.float() / eff_scale + code = nvfp4_cast_to_e2m1(scaled) + return nvfp4_dequantize_e2m1(code, dtype=torch.float32) * eff_scale + + class NVFP4QDQModule(torch.nn.Module): def __init__( self, diff --git a/angelslim/compressor/quant/ptq.py b/angelslim/compressor/quant/ptq.py index ec88ccdf..37248947 100644 --- a/angelslim/compressor/quant/ptq.py +++ b/angelslim/compressor/quant/ptq.py @@ -55,12 +55,20 @@ def __init__(self, model, slim_config=None): # trasform first, then run quantization self.transform_runner.run() - if "fp8" in self.quant_algo or "int8" in self.quant_algo or "nvfp4" in self.quant_algo: + # NVFP4 weight-only GPTQ (``nvfp4_gptq``) is driven by the GPTQ runner + # and computes its scales directly from weights, so it must NOT install + # the activation/weight observer hook even though "nvfp4" is in algo. + is_gptq = "gptq" in self.quant_algo or "gptaq" in self.quant_algo + if ( + "fp8" in self.quant_algo + or "int8" in self.quant_algo + or ("nvfp4" in self.quant_algo and not is_gptq) + ): # Add ptq observer hook self.ptq_hook = PTQHook(self.quant_model) self.ptq_hook.apply_hook() - if "gptq" in self.quant_algo or "gptaq" in self.quant_algo: + if is_gptq: max_seq_length = self.quant_model.quant_config.max_seq_length hidden_size = self.quant_model.quant_config.hidden_size self.gptq = GPTQ( diff --git a/angelslim/models/llm/hunyuan_v3_moe.py b/angelslim/models/llm/hunyuan_v3_moe.py index f44bdaf0..dda0215c 100644 --- a/angelslim/models/llm/hunyuan_v3_moe.py +++ b/angelslim/models/llm/hunyuan_v3_moe.py @@ -271,7 +271,12 @@ def from_pretrained( use_cache=False, using_multi_nodes=False, ): - attn_implementation = "eager" + # Attention implementation. Default "eager" so KV-cache / activation + # calibration can hook the explicit attention weights. For weight-only + # GPTQ (no attention hooks) "eager" materializes a [heads, seq, seq] + # fp32 matrix (16 GiB at seq=8192) and OOMs; set ANGELSLIM_ATTN_IMPL=sdpa + # to use the flash-style SDPA path that never materializes it. + attn_implementation = os.environ.get("ANGELSLIM_ATTN_IMPL", "eager") torch_dtype = torch.bfloat16 if is_deepspeed_zero3_enabled(): _patch_hyv3_router_for_zero3() @@ -611,9 +616,8 @@ def _enable_expert_parallel(self): ), f"num_experts {num_experts} must be divisible by world_size {self.world_size}" print_info( - "Enable HYV3 expert parallel: " - f"rank={self.rank}, world_size={self.world_size}, " - f"num_experts={num_experts}" + f"Enable HYV3 expert parallel: rank={self.rank}, " + f"world_size={self.world_size}, num_experts={num_experts}" ) for layer_idx, layer in enumerate(self.model.model.layers): moe_module = getattr(layer, "mlp", None) diff --git a/angelslim/utils/config_parser.py b/angelslim/utils/config_parser.py index 3be83400..3378395a 100644 --- a/angelslim/utils/config_parser.py +++ b/angelslim/utils/config_parser.py @@ -50,6 +50,8 @@ class QuantizationMethod(str, Enum): INT4_GPTAQ = "int4_gptaq" NVFP4 = "nvfp4" NVFP4_WEIGHT_ONLY = "nvfp4_weight_only" + NVFP4_GPTQ = "nvfp4_gptq" + NVFP4_GPTAQ = "nvfp4_gptaq" W4A8_INT8 = "w4a8i8" diff --git a/configs/Hy3/ptq/nvfp4_gptq/hunyuanv3_a20b_nvfp4_gptq.yaml b/configs/Hy3/ptq/nvfp4_gptq/hunyuanv3_a20b_nvfp4_gptq.yaml new file mode 100644 index 00000000..8e17e7ca --- /dev/null +++ b/configs/Hy3/ptq/nvfp4_gptq/hunyuanv3_a20b_nvfp4_gptq.yaml @@ -0,0 +1,47 @@ +# Global configuration of pipeline +global: + save_path: ./output + +# Simplified Configuration for LLM compression +model: + name: HYV3MoE + model_path: tencent/Hy3-preview + low_cpu_mem_usage: true + use_cache: false + torch_dtype: auto + device_map: cpu + enable_expert_parallel: true + +# Compression configuration +compression: + name: PTQ + quantization: + name: nvfp4_gptq # Supported: fp8_static, fp8_dynamic, int4_awq, int4_gptq, nvfp4, nvfp4_gptq + bits: 4 + quant_method: + weight: "per-group" + group_size: 16 # NVFP4 micro-scaling block (fixed at 16; -1 also resolves to 16) + dequant_to_bf16: false + actorder: true + share_gate_up_weight_scale_2: true # share one per-tensor scale_2 across each fused gate/up (q/k/v) group; matches fused-GEMM deployment + ignore_layers: # Skip quantization for these layers + - "lm_head" + - "self_attn.q_proj" + - "self_attn.k_proj" + - "self_attn.v_proj" + - "self_attn.o_proj" + - "mlp.router.gate" + - "mlp.gate_proj" + - "mlp.up_proj" + - "mlp.down_proj" + - "mlp.shared_experts.gate_proj" + - "mlp.shared_experts.up_proj" + - "mlp.shared_experts.down_proj" + +# Dataset for calibration +dataset: + name: TextDataset + data_path: ./dataset/sharegpt_gpt4/sharegpt_gpt4_256.jsonl + max_seq_length: 2048 + num_samples: 128 + batch_size: 1 diff --git a/configs/hunyuan/nvfp4_weight_only/hunyuan_a20b_nvfp4_weight_only.yaml b/configs/Hy3/ptq/nvfp4_weight_only/hunyuan_a13b_nvfp4_weight_only.yaml similarity index 100% rename from configs/hunyuan/nvfp4_weight_only/hunyuan_a20b_nvfp4_weight_only.yaml rename to configs/Hy3/ptq/nvfp4_weight_only/hunyuan_a13b_nvfp4_weight_only.yaml diff --git a/configs/Hy3/ptq/nvfp4_weight_only/hunyuan_a20b_nvfp4_weight_only.yaml b/configs/Hy3/ptq/nvfp4_weight_only/hunyuan_a20b_nvfp4_weight_only.yaml new file mode 100644 index 00000000..edd8697a --- /dev/null +++ b/configs/Hy3/ptq/nvfp4_weight_only/hunyuan_a20b_nvfp4_weight_only.yaml @@ -0,0 +1,38 @@ +# Global configuration of pipeline +global: + save_path: ./output + +# Simplified Configuration for LLM compression +model: + name: HYV3MoE + model_path: tencent/Hy3-preview + trust_remote_code: true + low_cpu_mem_usage: true + use_cache: false + torch_dtype: auto + device_map: auto + +# Compression configuration +compression: + name: PTQ + quantization: + name: nvfp4_weight_only + bits: 4 + cpu_convert: true + quant_method: + weight: "per-group" + group_size: 16 + dequant_to_bf16: false + ignore_layers: # Skip quantization for these layers + - "lm_head" + - "self_attn.q_proj" + - "self_attn.k_proj" + - "self_attn.v_proj" + - "self_attn.o_proj" + - "mlp.router.gate" + - "mlp.gate_proj" + - "mlp.up_proj" + - "mlp.down_proj" + - "mlp.shared_mlp.gate_proj" + - "mlp.shared_mlp.up_proj" + - "mlp.shared_mlp.down_proj" diff --git a/tools/run.py b/tools/run.py index fd20b58b..f17a2c84 100644 --- a/tools/run.py +++ b/tools/run.py @@ -28,6 +28,10 @@ from angelslim.utils import get_yaml_prefix_simple, print_info # noqa: E402 from angelslim.utils.config_parser import SlimConfigParser, print_config # noqa: E402 +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + def get_args(): parser = argparse.ArgumentParser(description="AngelSlim")