Новый стандарт хранения AI моделей: концепция и преимущества

Если вы будете блуждать по хабам AI моделей, таким как Hugging Face, вы начнете всё чаще встречать файлы с расширением .safetensors вместо привычных .bin или .pth.

В этом посте мы подробно разберем, что такое .safetensors, почему он появился и какие мощные преимущества он имеет по сравнению с традиционными способами хранения с технической точки зрения.

1. Что такое .safetensors?

.safetensors — это новый формат хранения тензоров, разработанный компанией Hugging Face.

Глубокие нейронные сети состоят из миллиардов параметров (весов), и данный формат выполняет функции их сохранения и загрузки. Он был создан для устранения серьезных недостатков (безопасность и скорость) традиционного формата хранения на основе модуля pickle в Python, который на протяжении долгого времени использовался как де-факто стандарт.

Проще говоря, это «более безопасный и быстрый формат для хранения моделей».

2. Почему он появился: проблемы с традиционным методом (Pickle)

Существующие модели PyTorch (.bin, .pth) используют для сериализации данных модуль pickle в Python. Однако у pickle есть серьезная проблема.

Уязвимость безопасности (Выполнение произвольного кода)

pickle не просто сохраняет данные, но и сохраняет сами объекты Python. В этом процессе существует возможность включения кода на Python, что позволяет злоумышленнику внедрить в модель код, разрушающий систему или украдущий личные данные. Как только пользователь без задней мысли выполнит load(), этот вредоносный код будет запущен.

Пример проблемной ситуации:

Пользователь загружает model.bin из интернета -> скрытый код внутри файла выполняется -> SSH ключи или пароли пользователя отправляются на сервер хакера.

.safetensors был разработан для полного предотвращения таких угроз безопасности.

3. Ключевые особенности .safetensors

3.1. Безопасность

.safetensors — это, как следует из названия, безопасно. Этот формат сохраняет только чистые данные тензоров и метаданные (в формате JSON). Поскольку в него нет пространства для исполняемого кода, вы можете спокойно загружать файлы даже из ненадежных источников.

3.2. Zero-Copy и скорость

При загрузке больших моделей LLM (моделей с большим объемом данных) или моделей Stable Diffusion скорость значительно повышается.

• Традиционный метод: копирование файла в память CPU -> обратная сериализация (Unpickling) -> преобразование в тензор -> перемещение на GPU. (Возникает ненужный процесс копирования)

• safetensors: использование технологии отображения памяти (Memory Mapping, mmap). Операционная система напрямую отображает файл в адрес памяти, поэтому данные можно использовать немедленно без ненужного процесса копирования. Это называется Zero-Copy.

3.3. Ленивая загрузка (Lazy Loading)

Вместо полной загрузки модели в память можно быстро загружать только необходимые части.

Например, если вы хотите проверить веса определенного слоя в 100-гигабайтном файле модели, вам придется прочитать все 100 ГБ с традиционным методом, тогда как .safetensors позволяет загружать только нужные части. Это очень полезно в условиях распределенного обучения или оптимизации вывода.

3.4. Совместимость с фреймворками

Не зависит от конкретных фреймворков глубокого обучения (таких как PyTorch и др.).

• PyTorch

• TensorFlow

• JAX

• PaddlePaddle

Он разработан для легкого чтения и записи во всех вышеперечисленных фреймворках.

4. Структура файла

Файл .safetensors имеет очень простую структуру.

1. Заголовок (Header): размещается в начале файла и имеет формат JSON. Содержит информацию о каждом тензоре, такую как имя, тип данных (dtype), форма (shape) и положение данных (offset).

2. Данные (Data): двоичные данные, следующие за заголовком. Содержат чистые значения тензоров, плотно заполненные внутри.

Благодаря этой структуре можно понять структуру модели, прочитав только заголовок, не читая весь файл.

4-1. Советы по практике: Как проверить заголовок и метаданные в терминале

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

Библиотека safetensors позволяет быстро просканировать только заголовочную информацию, не читая весь файл. Введите следующую однострочную команду Python в терминале, чтобы немедленно получить доступ к информации.

Конечно, вы также можете перейти на страницу загрузки модели Hugging Face и нажать на file info в правом верхнем углу для проверки. Однако здесь я объясню, как сделать это прямо в моем терминале.

Предварительная подготовка

Сначала необходимо установить библиотеку.

pip install safetensors

Команда (ввод в терминале)

Введите фактический путь к файлу на месте model.safetensors.

python -c "from safetensors import safe_open; \
with safe_open('model.safetensors', framework='pt', device='cpu') as f: \
    print('--- Metadata ---'); \
    print(f.metadata()); \
    print('\n--- Tensor Keys (Layers) ---'); \
    print(list(f.keys())[:5])" # Выводим только верхние 5 элементов, так как их много

Интерпретация результатов вывода

После выполнения этой команды вы получите две важные информации.

1. Metadata: информация, встроенная создателем модели. Если здесь указано format: gptq или quantization: int4, значит, это квантованная модель, даже если в имени файла это не указано. (Однако, если создатель оставил метаданные пустыми, может появиться None.)

2. Keys: названия слоев, из которых состоит модель. Это дает вам возможность понять структуру модели.

5. Сравнительная сводка: .bin (Pickle) vs .safetensors

6. Пример использования (Python)

Простой пример сохранения и загрузки тензора с использованием библиотеки safetensors.

import torch
from safetensors.torch import save_file, load_file

# 1. Создание и сохранение тензора
tensors = {
    "embedding": torch.zeros((1024, 512)),
    "attention": torch.rand((512, 512))
}

# Сохранение тензора в виде словаря в файл
save_file(tensors, "model.safetensors")

# 2. Загрузка файла
loaded = load_file("model.safetensors")
print(loaded["embedding"].shape) 
# Вывод: torch.Size([1024, 512])

7. Заключение

.safetensors — это не просто изменение расширения файла, а необходимая эволюция для безопасности и эффективности AI моделей. Основные сообщества, такие как Hugging Face, уже приняли его в качестве стандартного формата.

В будущем, когда вы будете загружать или распространять модели, желательно использовать .safetensors вместо .bin. Это обеспечит безопасность от угроз и значительно сократит время загрузки моделей.