compressed-tensors

compressed-tensors расширяет файлы safetensors сжатыми типами данных тензоров, обеспечивая единый формат контрольных точек для сохранения и загрузки различных форматов квантования, таких как dense, int-квантование (int8), float-квантование (fp8) и pack-квантование (int4 или int8, весовое квантование, упакованное в int32).

compressed-tensors поддерживает дообучение с PEFT и включает в себя следующие возможности.

• Точности весов и активаций fp8, int4, int8.
• Стратегии масштабов квантования и нулевых точек для tensor, channel, group, block, token.
• Динамическое квантование активаций для каждого токена (или любая статическая стратегия).
• Квантование произвольных модулей, а не только модулей nn.Linear.
• Целевая поддержка конкретных модулей по имени или классу.

Установите compressed-tensors из PyPI, чтобы получить последнюю стабильную версию (рекомендуется), или установите пакет из исходного кода, чтобы получить доступ к новейшим функциям.

pip install compressed-tensors

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

Используйте тег compressed-tensors для поиска совместимой модели на Hugging Face Hub.

Предварительно квантованные модели можно загружать напрямую. Чтобы квантовать модель в формат compressed-tensors, см. llm-compressor. Альтернативно, модели можно создавать независимо и сериализовать с конфигурацией 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

Режимы загрузки

Чекпоинт compressed-tensors хранит веса в сжатом виде (fp8 или упакованные int4/int8). Способ их выполнения определяется двумя аргументами CompressedTensorsConfig.

По умолчанию save_pretrained() записывает сжатые тензоры без изменений, поэтому загруженная модель сохраняется в том же формате compressed-tensors.

Ускорение FP8-ядер

Передайте use_optimized_inference=True, чтобы модель compressed-tensors в формате FP8 оставалась в FP8, а её операции matmul выполнялись через аппаратно-ускоренные ядра FP8 (torch.nn.functional.scaled_mm, которое вызывает torch._scaled_mm_v2; в более старых версиях torch используется torch._scaled_mm), а не через деквантование весов обратно в BF16. Хранение весов в FP8 на протяжении всего инференса снижает потребление памяти и ускоряет вычисления. Это работает только для инференса, поэтому отключите этот параметр при дообучении.

Ядро FP8 поддерживает следующие схемы квантования.

Загрузка предварительно квантованной модели FP8

Ядра FP8 включаются по запросу: укажите use_optimized_inference=True, и они будут использоваться, если конфигурация модели задаёт квантование FP8 и доступно поддерживаемое GPU.

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

Декуантизация при загрузке

Без use_optimized_inference=True модель использует стандартный путь compressed-tensors: веса остаются сжатыми, а compressed-tensors разжимает их при первом проходе. Передайте dequantize=True, чтобы разжать их при загрузке, что необходимо для дообучения модели или сохранения её в исходной точности (например, BF16). В результате получается обычная плотная модель: схема compressed-tensors не привязывается, поэтому квантование активаций (например, 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",
)

Чекпоинт модели

Модели compressed-tensor определяются через соответствующий параметр в конфигурации. Приведённый ниже пример взят из файла config.json модели nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf.

Список содержит множество записей, что позволяет гибко настраивать процесс как во время сжатия, так и после него, однако записи, отвечающие за загрузку и инференс, можно упростить, сосредоточившись лишь на нескольких ключевых параметрах.

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

Файл конфигурации задаёт квантование группы конфигурации (group_0), которая включает квантование весов и активаций до fp8 со статической стратегией per-tensor. Модуль lm_head не квантуется, как показано в ключе ignore.

Для более подробного изучения весов модели используйте просмотрщик safetensors на карточке модели, чтобы увидеть квантованные веса, масштаб входных данных и масштаб весов для всех модулей nn.Linear.

При загрузке модели compressed-tensors с интеграцией ~quantizers.HFQuantizer целевые модули передаются в compressed-tensors: он присоединяет разрешённую quantization_scheme, устанавливает quantization_status, регистрирует параметры, которые хранятся в чекпоинте (weight в fp8, а также weight_scale и, для статической стратегии, input_scale), и устанавливает собственный forward-проход для них. Они остаются экземплярами nn.Linear, поэтому именно это показывает print — в последних версиях compressed-tensors они больше не оборачиваются в подкласс CompressedLinear. Модули, перечисленные в ignore, например lm_head, остаются без изменений.

При dequantize=False (значение по умолчанию) веса остаются сжатыми после завершения загрузки, а compressed-tensors разжимает всю модель при первом проходе. dequantize=True разжимает каждый вес в процессе его загрузки, поэтому для получения плотных весов не требуется пропускать данные через модель, а модули становятся обычными плотными слоями без 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

При use_optimized_inference=True слои, охваченные группой конфигурации fp8, заменяются на CompressedTensorsFP8Linear, который хранит вес fp8 и его масштаб в том расположении, которое ожидает его ядро row-wise matmul. Эти веса остаются в fp8, включая прохождение вперед.

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