runner.py

# runner.py - PTQ Runner - Post-Training Quantization for KAN models

import os, pdb
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import argparse, yaml, os, json, torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision as tv
import torchvision.transforms as T
import munch  # Required for QAT-style config
import logging
from pathlib import Path

# Import QAT utilities for unified data loading and model creation
import util
from model import create_model
from util.ptq_helpers import set_seed, load_checkpoint, save_checkpoint, _clean_state_dict, _clean_state_dict, setup_logging, get_logger


# Quant APIs already in your repo
from ptq.gptq import apply_gptq as apply_gptq_block   # GPTQ (block/diag)
from ptq.gptq_strict import apply_gptq as apply_gptq_strict  # Strict GPTQ (rfactor/triangular)

from ptq.adaround import apply_adaround          # AdaRound
from ptq.awq import apply_awq                    # AWQ
from ptq.hawq_v2 import apply_hawq_v2            # HAWQ-V2
from ptq.brecq import apply_brecq, save_act_scales_json  # BRECQ
from ptq.uniform import (                        # Uniform
    quantize_weights_uniform, calibrate_activation_uniform, save_activation_json
)
from ptq.smoothquant import (                    # SmoothQuant
    compute_scales as sq_compute, apply_smoothquant as sq_apply, save_scales_json as sq_save,
    _is_smoothquant_wrapped_state_dict, _maybe_unwrap_state_dict_for_ptq, _maybe_prepare_model_for_state_dict,
)
from ptq.zeroq import generate_zeroq_data        # ZeroQ

# Activation PTQ helpers 
from ptq.actquant import collect_input_activation_scales, wrap_input_quant

# (optional): sidecar params writer for ACTQ
try:
    from ptq.actquant import save_act_params  # may not exist in older versions
except Exception:
    save_act_params = None







OWNER_CLASSES = {
    "GRAMLayer",
    "KAGNConvNDLayer", "KAGNConv1DLayer", "KAGNConv2DLayer", "KAGNConv3DLayer",
}


# =====================================================================================
# CONFIG ADAPTER: Convert QAT-style config to args object for unified data/model loading
# =====================================================================================

