refine awq, bugfix

auto_round/algorithms/transforms/awq/config.py

@@ -47,6 +47,9 @@ def __init__(
         clip_n_grid: int = 20,
         clip_max_shrink: float = 0.5,
         clip_n_sample_token: int = 512,
+        smooth_seqlen: int = 512,
+        smooth_batch_size: int | None = None,
+        skip_moe: bool = True,
         mappings: list[dict] | None = None,
         **kwargs,
     ):
@@ -104,6 +107,22 @@ def __init__(
             clip_n_sample_token: Maximum number of calibration tokens used per
                 balance layer when searching the clip threshold (subsampled to
                 bound memory).
+            smooth_seqlen: Maximum sequence length (number of tokens) used per
+                calibration sample during the AWQ scale grid search. Defaults to
+                ``512``. Set a larger positive integer to use longer sequences,
+                or a value ``<= 0`` to disable truncation entirely.
+            smooth_batch_size: Optional microbatch size used when replaying AWQ
+                parent modules during scale grid search. Smaller values reduce
+                peak VRAM while preserving the AWQ parent-output loss, at the
+                cost of more parent forward calls. ``None`` or ``<= 0`` replays
+                the cached calibration batch as-is.
+            skip_moe: Whether to exclude routed MoE experts from AWQ smoothing.
+                When True, balance layers belonging to routed experts (module
+                names matching ``.experts.<N>.``) are dropped from the resolved
+                mappings, so AWQ only smooths attention and dense/shared paths
+                and leaves routed experts to the downstream block quantizer. This
+                has no effect on dense models and is ignored when explicit
+                ``mappings`` are provided.
             mappings: Optional explicit AWQ smooth/balance mappings. Each
                 item should contain ``smooth_layer`` and
                 ``balance_layers`` entries. If None, mappings are inferred
@@ -135,9 +154,14 @@ def __init__(
             raise ValueError(f"`clip_max_shrink` must be in (0, 1), got {clip_max_shrink!r}")
         if clip_n_sample_token is None or clip_n_sample_token < 1:
             raise ValueError(f"`clip_n_sample_token` must be a positive integer, got {clip_n_sample_token!r}")
+        if smooth_batch_size is not None and smooth_batch_size < 0:
+            raise ValueError(f"`smooth_batch_size` must be a non-negative integer or None, got {smooth_batch_size!r}")
         self.clip_n_grid = clip_n_grid
         self.clip_max_shrink = clip_max_shrink
         self.clip_n_sample_token = clip_n_sample_token
+        self.smooth_seqlen = smooth_seqlen
+        self.smooth_batch_size = smooth_batch_size
+        self.skip_moe = skip_moe
         self.mappings = mappings
         self.infer_bs_coeff = 1
         self.batch_dim = None
@@ -159,6 +183,8 @@ def __repr__(self) -> str:
             f"AWQConfig(duo_scaling={self.duo_scaling!r}, n_grid={self.n_grid}, "
             f"smooth_iters={self.smooth_iters}, "
             f"apply_clip={self.apply_clip}, clip_as_init={self.clip_as_init}, "
+            f"smooth_seqlen={self.smooth_seqlen}, smooth_batch_size={self.smooth_batch_size}, "
+            f"skip_moe={self.skip_moe}, "
             f"bits={self.bits}, group_size={self.group_size}, sym={self.sym}, "
             f"mappings={'<explicit>' if self.mappings else 'auto'})"
         )

auto_round/algorithms/transforms/awq/mappings.py

@@ -70,6 +70,10 @@ class ResolvedMapping:
     activation_hook_target: str | None = None
 
 
+# Matches routed MoE expert modules, e.g. "...mlp.experts.3.gate_proj".
+_ROUTED_EXPERT_RE = re.compile(r"\.experts\.\d+\.")
+
+
 # ── Mapping definitions ─────────────────────────
 # Reference: vllm-project/llm-compressor src/llmcompressor/modifiers/awq/mappings.py
 
@@ -173,6 +177,15 @@ class ResolvedMapping:
     AWQMapping(r"up_proj$", [r"down_proj$"]),
 ]
 
