compressed-tensors

compressed-tensors étend les fichiers safetensors avec des types de tenseurs compressés afin de fournir un format de checkpoint unifié pour le stockage et le chargement de divers formats de quantification, tels que dense, quantifié en entier (int8), quantifié en virgule flottante (fp8) et quantifié par paquets (int4 ou int8, poids quantifiés regroupés en int32).

compressed-tensors prend en charge le fine-tuning avec PEFT et inclut également les fonctionnalités suivantes.

• Précisions de poids et d'activation fp8, int4, int8.
• Stratégies d'échelles de quantification et de points zéro pour tensor, channel, group, block, token.
• Quantification dynamique d'activation par jeton (ou toute stratégie statique).
• Quantification de modules arbitraires, et non uniquement des modules nn.Linear.
• Prise en charge ciblée de modules spécifiques par nom ou par classe.

Installez compressed-tensors depuis PyPI pour obtenir la dernière version stable (recommandé) ou installez-le depuis le code source pour bénéficier des dernières fonctionnalités.

pip install compressed-tensors

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

Recherchez à l'aide de la balise tag compressed-tensors pour trouver un modèle compatible sur le Hugging Face Hub.

Les modèles pré-quantifiés peuvent être chargés directement. Pour quantifier un modèle au format compressed-tensors, consultez llm-compressor. Alternativement, les modèles peuvent être créés de manière indépendante et sérialisés avec une configuration 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

Modes de chargement

Un point de contrôle compressed-tensors stocke ses poids compressés (fp8, ou int4/int8 empilés). Leur exécution est déterminée par deux arguments de CompressedTensorsConfig.

Par défaut, save_pretrained() réécrit les tenseurs compressés sans les modifier, de sorte qu'un modèle chargé est enregistré dans le même format compressed-tensors.

Accélération des noyaux FP8

Passez use_optimized_inference=True pour conserver un modèle compressed-tensors en FP8 et exécuter ses produits matriciels via des noyaux FP8 accélérés par le matériel (torch.nn.functional.scaled_mm, qui délègue à torch._scaled_mm_v2 ; les versions antérieures de torch basculent sur torch._scaled_mm), au lieu de déquantiser les poids en BF16. Conserver les poids en FP8 tout au long de l'inférence réduit l'utilisation de la mémoire et accélère les calculs. Cette option est réservée à l'inférence ; désactivez-la pour l'ajustement fin.

Le chemin du noyau FP8 prend en charge ces dispositions de quantification.

Chargement d'un modèle FP8 pré-quantifié

Les noyaux FP8 sont activés sur demande : demandez-les avec use_optimized_inference=True, et ils sont utilisés lorsque la configuration du modèle spécifie une quantification FP8 et qu’un GPU pris en charge est 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))

Désquantification au moment du chargement

Sans use_optimized_inference=True, le modèle suit la voie standard de compressed-tensors : les poids restent compressés et compressed-tensors les décompresse lors du premier passage avant. Passez dequantize=True pour les décompresser au moment du chargement, ce qui est souhaitable si vous souhaitez affiner le modèle ou le sauvegarder dans sa précision d’origine (par exemple BF16). Le résultat est un modèle dense classique : aucun schéma compressed-tensors n’est associé, de sorte que la quantification des activations (par exemple w8a8) est également abandonnée.

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

Point de contrôle du modèle

Les modèles compressed-tensor sont définis par leur entrée de configuration. L’exemple suivant est extrait du fichier config.json de nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf.

Il existe de nombreuses entrées pour permettre une expression flexible, tant pendant qu’après la compression, mais les entrées relatives au chargement et à l’inférence peuvent être simplifiées afin de se concentrer sur quelques entrées clés.

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

Le fichier de configuration spécifie la quantification d’un groupe de configuration (group_0), qui comprend la quantification des poids et des activations en fp8 avec une stratégie statique par tenseur. Le module lm_head n’est pas quantifié, comme indiqué par la clé ignore.

Pour un examen plus détaillé des poids du modèle, utilisez le safetensors viewer sur la carte du modèle pour visualiser les poids quantifiés, l'échelle d'entrée et l'échelle des poids de tous les modules nn.Linear.

Lors du chargement d'un modèle compressed-tensors avec l'intégration ~quantizers.HFQuantizer, les modules cibles sont transmis à compressed-tensors : celui-ci attache le quantization_scheme résolu, définit quantization_status, enregistre les paramètres stockés dans le checkpoint (weight en fp8, ainsi que weight_scale et, pour une stratégie statique, input_scale) et installe son propre passage avant sur ces derniers. Ils restent des instances de nn.Linear, c'est donc ce que montre print — les versions récentes de compressed-tensors ne les encapsulent plus dans une sous-classe CompressedLinear. Les modules répertoriés sous ignore, tels que lm_head, sont laissés inchangés.

Avec dequantize=False (par défaut), les poids restent compressés une fois le chargement terminé, et compressed-tensors décompresse l’ensemble du modèle lors du premier passage avant. dequantize=True décompresse chaque poids au moment de son chargement, de sorte qu’aucun passage avant n’est nécessaire pour obtenir des poids denses, et les modules sont de simples couches denses sans 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

Avec use_optimized_inference=True, les couches couvertes par un groupe de configuration fp8 sont remplacées par CompressedTensorsFP8Linear, qui conserve le poids fp8 et son échelle dans la disposition attendue par son noyau de multiplication matricielle par ligne. Ces poids restent en fp8, y compris lors des passes avant.

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