def _adapt_config_to_args(cfg: dict) -> munch.Munch:
    """
    Convert the QAT-style YAML config dict into a munch.Munch object
    that mimics the args structure expected by:
      - util.load_data(args.dataloader)
      - create_model(args)
    
    This bridges the gap between the new unified config format and the
    QAT utilities.
    
    The QAT's load_data(cfg) expects these fields DIRECTLY on cfg:
      - cfg.dataset
      - cfg.path
      - cfg.batch_size
      - cfg.val_split
      - cfg.workers
      - cfg.deterministic
      - cfg.num_classes (optional)
    
    In main.py, this is called as: util.load_data(args.dataloader)
    So args.dataloader must have all these fields.
    """
    # Start with munchifying the entire config
    args = munch.munchify(cfg)
    
    # Ensure dataloader section exists
    if not hasattr(args, 'dataloader') or args.dataloader is None:
        args.dataloader = munch.Munch()
    
    dl = args.dataloader
    
    # =========================================================================
    # Required fields for load_data(cfg) - must be DIRECTLY on dataloader
    # =========================================================================
    
    # dataset: required
    if not hasattr(dl, 'dataset'):
        dl.dataset = 'cifar10'  # fallback default
    
    # path: QAT uses 'path' for dataset root directory
    if not hasattr(dl, 'path'):
        # Check if old-style 'data.root' exists
        if 'data' in cfg and 'root' in cfg['data']:
            dl.path = cfg['data']['root']
        else:
            dl.path = './datasets'  # fallback default
    
    # batch_size: required
    if not hasattr(dl, 'batch_size'):
        # Check if in old-style 'data' section
        if 'data' in cfg and 'batch_size' in cfg['data']:
            dl.batch_size = cfg['data']['batch_size']
        else:
            dl.batch_size = 128
    
    # val_split: required (0.0 means use test set as validation)
    if not hasattr(dl, 'val_split'):
        dl.val_split = 0.0
    
    # workers: required for DataLoader num_workers
    # This is expected DIRECTLY on the dataloader config object
    if not hasattr(dl, 'workers'):
        # Check old-style locations
        if 'data' in cfg and 'num_workers' in cfg['data']:
            dl.workers = cfg['data']['num_workers']
        else:
            dl.workers = 4  # reasonable default
    
    # deterministic: required for worker_init_fn
    # This is expected DIRECTLY on the dataloader config object
    if not hasattr(dl, 'deterministic'):
        # Check root level (common in default config)
        dl.deterministic = cfg.get('deterministic', True)
    
    # num_classes: optional but useful
    if not hasattr(dl, 'num_classes'):
        # Infer from dataset
        dataset_classes = {
            'mnist': 10,
            'cifar10': 10,
            'cifar100': 100,
            'tinyimagenet': 200,
            'imagenet': 1000,
        }
        dl.num_classes = dataset_classes.get(str(dl.dataset).lower(), 10)
    
    # =========================================================================
    # Required fields for create_model(args)
    # =========================================================================
    
    # arch: required at top level
    if not hasattr(args, 'arch'):
        if 'model' in cfg and 'arch' in cfg['model']:
            args.arch = cfg['model']['arch']
        else:
            args.arch = 'kan_test_all_mnist'  # fallback
    
    # pre_trained: required (create_model logs this)
    if not hasattr(args, 'pre_trained'):
        # Check pretrained section
        if 'pretrained' in cfg:
            args.pre_trained = cfg['pretrained'].get('pretrain_from_fp32', False)
        else:
            args.pre_trained = False
    
    return args


def _build_data(cfg: dict):
    """
    Build data loaders using the QAT-style unified data loader.
    """
    args = _adapt_config_to_args(cfg)

    train_loader, val_loader, test_loader = util.load_data(args.dataloader)
    
    # For PTQ, we typically use train_loader for calibration and test_loader for eval
    return train_loader, test_loader


def _build_model(cfg: dict):
    """
    Build model using QAT's unified create_model function.
    """
    args = _adapt_config_to_args(cfg)
    
    return create_model(args)


def _layer_types(cfg):
    return tuple(cfg["ptq"]["layer_types"])


def _layer_types_for_act(cfg):
    """
    For activation PTQ, optionally skip owner layers.
    """
    base = list(_layer_types(cfg))
    if bool(cfg["ptq"].get("skip_owner_actq", False)):
        base = [t for t in base if t not in OWNER_CLASSES]
    return tuple(base)


# ----------------------------- subcommands -----------------------------
def enforce_torch_determinism(deterministic: bool = True):
    if not deterministic:
        return
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    try:
        torch.use_deterministic_algorithms(True)
    except Exception:
        pass
    try:
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
    except Exception:
        pass


