[Transform] Serialize transforms config

[kylesayrs]

Purpose

• Enable saving models with applied transforms
  • Transform config encodes both online and offline (fused) rotations

config.json

{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128009,
  "head_dim": 128,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "mlp_bias": false,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "quantization_config": {
    "config_groups": {
      "group_0": {
        "input_activations": null,
        "output_activations": null,
        "targets": [
          "Linear"
        ],
        "weights": {
          "actorder": null,
          "block_structure": null,
          "dynamic": false,
          "group_size": 128,
          "num_bits": 4,
          "observer": "minmax",
          "observer_kwargs": {},
          "strategy": "group",
          "symmetric": true,
          "type": "int"
        }
      }
    },
    "global_compression_ratio": null,
    "ignore": [
      "lm_head"
    ],
    "kv_cache_scheme": null,
    "quant_method": "compressed-tensors",
    "quantization_status": "compressed",
    "sparsity_config": {},
    "transform_config": {
      "config_groups": {
        "u": {
          "apply": [
            {
              "ignore": [
                "lm_head"
              ],
              "inverse": false,
              "location": "weight_output",
              "targets": [
                "Linear"
              ]
            },
            {
              "ignore": [
                "lm_head"
              ],
              "inverse": true,
              "location": "output",
              "targets": [
                "Linear"
              ]
            }
          ],
          "head_dim": null,
          "randomize": false,
          "requires_grad": false,
          "type": "random-hadamard"
        },
        "v": {
          "apply": [
            {
              "ignore": [
                "lm_head"
              ],
              "inverse": false,
              "location": "input",
              "targets": [
                "Linear"
              ]
            },
            {
              "ignore": [
                "lm_head"
              ],
              "inverse": true,
              "location": "weight_input",
              "targets": [
                "Linear"
              ]
            }
          ],
          "head_dim": null,
          "randomize": false,
          "requires_grad": false,
          "type": "random-hadamard"
        }
      }
    },
    "version": "0.10.3.dev146+ga3cd59d"
  },
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.55.0.dev0",
  "use_cache": true,
  "vocab_size": 128256
}

Prerequisites

• [Transform] Serialize with tied weights #370

Changes

• Implement transform_config similar to sparsity config, as a subconfig of the quantization config
  • This aligns with HF's pattern of treating a "quantization config" as a compression or optimization config
• Transform config is passed to serialization by attaching the transform config to the model when it is applied to the model
• Refactor ModelCompressor.update_config to support writing q/s/t configs

Follow ups

• Some work will need to be done if we want to support users passing with CompressedTensorsConfig
• Right now there are 3 ways we pass configs. Some work could be done to consolidate these methods [WIP] Refactor serialization of qconfig #410
  • qconfig is reconstructed from attached schemes
  • sconfig is inferred from the model in LC and passed as an argument
  • qconfig is attached to the model directly

Testing

• Sample model is located at https://huggingface.co/nm-testing/Meta-Llama-3-8B-Instruct-quip-w4a16

[brian-dellabetta]

changes to transform config look good to me, so approving, but definitely need to confirm with @dsikka and @rahul-tuli the changes to quantization life cycle

The base branch was changed.

[dsikka]

LGTM but needs rebase

[brian-dellabetta]

one nit question, otherwise LGTM