quantize_opera_nvfp4.py

#!/usr/bin/env python3
"""Lượng tử hoá beyoru/Opera sang NVFP4 theo ĐÚNG recipe của unsloth/Qwen3.6-27B-NVFP4.

Recipe đọc ngược từ `quantization_config` của unsloth (compressed-tensors,
format=mixed-precision, version 0.17.2.a20260707):

  group_0 — FP8 W8A8
    W: 8b float, symmetric, strategy=channel,      dynamic=False
    A: 8b float, symmetric, strategy=token,        dynamic=True
    targets: self_attn.(q|k|v|o)_proj
             linear_attn.(in_proj_qkv|in_proj_z|out_proj)
             lm_head
             layers.(56..63).mlp.(gate|up|down)_proj      <-- 8 LỚP CUỐI giữ FP8

  group_1 — NVFP4 W4A4
    W: 4b float, gs=16, symmetric, strategy=tensor_group, dynamic=False
    A: 4b float, gs=16, symmetric, strategy=tensor_group, dynamic=local
    targets: mlp.(gate|up|down)_proj

  ignore: toàn bộ model.visual.* (vision tower giữ bf16)

Ba mẹo giữ chất lượng, đều suy ra từ config chứ không phải đoán:
  1. 8 lớp cuối MLP giữ FP8 — lớp cuối nhạy nhất với lượng tử hoá.
  2. Attention + GDN projection giữ FP8; chỉ MLP thân giữa xuống W4A4.
  3. Vision tower nguyên bf16.

⚠️ QUAN TRỌNG — recipe này DATA-FREE: cả hai group đều có activation dynamic
(`token` / `local`), và scale của weight suy ra từ chính weight. Nên KHÔNG cần
dataset hiệu chuẩn. Đừng thêm calibration vào, sẽ lệch khỏi bản của unsloth.

Opera cùng kiến trúc với Qwen3.6-27B (64 lớp, hidden 5120, 24 head, interval 4)
nên các regex chuyển sang dùng được nguyên, KHÔNG cần sửa chỉ số lớp.

Chạy:
    pip install llmcompressor compressed-tensors
    python3 quantize_opera_nvfp4.py \
        --model beyoru/Opera \
        --output ./Opera-NVFP4
"""

import argparse
import re

# 8 lớp cuối giữ FP8. Opera có 64 lớp (0..63) => 56..63, y hệt unsloth.
LAST_LAYERS_FP8 = list(range(56, 64))

FP8_TARGETS = [
    r"re:.*self_attn\.(q|k|v|o)_proj$",
    r"re:.*linear_attn\.(in_proj_qkv|in_proj_z|out_proj)$",
    r"re:.*lm_head",
    r"re:.*layers\.(%s)\.mlp\.(gate|up|down)_proj$"
    % "|".join(str(i) for i in LAST_LAYERS_FP8),
]

NVFP4_TARGETS = [r"re:.*mlp\.(gate|up|down)_proj$"]


def build_recipe():
    """Recipe llmcompressor khớp từng trường với config của unsloth."""
    from llmcompressor.modifiers.quantization import QuantizationModifier

    return QuantizationModifier(
        config_groups={
            # ---- group_0: FP8 W8A8 -------------------------------------------
            "group_0": {
                "targets": FP8_TARGETS,
                "weights": {
                    "num_bits": 8,
                    "type": "float",
                    "symmetric": True,
                    "strategy": "channel",
                    "dynamic": False,
                },
                "input_activations": {
                    "num_bits": 8,
                    "type": "float",
                    "symmetric": True,
                    "strategy": "token",
                    "dynamic": True,
                },
            },
            # ---- group_1: NVFP4 W4A4 -----------------------------------------
            "group_1": {
                "targets": NVFP4_TARGETS,
                "weights": {
                    "num_bits": 4,
                    "type": "float",
                    "group_size": 16,
                    "symmetric": True,
                    "strategy": "tensor_group",
                    "dynamic": False,
                },
                "input_activations": {
                    "num_bits": 4,
                    "type": "float",
                    "group_size": 16,
                    "symmetric": True,
                    "strategy": "tensor_group",
                    "dynamic": "local",
                },
            },
        },
        # Vision tower giữ bf16. Dùng regex thay vì liệt kê 303 tên như unsloth --
        # tương đương về hiệu lực nhưng không phụ thuộc số block của vision.
        # Lưu ý: KHÔNG đưa lm_head vào đây -- nó thuộc group_0 (FP8) trong recipe gốc.
        ignore=["re:.*visual\\..*"],
    )


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="beyoru/Opera")
    ap.add_argument("--output", default="./Opera-NVFP4")
    ap.add_argument(
        "--dry-run",
        action="store_true",
        help="chi in recipe + doi chieu voi unsloth, khong nap model",
    )
    args = ap.parse_args()

    print("=== recipe se ap dung ===")
    print("  group_0 (FP8 W8A8) targets:")
    for t in FP8_TARGETS:
        print("   ", t)
    print("  group_1 (NVFP4 W4A4) targets:")
    for t in NVFP4_TARGETS:
        print("   ", t)
    print("  ignore: re:.*visual\\..*   (vision tower giu bf16)")
    print("  8 lop cuoi giu FP8:", LAST_LAYERS_FP8)
    print("  KHONG dung calibration dataset (recipe data-free)")

    if args.dry_run:
        print("\n[dry-run] dung o day.")
        return

    from llmcompressor import oneshot
    from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

    # Opera là `Qwen3_5ForConditionalGeneration` (có vision tower 333 tensor), KHÔNG phải
    # CausalLM thuần. Dùng AutoModelForCausalLM sẽ nạp thiếu/sai phần visual. Chọn class
    # theo `architectures` trong config thay vì đoán.
    cfg = AutoConfig.from_pretrained(args.model)
    arch = (getattr(cfg, "architectures", None) or [""])[0]
    print(f"\n=== nap {args.model} (bf16, ~55 GB) ===")
    print(f"  architectures: {arch}")
    if "ConditionalGeneration" in arch or "ImageTextToText" in arch:
        from transformers import AutoModelForImageTextToText

        loader = AutoModelForImageTextToText
        print("  -> dung AutoModelForImageTextToText (model da mode)")
    else:
        loader = AutoModelForCausalLM
        print("  -> dung AutoModelForCausalLM")

    model = loader.from_pretrained(args.model, dtype="auto", device_map="auto")
    tokenizer = AutoTokenizer.from_pretrained(args.model)

    # ⚠️ PHẢI truyền pipeline="datafree". Mặc định là "independent", nó dispatch sang
    # "sequential" — pipeline đó cần dataloader để trace model qua từng lớp, không có
    # dataset thì vỡ ngay: `TypeError: 'NoneType' object is not iterable` tại
    # pipelines/sequential/pipeline.py:103 (`next(iter(dataloader))`).
    # Recipe này data-free thật (activation dynamic, scale weight suy từ chính weight)
    # nhưng llmcompressor không tự suy ra được — phải nói thẳng.
    print("=== oneshot quantize (pipeline=datafree) ===")
    oneshot(model=model, recipe=build_recipe(), pipeline="datafree")

    print(f"=== luu {args.output} ===")
    model.save_pretrained(args.output, save_compressed=True)
    tokenizer.save_pretrained(args.output)
    print("XONG")


if __name__ == "__main__":
    main()