def cmd_ptq(args):
    logger = get_logger()
    
    cfg = yaml.safe_load(open(args.config, "r"))
    logger.info(f"Loaded config from: {args.config}")

    # Handle seed - check multiple possible locations
    seed = cfg.get("seed", 42) # Config can specify seed as well.
    det = bool(cfg.get("deterministic", True))
    set_seed(int(seed), deterministic=det)
    enforce_torch_determinism(det)
    logger.info(f"Set seed: {seed}, deterministic: {det}")

    device = args.device or cfg.get("device", "cuda")
    device = device if torch.cuda.is_available() else "cpu"
    logger.info(f"Using device: {device}")

    # Use unified data loader
    logger.info("Building data loaders...")
    tr_loader, _ = _build_data(cfg)
    logger.info(f"Training loader: {len(tr_loader)} batches, batch_size={tr_loader.batch_size}")
    
    layer_types = _layer_types(cfg)
    logger.info(f"Target layer types: {layer_types}")

    method = args.method.lower()
    logger.info(f"PTQ method: {method}")

    # Use unified model builder
    logger.info("Building student and teacher models...")
    student = _build_model(cfg)
    teacher = _build_model(cfg)
    
    # Log model info
    total_params = sum(p.numel() for p in student.parameters())
    trainable_params = sum(p.numel() for p in student.parameters() if p.requires_grad)
    logger.info(f"Model parameters: {total_params:,} total, {trainable_params:,} trainable")

    logger.info(f"Loading checkpoint from: {args.ckpt}")
    ckpt = load_checkpoint(args.ckpt, map_location="cpu")
    sd_clean = _clean_state_dict(ckpt)
    logger.info(f"Checkpoint contains {len(sd_clean)} keys")

    sd_to_load, did_unwrap = _maybe_unwrap_state_dict_for_ptq(sd_clean, method=method)
    if did_unwrap:
        logger.info("SmoothQuant checkpoint detected and unwrapped for PTQ compatibility")

    try:
        teacher.load_state_dict(sd_to_load, strict=True)
        student.load_state_dict(sd_to_load, strict=True)
        logger.info("Loaded checkpoint into teacher and student models (strict=True)")
    except Exception as e:
        logger.warning(f"Strict load failed after SQ-unwrapping attempt: {e}")
        logger.warning("Falling back to wrapper-rebuild strict load.")
        _maybe_prepare_model_for_state_dict(teacher, sd_clean, context="PTQ teacher")
        _maybe_prepare_model_for_state_dict(student, sd_clean, context="PTQ student")
        teacher.load_state_dict(sd_clean, strict=True)
        student.load_state_dict(sd_clean, strict=True)
        did_unwrap = False

    # choose input (real vs ZeroQ)
    calib_loader = tr_loader
    if args.zeroq:
        logger.info("Using ZeroQ for data-free calibration...")
        # Determine image size from config
        img_size = 32  # default
        if 'dataloader' in cfg:
            dataset = cfg['dataloader'].get('dataset', 'cifar10').lower()
            img_size_map = {'mnist': 28, 'cifar10': 32, 'cifar100': 32, 'tinyimagenet': 64, 'imagenet': 224}
            img_size = img_size_map.get(dataset, 32)
        
        logger.info(f"ZeroQ params: num_images={args.z_num_images}, img_size={img_size}, iters={args.z_iters}")
        _, calib_loader = generate_zeroq_data(
            teacher, num_images=int(args.z_num_images), batch_size=int(args.z_batch_size),
            img_size=img_size, iters=int(args.z_iters), lr=float(args.z_lr),
            tv_weight=float(args.z_tv), l2_weight=float(args.z_l2), bn_weight=float(args.z_bnw),
            input_clamp=float(args.z_input_clamp), device=device
        )
        logger.info("ZeroQ synthetic data generated successfully")

    logger.info(f"Starting {method.upper()} quantization...")
    
    if method == "gptq":
        gptq_impl = str(getattr(args, "gptq_impl", "block")).lower()
        logger.info(f"GPTQ implementation: {gptq_impl}, w_bit={args.w_bit}, damping={args.damping}")

        if gptq_impl == "block":
            logger.info(f"GPTQ block mode: {args.gptq_mode}, block_size={args.block_size}")
            apply_gptq_block(
                model=student,
                calib_loader=calib_loader,
                bit=int(args.w_bit),
                device=device,
                mode=str(args.gptq_mode),
                nsamples=int(args.nsamples),
                damping=float(args.damping),
                block_size=int(args.block_size),
                layer_types=layer_types,
                max_batches=int(args.calib_batches),
            )

        elif gptq_impl in ("rfactor", "triangular"):
            apply_gptq_strict(
                model=student,
                calib_loader=calib_loader,
                bit=int(args.w_bit),
                device=device,
                nsamples=int(args.nsamples),
                damping=float(args.damping),
                impl=gptq_impl,
                layer_types=layer_types,
                max_batches=int(args.calib_batches),
            )
        else:
            raise ValueError(f"Unknown --gptq_impl: {gptq_impl}")
        logger.info("GPTQ quantization completed")

    elif method == "adaround":
        logger.info(f"AdaRound params: w_bit={args.w_bit}, iters={args.iters}, lr={args.lr}")
        apply_adaround(
            model=student, teacher=teacher, calib_loader=calib_loader, bit=int(args.w_bit),
            per_channel=bool(cfg["ptq"]["per_channel"]), iters=int(args.iters), lr=float(args.lr),
            reg_param=float(args.reg_param), start_beta=float(cfg["ptq"]["start_beta"]),
            end_beta=float(cfg["ptq"]["end_beta"]), gamma=float(cfg["ptq"]["gamma"]), zeta=float(cfg["ptq"]["zeta"]),
            calib_size=int(args.calib_size), layer_types=layer_types, device=device,
            seed=int(seed),
            ptq_batch_size=int(cfg["ptq"].get("adaround_batch_size", 32))
        )
        logger.info("AdaRound quantization completed")

    elif method == "awq":
        logger.info(f"AWQ params: w_bit={args.w_bit}, iters={args.iters}, lr={args.lr}, s_min={args.s_min}")
        apply_awq(
            model=student, teacher=teacher, calib_loader=calib_loader, bit=int(args.w_bit),
            iters=int(args.iters), lr=float(args.lr), s_min=float(args.s_min),
            reg_lambda=float(args.reg_lambda), calib_size=int(args.calib_size),
            layer_types=layer_types, device=device, batch_size=int(args.batch_size)
        )
        logger.info("AWQ quantization completed")

    elif method == "hawqv2":
        logger.info(f"HAWQ-V2 params: bit_candidates={args.bit_candidates}, target_avg_bit={args.target_avg_bit}")
        allocation, traces, qerr = apply_hawq_v2(
            model=student, calib_loader=calib_loader,
            bit_candidates=list(map(int, args.bit_candidates.split(","))),
            target_avg_bit=float(args.target_avg_bit) if args.target_avg_bit > 0 else None,
            keep_first_last=not args.no_keep_first_last, layer_types=layer_types,
            device=device, num_trace_samples=int(args.num_trace_samples),
            max_trace_batches=int(args.max_trace_batches)
        )
        logger.info(f"HAWQ-V2 completed. Allocation: {allocation}")

    elif method == "brecq":
        logger.info(f"BRECQ params: w_bit={args.w_bit}, a_bit={args.a_bit}, iters={args.iters}")
        act_scales = apply_brecq(
            model=student, teacher=teacher, calib_loader=calib_loader,
            w_bits=int(args.w_bit), a_bits=int(args.a_bit),
            iters_per_block=int(args.iters), lr=float(args.lr),
            block_types=tuple(args.block_types.split(",")),
            layer_types=layer_types, device=device,
            calib_size=int(args.calib_size), max_batches=int(args.calib_batches),
        )
        scales_json = os.path.splitext(args.output_ckpt)[0] + "_act_scales.json"
        save_act_scales_json(act_scales, scales_json)
        logger.info(f"BRECQ completed. Activation scales saved to: {scales_json}")

    elif method == "uniform":
        logger.info(f"Uniform params: w_bit={args.w_bit}, a_bit={args.a_bit}, w_mode={args.w_mode}")
        quantize_weights_uniform(
            model=student, w_bit=int(args.w_bit),
            per_channel=bool(args.w_per_channel == "pc"),
            mode=str(args.w_mode), layer_types=layer_types
        )
        act_params = calibrate_activation_uniform(
            model=student, calib_loader=calib_loader, a_bit=int(args.a_bit),
            per_channel=bool(args.a_per_channel == "pc"),
            mode=str(args.a_mode), method=str(args.a_calib), percentile=float(args.a_percentile),
            layer_types=layer_types, device=device, max_batches=int(args.calib_batches)
        )
        act_json = os.path.splitext(args.output_ckpt)[0] + "_act_scales.json"
        save_activation_json(act_params, act_json)
        logger.info(f"Uniform quantization completed. Activation params saved to: {act_json}")

    else:
        raise ValueError(f"Unknown PTQ method: {method}")

    os.makedirs(os.path.dirname(args.output_ckpt), exist_ok=True)
    save_checkpoint(
        {
            "state_dict": student.state_dict(),
            "cfg": cfg,
            "ptq": {
                "method": method,
                "source_ckpt": args.ckpt,
                "sq_unwrapped_for_ptq": bool(did_unwrap),
            },
        },
        args.output_ckpt,
        is_best=False
    )
    logger.info(f"[PTQ] Quantized model saved to: {args.output_ckpt}")


