Add Export Format

auto_round/export/export_to_yourformat/
├── __init__.py
└── export.py

def pack_layer(layer_name, model, backend, output_dtype=torch.float16):
    """Pack a quantized layer for serialization.

    Args:
        layer_name: Full module path (e.g., "model.layers.0.self_attn.q_proj")
        model: The quantized model
        backend: Backend configuration string
        output_dtype: Output tensor dtype

    Returns:
        dict: Packed tensors ready for serialization
    """
    layer = get_module(model, layer_name)
    device = layer.weight.device

    # Get quantization parameters from layer
    bits = layer.bits
    group_size = layer.group_size
    scale = layer.scale
    zp = layer.zp
    weight = layer.weight

    # Pack weights according to your format
    packed_weight = _pack_weights(weight, bits, group_size)

    return {
        f"{layer_name}.qweight": packed_weight,
        f"{layer_name}.scales": scale,
        f"{layer_name}.qzeros": zp,
    }

def save_quantized_as_yourformat(output_dir, model, tokenizer, layer_config, serialization_dict=None, **kwargs):
    """Save quantized model in your format.

    Args:
        output_dir: Directory to save to
        model: The quantized model
        tokenizer: Model tokenizer
        layer_config: Per-layer quantization configuration
        serialization_dict: Pre-packed layer tensors (optional)
        **kwargs: Additional format-specific arguments
    """
    import os
    from safetensors.torch import save_file

    os.makedirs(output_dir, exist_ok=True)

    # 1. Pack all quantized layers (if not pre-packed)
    if serialization_dict is None:
        serialization_dict = {}
        for layer_name, config in layer_config.items():
            serialization_dict.update(pack_layer(layer_name, model, ...))

    # 2. Save weights
    save_file(serialization_dict, os.path.join(output_dir, "model.safetensors"))

    # 3. Save quantization config
    quant_config = {
        "quant_method": "yourformat",
        "bits": ...,
        "group_size": ...,
        # format-specific metadata
    }
    # Write config to output_dir

    # 4. Save tokenizer
    tokenizer.save_pretrained(output_dir)

@OutputFormat.register("yourformat")
class YourFormat(OutputFormat):
    format_name = "yourformat"
    support_schemes = ["W4A16", "W8A16"]  # List supported scheme names

    def __init__(self, format: str, ar):
        super().__init__(format, ar)

    @classmethod
    def check_scheme_args(cls, scheme: QuantizationScheme) -> bool:
        """Check if a QuantizationScheme is compatible with this format."""
        return scheme.bits in [4, 8] and scheme.data_type == "int" and scheme.act_bits >= 16

    def pack_layer(self, layer_name, model, output_dtype=torch.float16):
        from auto_round.export.export_to_yourformat.export import pack_layer

        return pack_layer(layer_name, model, self.get_backend_name(), output_dtype)

    def save_quantized(self, output_dir, model, tokenizer, layer_config, serialization_dict=None, **kwargs):
        from auto_round.export.export_to_yourformat.export import save_quantized_as_yourformat

        return save_quantized_as_yourformat(
            output_dir, model, tokenizer, layer_config, serialization_dict=serialization_dict, **kwargs
        )

class SupportedFormats:
    def __init__(self):
        self._support_format = (
            "auto_round",
            "auto_gptq",
            # ...
            "yourformat",  # Add your format here
        )

BackendInfos["auto_round:yourformat"] = BackendInfo(
    device=["cuda"],
    sym=[True, False],
    packing_format=["yourformat"],
    bits=[4, 8],
    group_size=[32, 64, 128],
    priority=2,
)

def test_yourformat_export(tiny_opt_model_path, dataloader):
    ar = AutoRound(
        tiny_opt_model_path,
        bits=4,
        group_size=128,
        dataset=dataloader,
        iters=2,
        nsamples=2,
    )
    compressed_model, _ = ar.quantize()
    ar.save_quantized(output_dir="./tmp_yourformat", format="yourformat")

    # Verify saved files exist
    assert os.path.exists("./tmp_yourformat/model.safetensors")

    # Verify model can be loaded back
    from transformers import AutoModelForCausalLM

    loaded = AutoModelForCausalLM.from_pretrained("./tmp_yourformat")