compressed-tensors

compressed-tensors extiende los archivos safetensors a tipos de datos de tensores comprimidos para ofrecer un formato de checkpoint unificado que permite guardar y cargar diversos formatos de cuantización, como denso, cuantizado en enteros (int8), cuantizado en coma flotante (fp8) y cuantizado empaquetado (int4 o int8, con pesos cuantizados empaquetados en int32).

compressed-tensors permite el ajuste fino con PEFT e incluye las siguientes características.

• precisones de pesos y activaciones fp8, int4 e int8.
• Estrategias de escalas de cuantización y puntos cero para tensor, canal, grupo, bloque, token.
• Cuantización dinámica de activaciones por token (o cualquier estrategia estática).
• Cuantización de módulos arbitrarios, no solo de módulos nn.Linear.
• Soporte específico para módulos concretos por nombre o clase.

Instale compressed-tensors desde PyPI para obtener la última versión estable (recomendado) o instálelo desde el código fuente para acceder a las funciones más recientes.

pip install compressed-tensors

git clone https://github.com/neuralmagic/compressed-tensors
cd compressed-tensors
pip install -e .

Busca usando la etiqueta compressed-tensors para encontrar un modelo compatible en el Hugging Face Hub.

Los modelos pre-cuantizados se pueden cargar directamente. Para cuantizar un modelo al formato compressed-tensors, consulte llm-compressor. Alternativamente, los modelos se pueden crear de forma independiente y serializar con una configuración de compressed-tensors.

from transformers import AutoModelForCausalLM

ct_model = AutoModelForCausalLM.from_pretrained("nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf", device_map="auto")

# measure memory usage
mem_params = sum([param.nelement()*param.element_size() for param in ct_model.parameters()])
print(f"{mem_params/2**30:.4f} GB")
# 8.4575 GB

Modos de carga

Un checkpoint de compressed-tensors almacena sus pesos comprimidos (fp8, o int4/int8 empaquetados). La forma en que se ejecutan depende de dos argumentos de CompressedTensorsConfig.

Con el valor predeterminado, save_pretrained() escribe los tensores comprimidos sin cambios, por lo que un modelo cargado se guarda en el mismo formato de tensores comprimidos.

Aceleración de kernels FP8

Pase use_optimized_inference=True para mantener un modelo de compressed-tensors en FP8 y ejecutar sus matmuls a través de kernels de FP8 acelerados por hardware (torch.nn.functional.scaled_mm, que despacha a torch._scaled_mm_v2; las versiones anteriores de torch recurren a torch._scaled_mm), en lugar de descuantizar los pesos de vuelta a BF16. Mantener los pesos en FP8 durante toda la inferencia reduce el uso de memoria y acelera el cálculo. Esto es solo para inferencia, así que desactívelo para hacer fine-tuning.

La ruta del kernel FP8 admite estos diseños de cuantización.

Carga de un modelo FP8 pre-cuantizado

Los núcleos FP8 son opcionales: solicítalos con use_optimized_inference=True, y se utilizan cuando la configuración del modelo especifica cuantización FP8 y hay una GPU compatible disponible.

from transformers import AutoModelForCausalLM, AutoTokenizer, CompressedTensorsConfig

