Зачем это вообще нужно?
У вас есть RTX 3090. 24 ГБ VRAM. Кажется много, пока не пытаешься запустить Llama 3.1 405B или даже Qwen 2.5 72B в полной точности. Апгрейд на H200 с аппаратной поддержкой FP8 стоит как подержанный автомобиль. Но что если я скажу, что можно выжать из 3090 почти вдвое больше VRAM?
Что это за зверь — самодельный экстеншен?
Не ждите готового решения от NVIDIA. Они заинтересованы продавать новые карты. Мой экстеншен — это набор кастомных CUDA-ядер, которые эмулируют FP8 вычисления на RTX 30-й серии. Работает через PyTorch Custom Ops, подключается как обычный питоновский модуль.
Суть в трех словах: квантуем веса в FP8 на лету, вычисляем в FP16, результаты возвращаем в FP8. Память экономится, точность страдает минимально (если правильно настроить калибровку).
Внимание: это не магия. Производительность в вычислениях не увеличится — ускоряется только передача данных и экономится память. Для инференса LLM это часто важнее, чем raw FLOPS.
Как работает под капотом
Детали, от которых зависит, сломается ли ваша модель или заработает быстрее.
Динамическое квантование с калибровкой
Самое важное — правильно выбрать масштабирующий коэффициент. Берете небольшую часть датасета (100-200 примеров), пропускаете через модель в FP16, смотрите распределение активаций. На основе этого вычисляете scale для каждого слоя.
| Тип квантования | Точность | Сложность |
|---|---|---|
| Static (статика) | Низкая | Простая |
| Dynamic (динамика) | Высокая | Средняя |
| Per-token (по токену) | Максимальная | Сложная |
Ядра на Triton или чистый CUDA?
Попробовал оба варианта. Triton проще писать, но на RTX 3090 дает overhead в 10-15%. Чистый CUDA сложнее, зато работает ближе к металлу. В моем экстеншене — гибрид: критичные части на CUDA, остальное на Triton для простоты поддержки.
Что получилось в цифрах
Тестировал на Llama 3 70B (квантованная версия) и Mistral 8x22B. Конфигурация: одна RTX 3090, PyTorch 2.4.0, CUDA 12.4.
- Память под веса модели: с 13.5 ГБ (FP16) до 7.2 ГБ (FP8)
- Активации во время инференса: с 8.3 ГБ до 4.7 ГБ
- Максимальный контекст: увеличился с 4096 до 8192 токенов
- Perplexity на WikiText: рост на 0.8-1.2% против FP16
Не 2x экономии, как в теории, но 1.7-1.8x — уже отлично. Особенно когда каждый гигабайт на счету.
Альтернативы, которые не работают так же
Пробовал всё, что нашел. Вот что получилось.
Bitsandbytes 8-bit
Хорош для загрузки моделей, но не оптимизирует активации. Память экономит только на весах. Мой экстеншен работает и с весами, и с активациями.
TensorRT с fake quantization
Мощно, но требует конвертации модели в специальный формат. Каждый раз, когда меняете prompt или параметры генерации — боль. Мой solution работает напрямую в PyTorch.
Две RTX 3090 с NVLink
Решает проблему памяти, но стоит дорого (две карты + мост). Мой метод — бесплатно, если не считать времени на настройку.
Где это реально пригодится
Инференс больших LLM на одной карте
Те самые случаи, когда модель почти влезает в 24 ГБ, но не совсем. С FP8 — влезает с запасом. Особенно актуально для моделей с длинным контекстом.
Fine-tuning с ограниченной VRAM
Даже с LoRA иногда не хватает памяти под градиенты и оптимизаторы. FP8 освобождает место для большего batch size или более сложных адаптеров.
Мультимодальные модели
Когда нужно держать в памяти и текстовый энкодер, и визуальный, и LLM — каждый байт на счету. Здесь экономия в 1.8x меняет всё.
Не используйте для обучения с нуля. Градиенты в FP8 слишком шумные, модель может не сойтись. Только инференс и fine-tuning с замороженными весами.
Как начать использовать
Не буду врать — просто pip install не сработает. Нужно собрать из исходников. Но если вы читаете эту статью, скорее всего, у вас уже есть опыт с CUDA.
1 Установка зависимостей
PyTorch 2.4.0+, CUDA 12.4, Triton 3.0.0. Меньшие версии не поддерживают все фичи, нужные для нормальной работы.
2 Клонирование и сборка
git clone, затем python setup.py build_ext. Самая частая ошибка — несовместимость версий CUDA. Проверяйте, что версия в системе совпадает с той, под которую собран PyTorch.
3 Калибровка модели
Запускаете скрипт calib.py на вашем датасете. Занимает 10-15 минут. Результат — файл scales.json с коэффициентами для каждого слоя.
4 Загрузка с FP8
Вместо стандартной загрузки через transformers используете wrapper из экстеншена. Модель автоматически конвертирует веса в FP8 при загрузке.
Подводные камни, о которых молчат
- LayerNorm и внимания хуже всего переносят квантование. Приходится оставлять их в FP16
- Некоторые операции (особенно с очень маленькими числами) дают NaN в FP8. Нужно добавлять clipping
- Совместимость с другими оптимизациями (например, MoE) не гарантирована
- Debugging адский — ошибки в CUDA-ядре крешат весь Python без stack trace
Кому подойдет этот костыль
Если вы:
- Владелец RTX 3090/4090, который не хочет покупать H100
- Готовы потратить день на настройку ради экономии VRAM
- Работаете с инференсом, а не обучением с нуля
- Не боитесь покопаться в CUDA и Triton
Если вам нужно простое решение «из коробки» — посмотрите на Software FP8 или готовые квантованные модели. Мой экстеншен — для тех, кому нужен максимальный контроль.
Что будет дальше
NVIDIA медленно, но верно добавляет поддержку FP8 в драйверы для старых карт. В CUDA 12.5 уже есть экспериментальные API. Через год, возможно, всё это станет ненужным. Но пока — это единственный способ выжать из RTX 3090 почти H100-уровень экономии памяти.
И последнее: не верьте тем, кто говорит, что FP8 на старом железе невозможен. Возможен. Просто нужен правильный костыль. И терпение. Много терпения.
