compressed-tensors

compressed-tensors estende os arquivos safetensors para tipos de dados de tensores comprimidos, fornecendo um formato de checkpoint unificado para armazenar e carregar diversos formatos de quantização, como denso, quantizado em int (int8), quantizado em float (fp8) e quantizado em pacote (int4 ou int8 com pesos quantizados empacotados em int32).

O compressed-tensors suporta fine-tuning com PEFT e inclui os seguintes recursos.

• precisões de pesos e ativações fp8, int4 e int8.
• Estratégias de escalas de quantização e pontos zero para tensor, canal, grupo, bloco, token.
• Quantização dinâmica de ativação por token (ou qualquer estratégia estática).
• Quantização de módulos arbitrários, não apenas de módulos nn.Linear.
• Suporte direcionado a módulos específicos por nome ou classe.

Instale o compressed-tensors a partir do PyPI para obter a versão estável mais recente (recomendado) ou instale-o a partir do código-fonte para obter os recursos mais recentes.

pip install compressed-tensors

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

Pesquise usando a tag compressed-tensors para encontrar um modelo compatível no Hugging Face Hub.

Modelos pré-quantizados podem ser carregados diretamente. Para quantizar um modelo para o formato compressed-tensors, consulte llm-compressor. Alternativamente, os modelos podem ser criados de forma independente e serializados com uma configuração 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 carregamento

Um checkpoint do compressed-tensors armazena seus pesos em formato comprimido (fp8, ou int4/int8 empacotados). A forma como eles são executados depende de dois argumentos de CompressedTensorsConfig.

Com o padrão, save_pretrained() grava os tensores comprimidos de volta sem alterações, de modo que um modelo carregado é salvo no mesmo formato de tensores comprimidos.

Aceleração de kernels FP8

Passe use_optimized_inference=True para manter um modelo de compressed-tensors em FP8 em FP8 e executar seus matmuls por meio de kernels de FP8 acelerados por hardware (torch.nn.functional.scaled_mm, que despacha para torch._scaled_mm_v2; versões mais antigas do torch fazem fallback para torch._scaled_mm), em vez de desquantizar os pesos de volta para BF16. Manter os pesos em FP8 durante todo o inferência reduz o uso de memória e acelera o cálculo. Isso é apenas para inferência, então deixe desativado para fazer fine-tuning.

O caminho do kernel FP8 suporta esses layouts de quantização.

Carregando um modelo FP8 pré-quantizado

Os kernels FP8 são opcionais: solicite-os com use_optimized_inference=True, e eles serão utilizados quando a config do modelo especificar quantização FP8 e uma GPU compatível estiver disponível.

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))

Desquantização no momento do carregamento

Sem use_optimized_inference=True, o modelo segue a rota regular do compressed-tensors: os pesos permanecem comprimidos e o compressed-tensors os descomprime na primeira passada de forward. Passe dequantize=True para descomprimi-los durante o carregamento, o que é o que você deseja ao fazer fine-tuning do modelo ou salvá-lo em sua precisão original (por exemplo, BF16). O resultado é um modelo denso simples: nenhum esquema do compressed-tensors é anexado, portanto, a quantização de ativação (por exemplo, w8a8) também é descartada.

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",
)

Checkpoint do modelo

Modelos compressed-tensor são definidos por meio de sua entrada de configuração. O exemplo a seguir foi extraído do arquivo config.json do nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf.

Há muitas entradas para permitir uma expressão flexível tanto durante quanto após a compressão, mas as entradas para carregamento e inferência podem ser simplificadas para focar apenas em algumas entradas-chave.

"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"
},

O arquivo de configuração especifica a quantização de um grupo de configuração (group_0), que inclui a quantização de pesos e ativações para fp8 com uma estratégia estática por tensor. O módulo lm_head não é quantizado, conforme indicado na chave ignore.

Para uma análise mais detalhada dos pesos do modelo, use o safetensors viewer no cartão do modelo para ver os pesos quantizados, a escala de entrada e a escala de peso de todos os módulos nn.Linear.

Ao carregar um modelo do compressed-tensors com a integração ~quantizers.HFQuantizer, os módulos-alvo são transferidos para o compressed-tensors: ele anexa o quantization_scheme resolvido, define quantization_status, registra os parâmetros armazenados no checkpoint (weight em fp8, além de weight_scale e, para uma estratégia estática, input_scale) e instala seu próprio forward pass sobre eles. Eles permanecem como instâncias de nn.Linear, e é isso que o print exibe — as versões recentes do compressed-tensors não os envolvem mais em uma subclasse CompressedLinear. Os módulos listados em ignore, como lm_head, são deixados intactos.

Com dequantize=False (o padrão), os pesos ainda ficam comprimidos após o carregamento, e o compressed-tensors descomprime o modelo inteiro na primeira passada de forward. dequantize=True descomprime cada peso durante o carregamento, de modo que nenhuma passada de forward é necessária para obter pesos densos, e os módulos são camadas densas simples, sem 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

Com use_optimized_inference=True, as camadas cobertas por um grupo de configuração fp8 são substituídas por CompressedTensorsFP8Linear, que armazena o peso fp8 e sua escala no layout esperado pelo kernel de matmul por linha. Esses pesos permanecem em fp8, inclusive nos passes de forward.

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