model = AutoModelForCausalLM.from_pretrained(
    "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic",
    quantization_config=CompressedTensorsConfig(use_optimized_inference=True),
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic")
inputs = tokenizer("Hello, how are you?", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Descomprimir en el momento de la carga

Sin use_optimized_inference=True, el modelo sigue la ruta regular de compressed-tensors: los pesos se dejan comprimidos y compressed-tensors los descomprime en la primera pasada hacia adelante. Pasa dequantize=True para descomprimirlos durante la carga, que es lo que necesitas para afinar el modelo o guardarlo en su precisión original (p. ej. BF16). El resultado es un modelo denso simple: no se adjunta ningún esquema de compressed-tensors, por lo que también se descarta la cuantización de activaciones (p. ej. w8a8).

from transformers import AutoModelForCausalLM, CompressedTensorsConfig

model = AutoModelForCausalLM.from_pretrained(
    "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic",
    quantization_config=CompressedTensorsConfig(dequantize=True),
    device_map="auto",
)

Punto de control del modelo

Los modelos compressed-tensor se definen a través de su entrada de configuración. El siguiente ejemplo se toma del archivo config.json de nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf.

Hay muchas entradas para permitir una expresión flexible tanto durante como después de la compresión, pero las entradas para la carga y la inferencia se pueden simplificar para centrarse solo en algunas entradas clave.

"quantization_config": {
  "config_groups": {
    "group_0": {
      "input_activations": {
        "num_bits": 8,
        "strategy": "tensor",
        "type": "float"
      },
      "targets": ["Linear"],
      "weights": {
        "num_bits": 8,
        "strategy": "tensor",
        "type": "float"
      }
    }
  },
  "format": "naive-quantized",
  "ignore": ["lm_head"],
  "quant_method": "compressed-tensors",
  "quantization_status": "frozen"
},

El archivo de configuración especifica la cuantización de un grupo de configuración (group_0), que incluye la cuantización de pesos y activaciones a fp8 con una estrategia estática por tensor. El módulo lm_head no está cuantizado, como se muestra en la clave ignore.

Para una vista más detallada de los pesos del modelo, use el visor de safetensors en la tarjeta del modelo para ver los pesos cuantizados, la escala de entrada y la escala de pesos de todos los módulos nn.Linear.

Al cargar un modelo de compressed-tensors con la integración ~quantizers.HFQuantizer, los módulos objetivo se transfieren a compressed-tensors: este adjunta el quantization_scheme resuelto, establece quantization_status, registra los parámetros que almacena el checkpoint (weight en fp8, además de weight_scale y, para una estrategia estática, input_scale) e instala su propio paso de propagación hacia adelante sobre ellos. Siguen siendo instancias de nn.Linear, por lo que es lo que muestra print: las versiones recientes de compressed-tensors ya no los envuelven en una subclase CompressedLinear. Los módulos enumerados bajo ignore, como lm_head, se dejan sin modificar.

Con dequantize=False (el valor predeterminado), los pesos siguen comprimidos una vez finalizada la carga, y compressed-tensors descomprime todo el modelo en la primera pasada hacia adelante. dequantize=True descomprime cada peso mientras se carga, por lo que no se necesita una pasada hacia adelante para obtener pesos densos, y los módulos son capas densas simples sin quantization_scheme.

import torch
from transformers import AutoModelForCausalLM, CompressedTensorsConfig

model_id = "nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf"

ct_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=CompressedTensorsConfig(dequantize=False),
    device_map="auto",
)
q_proj = ct_model.model.layers[0].self_attn.q_proj
print(q_proj, q_proj.quantization_status)
# Linear(in_features=4096, out_features=4096, bias=False) QuantizationStatus.COMPRESSED
# ^ compressed-tensors module: fp8 weight, weight_scale, and its own forward

ct_model(input_ids=torch.tensor([[0, 1, 2]], device=ct_model.device))
print(q_proj, q_proj.quantization_status)
# Linear(in_features=4096, out_features=4096, bias=False) QuantizationStatus.DECOMPRESSED
# ^ weight is BF16 now, decompressed by that forward pass

ct_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=CompressedTensorsConfig(dequantize=True),
    device_map="auto",
)
print(ct_model.model.layers[0].self_attn.q_proj)
# Linear(in_features=4096, out_features=4096, bias=False)      weight: BF16

Con use_optimized_inference=True, las capas cubiertas por un grupo de configuración fp8 se reemplazan por CompressedTensorsFP8Linear, que almacena el peso fp8 y su escala en el formato que espera su kernel de matmul por filas. Esos pesos permanecen en fp8, incluidas las pasadas hacia adelante.

from transformers import AutoModelForCausalLM, CompressedTensorsConfig

ct_model = AutoModelForCausalLM.from_pretrained(
    "nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf",
    quantization_config=CompressedTensorsConfig(use_optimized_inference=True),
    device_map="auto",
)
print(ct_model.model.layers[0].self_attn.q_proj)
# CompressedTensorsFP8Linear(in_features=4096, out_features=4096, bias=False)      weight: F8_E4M3