def cmd_act(args):
    logger = get_logger()
    
    cfg = yaml.safe_load(open(args.config, "r"))
    logger.info(f"Loaded config from: {args.config}")
    
    seed = cfg.get("seed", 42)
    set_seed(int(seed), deterministic=bool(cfg.get("deterministic", True)))
    logger.info(f"Set seed: {seed}")
    
    device = args.device or cfg.get("device", "cuda")
    device = device if torch.cuda.is_available() else "cpu"
    logger.info(f"Using device: {device}")

    logger.info("Building data loaders...")
    tr_loader, _ = _build_data(cfg)
    logger.info(f"Training loader: {len(tr_loader)} batches")
    
    logger.info("Building model...")
    model = _build_model(cfg)

    logger.info(f"Loading checkpoint from: {args.ckpt}")
    ckpt = load_checkpoint(args.ckpt, map_location="cpu")
    sd_clean = _clean_state_dict(ckpt)

    _maybe_prepare_model_for_state_dict(model, sd_clean, context="ACT model")
    model.load_state_dict(sd_clean, strict=True)
    logger.info("Checkpoint loaded successfully")

    calib_batches = int(args.calib_batches or cfg["ptq"].get("calib_batches", 32))
    logger.info(f"Calibration batches: {calib_batches}")

    act_layer_types = _layer_types_for_act(cfg)
    logger.info(f"Target layer types for activation: {act_layer_types}")

    logger.info(f"Collecting activation scales (percentile={args.percentile})...")
    scales = collect_input_activation_scales(
        model, tr_loader, percentile=float(args.percentile),
        layer_types=act_layer_types,
        max_batches=calib_batches, device=device
    )
    logger.info(f"Collected scales for {len(scales)} layers")
    
    logger.info(f"Wrapping input quantization (abit={args.abit})...")
    model = wrap_input_quant(model, scales, abit=int(args.abit), layer_types=act_layer_types)

    meta = ckpt.get("ptq", {})
    meta.update({"act_ptq": {"abit": int(args.abit), "percentile": float(args.percentile),
                             "skip_owner_actq": bool(cfg["ptq"].get("skip_owner_actq", False))}})

    os.makedirs(os.path.dirname(args.output_ckpt), exist_ok=True)
    save_checkpoint({"state_dict": model.state_dict(), "cfg": cfg, "ptq": meta}, args.output_ckpt, is_best=False)

    base = os.path.splitext(args.output_ckpt)[0]
    scales_json_path = base + "_act_scales.json"
    with open(scales_json_path, "w") as f:
        json.dump(scales, f, indent=2)
    
    logger.info(f"[ACT] Model saved to: {args.output_ckpt}")
    logger.info(f"[ACT] Scales saved to: {scales_json_path}")

    if save_act_params is not None:
        params_json_path = base + "_act_params.json"
        try:
            save_act_params(
                params_json_path,
                abit=int(args.abit),
                signed=True,
                percentile=float(args.percentile),
                estimator="batch_max_percentile",
                layer_types=_layer_types(cfg),
            )
            logger.info(f"[ACT] Params saved to: {params_json_path}")
        except Exception as e:
            logger.warning(f"[ACT] Could not write act params sidecar: {e}")