+_bagel_mappings = [
+    AWQMapping(
+        r"input_layernorm$",
+        [r"\.self_attn\.q_proj$", r"\.self_attn\.k_proj$", r"\.self_attn\.v_proj$"],
+    ),
+    AWQMapping(r"post_attention_layernorm$", [r"\.mlp\.gate_proj$", r"\.mlp\.up_proj$"]),
+    AWQMapping(r"\.mlp\.up_proj$", [r"\.mlp\.down_proj$"]),
+]
+
 # ── Model class name → mappings registry ──────────────────────────────────────
 # Aligned with llm-compressor AWQ_MAPPING_REGISTRY (llmcompressor v0.10.0).
 # Models not in this registry fall back to default_mappings.
@@ -220,6 +233,10 @@ class ResolvedMapping:
     # Other models using default mappings
     "SeedOssForCausalLM": default_mappings,
     "Ernie4_5_MoeForCausalLM": default_mappings,
+    # BAGEL wraps a Qwen2 language model and carries parallel *_moe_gen modules
+    # for image generation. Keep AWQ smoothing on the normal text path only.
+    "BagelForQuantization": _bagel_mappings,
+    "BagelForConditionalGeneration": _bagel_mappings,
 }
 
 
@@ -331,6 +348,14 @@ def _build_hybrid_attention_mappings(model: torch.nn.Module) -> list[AWQMapping]
 
     layer_types, num_layers = result
 
+    if len(layer_types) < num_layers:
+        logger.warning(
+            "Hybrid attention model config has num_hidden_layers=%d but only %d layer_types entries. Falling back.",
+            num_layers,
+            len(layer_types),
+        )
+        return None
+
     full_indices = [i for i in range(num_layers) if layer_types[i] == "full_attention"]
     linear_indices = [i for i in range(num_layers) if layer_types[i] == "linear_attention"]
 
@@ -435,6 +460,7 @@ def _get_mappings_for_model(model: torch.nn.Module) -> list[AWQMapping]:
 def resolve_mappings(
     model: torch.nn.Module,
     user_mappings: list[dict] | None = None,
+    skip_moe: bool = False,
 ) -> list[ResolvedMapping]:
     """Resolve AWQ mappings for the given model.
 
@@ -444,15 +470,54 @@ def resolve_mappings(
       2. ``AWQ_MAPPING_REGISTRY`` — model-class-name lookup
       3. ``default_mappings`` — Llama-like fallback.
 
+    Args:
+        model: The model to resolve mappings against.
+        user_mappings: Optional explicit mappings; when provided they are used
+            verbatim and ``skip_moe`` is ignored.
+        skip_moe: When True, drop routed MoE experts (module names matching
+            ``.experts.<N>.``) from the resolved mappings so AWQ leaves each
+            routed expert to the downstream block quantizer.
+
     Returns:
         List of ``ResolvedMapping`` objects ready for AWQ grid search.
     """
     if user_mappings is not None:
-        mapping_defs = [AWQMapping(m["smooth_layer"], m["balance_layers"]) for m in user_mappings]
+        mapping_defs = [
+            AWQMapping(m["smooth_layer"], m["balance_layers"], m.get("activation_hook_target")) for m in user_mappings
+        ]
     else:
         mapping_defs = _get_mappings_for_model(model)
 