def cmd_smooth(args):
    logger = get_logger()
    
    cfg = yaml.safe_load(open(args.config, "r"))
    logger.info(f"Loaded config from: {args.config}")
    
    seed = cfg.get("seed", 42)
    set_seed(int(seed), deterministic=bool(cfg.get("deterministic", True)))
    logger.info(f"Set seed: {seed}")
    
    device = args.device or cfg.get("device", "cuda")
    device = device if torch.cuda.is_available() else "cpu"
    logger.info(f"Using device: {device}")

    logger.info("Building data loaders...")
    tr_loader, _ = _build_data(cfg)
    logger.info(f"Training loader: {len(tr_loader)} batches")
    
    logger.info("Building model...")
    model = _build_model(cfg)
    
    logger.info(f"Loading checkpoint from: {args.ckpt}")
    ckpt = load_checkpoint(args.ckpt, map_location="cpu")

    sd_clean = _clean_state_dict(ckpt)

    if _is_smoothquant_wrapped_state_dict(sd_clean):
        logger.warning("Input checkpoint already looks SmoothQuant-wrapped. Double-folding may occur.")
        _maybe_prepare_model_for_state_dict(model, sd_clean, context="SmoothQuant input model")

    model.load_state_dict(sd_clean, strict=True)
    logger.info("Checkpoint loaded successfully")

    sq_layer_types = _layer_types(cfg)
    logger.info(f"Target layer types: {sq_layer_types}")

    logger.info(f"Computing SmoothQuant scales (alpha={args.alpha}, percentile={args.percentile})...")
    scales = sq_compute(
        model, tr_loader, percentile=float(args.percentile), alpha=float(args.alpha),
        s_min=float(args.s_min), s_max=float(args.s_max),
        layer_types=sq_layer_types, device=device, max_batches=int(args.calib_batches)
    )
    logger.info(f"Computed scales for {len(scales)} layers")
    
    logger.info("Applying SmoothQuant transformations...")
    sq_apply(model, scales, layer_types=sq_layer_types)
    logger.info("SmoothQuant applied successfully")

    if args.apply_act:
        logger.info(f"Applying activation quantization (abit={args.abit}, percentile={args.act_percentile})...")
        act_layer_types = _layer_types_for_act(cfg)
        act_scales = collect_input_activation_scales(
            model, tr_loader, percentile=float(args.act_percentile),
            layer_types=act_layer_types, max_batches=int(args.calib_batches),
            device=device
        )
        model = wrap_input_quant(model, act_scales, abit=int(args.abit), layer_types=act_layer_types)
        act_json = os.path.splitext(args.output_ckpt)[0] + "_act_scales.json"
        with open(act_json, "w") as f:
            json.dump(act_scales, f, indent=2)
        logger.info(f"[SmoothQuant] Act scales saved to: {act_json}")

        if save_act_params is not None:
            try:
                params_json_path = os.path.splitext(args.output_ckpt)[0] + "_act_params.json"
                save_act_params(
                    params_json_path,
                    abit=int(args.abit),
                    signed=True,
                    percentile=float(args.act_percentile),
                    estimator="batch_max_percentile",
                    layer_types=_layer_types(cfg),
                )
                logger.info(f"[SmoothQuant] Act params saved to: {params_json_path}")
            except Exception as e:
                logger.warning(f"[SmoothQuant] Could not write act params sidecar: {e}")

    os.makedirs(os.path.dirname(args.output_ckpt), exist_ok=True)
    save_checkpoint({"state_dict": model.state_dict(), "cfg": cfg, "ptq": {"method": "SmoothQuant"}}, args.output_ckpt, is_best=False)
    scales_json = os.path.splitext(args.output_ckpt)[0] + "_smooth_scales.json"
    sq_save(scales, scales_json)
    
    logger.info(f"[SmoothQuant] Model saved to: {args.output_ckpt}")
    logger.info(f"[SmoothQuant] Scales saved to: {scales_json}")