-    return _resolve_mapping_defs(model, mapping_defs)
+    resolved = _resolve_mapping_defs(model, mapping_defs)
+
+    if skip_moe and user_mappings is None:
+        resolved = _drop_routed_experts(model, resolved)
+
+    return resolved
+
+
+def _drop_routed_experts(model: torch.nn.Module, resolved: list[ResolvedMapping]) -> list[ResolvedMapping]:
+    """Remove routed MoE experts from resolved mappings for ``skip_moe``."""
+    kept: list[ResolvedMapping] = []
+    dropped_experts = 0
+    for mapping in resolved:
+        if _ROUTED_EXPERT_RE.search(mapping.smooth_name):
+            dropped_experts += len(mapping.balance_names)
+            continue
+
+        keep_idx = [i for i, name in enumerate(mapping.balance_names) if not _ROUTED_EXPERT_RE.search(name)]
+        dropped_experts += len(mapping.balance_names) - len(keep_idx)
+        if not keep_idx:
+            continue
+        if len(keep_idx) < len(mapping.balance_names):
+            mapping.balance_names = [mapping.balance_names[i] for i in keep_idx]
+            mapping.balance_layers = [mapping.balance_layers[i] for i in keep_idx]
+            mapping.parent_name, mapping.parent = _find_parent(model, mapping.balance_names)
+        kept.append(mapping)
+
+    if dropped_experts:
+        logger.info(f"AWQ skip_moe: excluded {dropped_experts} routed-expert balance layer(s) from smoothing.")
+    return kept
 
 
 def _resolve_mapping_defs(
@@ -554,13 +619,7 @@ def _resolve_mapping_defs(
             "AWQConfig(mappings=[...])."
         )
     else:
-        first_prefix = next(iter(block_modules))
-        n_blocks = len(block_modules)
-        mappings_per_block = sum(1 for r in resolved if r.smooth_name.startswith(first_prefix))
-        logger.info(
-            f"AWQ resolved {matched_count} smooth-balance mappings "
-            f"({mappings_per_block} per block × {n_blocks} blocks)."
-        )
+        logger.info(f"AWQ resolved {matched_count} smooth-balance mappings.")
 
     return resolved
 
@@ -585,11 +644,11 @@ def check_model_compatibility(
     """
     warnings_list = []
     cls_name = _get_model_class_name(model)
-    in_registry = cls_name in AWQ_MAPPING_REGISTRY
+    in_registry = cls_name in AWQ_MAPPING_REGISTRY or cls_name in AWQ_DYNAMIC_MAPPING_REGISTRY
 
     if not in_registry and user_mappings is None:
         warnings_list.append(
-            f"Model class '{cls_name}' is not in AWQ_MAPPING_REGISTRY. "
+            f"Model class '{cls_name}' is not in any AWQ mapping registry. "
             f"Using default Llama-like mappings. If quantization quality is "
             f"poor, provide explicit mappings via AWQConfig(mappings=[...])."
         )

auto_round/autoround.py

@@ -379,6 +379,9 @@ def _select_rtn_compressor_base_cls(quant_config: "RTNConfig", scheme, format, b
     "clip_n_grid",
     "clip_max_shrink",
     "clip_n_sample_token",
+    "smooth_seqlen",
+    "smooth_batch_size",
+    "skip_moe",
     "mappings",
 }
 _ROTATION_FIELDS = {

auto_round/cli/algorithms.py

@@ -205,6 +205,20 @@ def register(self, group) -> None:
             type=int,
             help="Number of grid-search points for AWQ scaling ratio.",
         )
+        group.add_argument(
+            "--awq-smooth-seqlen",
+            dest="awq_smooth_seqlen",
+            default=None,
+            type=int,
+            help="Maximum sequence length used by AWQ scale-search parent replay.",
+        )
+        group.add_argument(
+            "--awq-smooth-batch-size",
+            dest="awq_smooth_batch_size",
+            default=None,
+            type=int,
+            help="Microbatch size for AWQ parent replay during scale search; <=0 disables microbatching.",
+        )
         group.add_argument(
             "--awq-apply-clip",
             dest="awq_apply_clip",
@@ -229,6 +243,8 @@ def build(self, args, common_kwargs: dict[str, Any]):
             n_grid=getattr(args, "n_grid", 20),
             apply_clip=getattr(args, "awq_apply_clip", False),
             clip_as_init=getattr(args, "awq_clip_as_init", False),
+            smooth_seqlen=getattr(args, "awq_smooth_seqlen", None) or 512,
+            smooth_batch_size=getattr(args, "awq_smooth_batch_size", None),
             **common_kwargs,
         )
 

docs/step_by_step.md

@@ -344,31 +344,53 @@ W2G64 Average Accuracy of 13 tasks and Time Cost Results(Testing was conducted o
 
 ### AWQ Algorithm
 
-**Experimental feature: our current implementation does not apply weight clipping yet, so accuracy may drop compared to the original AWQ algorithm.**
+**Experimental feature:** AWQ weight clipping is optional. Enable it with `--awq-apply-clip` when you want to match the original AWQ flow more closely.
 
 AWQ (Activation-Aware Weight Quantization) is available as an alternative quantization algorithm. AWQ protects salient weight channels by analyzing activation patterns and applying channel-wise scaling before standard RTN quantization.
 
-The canonical AWQ deployment path is **W4A16** served by vLLM's AWQ/Marlin CUDA kernels. **W8A8** with AWQ smoothing can also be served via vLLM's compressed_tensors backend (cutlass INT8 GEMM).
+The canonical AWQ deployment path is **W4A16** served by vLLM's AWQ/Marlin CUDA kernels. **INT8** is AutoRound's W8A8 scheme and can use AWQ smoothing before RTN quantization for vLLM's compressed_tensors backend (cutlass INT8 GEMM).
 
 #### CLI Usage
 
 ```bash
 auto-round --model Qwen/Qwen3-0.6B --scheme "W4A16" --algorithm awq --format "auto_round"
 ```
 
+INT8/W8A8 with AWQ smoothing and RTN:
+
+```bash
+auto-round \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --scheme INT8 \
+  --algorithm awq,rtn \
+  --nsamples 256 \
+  --seqlen 512 \
+  --awq-apply-clip \
+  --format auto_round:llm_compressor
+```
+
+For `INT8`, `disable_opt_rtn` defaults to `True`, so the command above uses plain RTN without requiring `--disable_opt_rtn`.
+
 AWQ-specific options:
-- `--duo_scaling`: Use both activations and weights for scaling. Options: `true`, `false`, or `both` (searches both modes and picks the best). (default: True).
-- `--n_grid`: Number of grid points for scaling ratio search (default: 20).
+- `--awq-duo-scaling`: Use both activations and weights for scaling. Options: `true`, `false`, or `both` (searches both modes and picks the best). (default: True).
+- `--awq-n-grid`: Number of grid points for scaling ratio search (default: 20).
+- `--awq-apply-clip`: Search and apply AWQ weight clipping after smoothing.
+
+API-only AWQ options:
+- `AWQConfig(smooth_seqlen=512)`: Caps the parent-forward replay length used by AWQ scale search. Set a value `<= 0` to use the full calibration sequence.
+- `AWQConfig(skip_moe=True)`: Skips routed MoE experts during AWQ smoothing while keeping attention and dense/shared paths. Explicit `mappings` are used as provided.
 
 #### API Usage
 
 ```python
-from auto_round import AWQConfig, AutoRound
+from auto_round import AWQConfig, AutoRound, RTNConfig
 
 ar = AutoRound(
-    "Qwen/Qwen3-0.6B",
+    "meta-llama/Llama-3.1-8B-Instruct",
     scheme="INT8",
-    alg_configs=AWQConfig(),
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig(disable_opt_rtn=True)],
+    nsamples=256,
+    seqlen=512,
 )
 
 output_dir = "./tmp_awq"
@@ -499,21 +521,43 @@ auto-round --model Qwen/Qwen3-0.6B --algorithm awq --scheme W4A16
 # AWQ + AutoRound optimization
 auto-round --model Qwen/Qwen3-0.6B --algorithm awq,auto_round --scheme W4A16
 
+# INT8/W8A8 + AWQ + RTN. disable_opt_rtn defaults to True for INT8.
+auto-round \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --scheme INT8 \
+  --algorithm awq,rtn \
+  --nsamples 256 \
+  --seqlen 512 \
+  --awq-apply-clip \
+  --format auto_round:llm_compressor
+
 # AWQ flags
---duo-scaling true|false|both  (default: true)
---n-grid 20                    (default: 20)
+--awq-duo-scaling true|false|both  (default: true)
+--awq-n-grid 20                    (default: 20)
+--awq-apply-clip
 ```
 
+`AWQConfig` also supports `smooth_seqlen=512` to cap AWQ scale-search replay length and `skip_moe=True` to leave routed MoE experts to the downstream block quantizer.
+
 #### API Usage
 ```python
-from auto_round import AWQConfig, AutoRound, SignRoundConfig
+from auto_round import AWQConfig, AutoRound, RTNConfig, SignRoundConfig
 
 # String alias (AWQ defaults, with RTN appended automatically)
 ar = AutoRound(model, tokenizer, alg_configs="awq", scheme="W4A16")
 
 # AWQ + default RTN (simplest)
 ar = AutoRound(model, tokenizer, alg_configs=AWQConfig(), scheme="W4A16")
 
+# INT8/W8A8 + AWQ + RTN
+ar = AutoRound(
+    "meta-llama/Llama-3.1-8B-Instruct",
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig(disable_opt_rtn=True)],
+    scheme="INT8",
+    nsamples=256,
+    seqlen=512,
+)
+
 # AWQ + AutoRound via alg_configs (explicit pipeline)
 ar = AutoRound(model, tokenizer, alg_configs=[AWQConfig(), SignRoundConfig(iters=200)], scheme="W4A16")
 ar.quantize_and_save(output_dir="./qmodel")