def cmd_eval(args):
    logger = get_logger()
    logger.info(f"Starting evaluation with config: {args.config}")
    logger.info(f"Checkpoint: {args.ckpt}")
    
    from eval import evaluate_ckpt
    acc = evaluate_ckpt(args.config, args.ckpt, device=args.device)
    
    logger.info(f"Evaluation completed - Top-1 Accuracy: {acc:.2f}%")


def build_parser():
    p = argparse.ArgumentParser(description="PTQ Runner - Post-Training Quantization for KAN models")
    p.add_argument("--log_dir", type=str, default=None,
                   help="Directory to save log files (default: derived from output_ckpt path)")
    p.add_argument("--log_level", type=str, default="INFO",
                   choices=["DEBUG", "INFO", "WARNING", "ERROR"],
                   help="Logging level (default: INFO)")
    
    sub = p.add_subparsers(dest="cmd", required=True)

    # PTQ
    q = sub.add_parser("ptq")
    q.add_argument("--config", default="configs/ptq_common.yaml")
    q.add_argument("--ckpt", required=True)
    q.add_argument("--output_ckpt", required=True)
    q.add_argument("--device", type=str, default=None)
    q.add_argument("--method", required=True,
                   choices=["gptq", "adaround", "awq", "hawqv2", "brecq", "uniform"])
    q.add_argument("--w_bit", type=int, default=4)
    q.add_argument("--a_bit", type=int, default=8)
    q.add_argument("--calib_size", type=int, default=1024)
    q.add_argument("--calib_batches", type=int, default=32)

    # method-specific knobs
    q.add_argument("--gptq_impl", default="block", choices=["block", "rfactor", "triangular"])
    q.add_argument("--gptq_mode", default="block", choices=["block", "diag"])
    q.add_argument("--block_size", type=int, default=128)
    q.add_argument("--nsamples", type=int, default=2048)
    q.add_argument("--damping", type=float, default=1e-4)

    q.add_argument("--iters", type=int, default=500)
    q.add_argument("--lr", type=float, default=5e-3)
    q.add_argument("--reg_param", type=float, default=1e-4)
    q.add_argument("--s_min", type=float, default=0.2)
    q.add_argument("--reg_lambda", type=float, default=1e-4)
    q.add_argument("--batch_size", type=int, default=32)
    q.add_argument("--bit_candidates", type=str, default="2,3,4,8")
    q.add_argument("--target_avg_bit", type=float, default=4.0)
    q.add_argument("--no_keep_first_last", action="store_true")
    q.add_argument("--num_trace_samples", type=int, default=24)
    q.add_argument("--max_trace_batches", type=int, default=16)
    q.add_argument("--block_types", type=str, default="BasicBlock")
    q.add_argument("--w_per_channel", type=str, default="pc", choices=["pc", "pt"])
    q.add_argument("--w_mode", type=str, default="symmetric", choices=["symmetric", "asymmetric"])
    q.add_argument("--a_per_channel", type=str, default="pc", choices=["pc", "pt"])
    q.add_argument("--a_mode", type=str, default="symmetric", choices=["symmetric", "asymmetric"])
    q.add_argument("--a_calib", type=str, default="percentile", choices=["percentile", "minmax"])
    q.add_argument("--a_percentile", type=float, default=0.999)

    # ZeroQ data-free toggle
    q.add_argument("--zeroq", action="store_true")
    q.add_argument("--z_num_images", type=int, default=1024)
    q.add_argument("--z_iters", type=int, default=400)
    q.add_argument("--z_batch_size", type=int, default=128)
    q.add_argument("--z_lr", type=float, default=0.05)
    q.add_argument("--z_tv", type=float, default=1e-5)
    q.add_argument("--z_l2", type=float, default=1e-6)
    q.add_argument("--z_bnw", type=float, default=1.0)
    q.add_argument("--z_input_clamp", type=float, default=3.0)

    # ACT
    a = sub.add_parser("act")
    a.add_argument("--config", default="configs/ptq_common.yaml")
    a.add_argument("--ckpt", required=True)
    a.add_argument("--output_ckpt", required=True)
    a.add_argument("--abit", type=int, default=8)
    a.add_argument("--percentile", type=float, default=0.999)
    a.add_argument("--calib_batches", type=int, default=None)
    a.add_argument("--device", type=str, default=None)

    # SmoothQuant
    s = sub.add_parser("smooth")
    s.add_argument("--config", default="configs/ptq_common.yaml")
    s.add_argument("--ckpt", required=True)
    s.add_argument("--output_ckpt", required=True)
    s.add_argument("--alpha", type=float, default=0.5)
    s.add_argument("--percentile", type=float, default=0.999)
    s.add_argument("--s_min", type=float, default=0.01)
    s.add_argument("--s_max", type=float, default=100.0)
    s.add_argument("--calib_batches", type=int, default=32)
    s.add_argument("--apply_act", action="store_true")
    s.add_argument("--abit", type=int, default=8)
    s.add_argument("--act_percentile", type=float, default=0.999)
    s.add_argument("--device", type=str, default=None)

    # Eval
    e = sub.add_parser("eval")
    e.add_argument("--config", default="configs/ptq_common.yaml")
    e.add_argument("--ckpt", required=True)
    e.add_argument("--device", type=str, default=None)

    return p




# =====================================================================================
# MAIN ENTRY POINT
# =====================================================================================

def main():
    p = build_parser()
    args = p.parse_args()
    
    # Determine log level
    log_level = getattr(logging, args.log_level.upper(), logging.INFO)
    
    # Determine log directory
    if args.log_dir:
        log_dir = Path(args.log_dir)
    elif hasattr(args, 'output_ckpt') and args.output_ckpt:
        log_dir = Path(args.output_ckpt).parent / "logs"
    else:
        log_dir = Path("./logs")
    
    # Determine experiment name from config or command
    experiment_name = f"PTQ_{args.cmd}"
    if hasattr(args, 'config') and args.config:
        cfg = yaml.safe_load(open(args.config, "r"))
        # experiment_name = cfg.get("name", f"PTQ_{args.cmd}")
        experiment_name = f"{cfg.get('name', f'PTQ_{args.cmd}')}__{Path(args.output_ckpt).stem if hasattr(args, 'output_ckpt') and args.output_ckpt else args.cmd}"

    
    # Setup logging
    logger, log_dir = setup_logging(
        log_dir=log_dir,
        experiment_name=experiment_name,
        level=log_level,
    )
    
    logger.info("=" * 70)
    logger.info(f"Starting PTQ Runner - Command: {args.cmd}")
    logger.info("=" * 70)
    
    # Log all arguments
    logger.debug("Command line arguments:")
    for arg_name, arg_value in vars(args).items():
        logger.debug(f"  {arg_name}: {arg_value}")
    
    try:
        if args.cmd == "ptq":
            cmd_ptq(args)
        elif args.cmd == "act":
            cmd_act(args)
        elif args.cmd == "smooth":
            cmd_smooth(args)
        elif args.cmd == "eval":
            cmd_eval(args)
        else:
            raise RuntimeError("unknown subcommand")
        
        logger.info("=" * 70)
        logger.info("PTQ Runner completed successfully")
        logger.info("=" * 70)
        
    except Exception as e:
        logger.error(f"PTQ Runner failed with error: {e}", exc_info=True)
        raise


if __name__ == "__main__":
    main()