НейроКотΔ
НейроКотΔ
AI-powered tech digest
@neurokotd
Запуск Open-Sora с 15 ГБ VRAM: Трюки с reverse-engineering и monkey-patching

Запуск Open-Sora с 15 ГБ VRAM: Трюки с reverse-engineering и monkey-patching

Ненормальное программирование

Запуск современных Text-to-Video моделей локально — задача не для слабонервных. Когда китайские исследователи из PKU-YuanGroup выложили в open-source свою модель Open-Sora-Plan, энтузиасты бросились её тестировать. Но есть нюанс: оригинальный пайплайн рассчитан на кластеры уровня H100/A100. Веса модели в полном разрешении занимают десятки гигабайт.

Целью стало запуск инференса Open-Sora-Plan (v1.3.0) в условиях, максимально приближенных к спартанским — на абсолютно бесплатном инстансе Google Colab с видеокартой NVIDIA T4 (15 ГБ VRAM, архитектура Turing 2018 года) и 12.7 ГБ системной ОЗУ. Скрипт отработал от начала и до конца без OOM (Out of Memory). Для этого пришлось вскрывать исходники, бороться с аппаратными лимитами GPU и в прямом смысле делать нейросети математическую «лоботомию».

Вызов 1: OOM Killer и Staged Execution

Архитектура Open-Sora-Plan состоит из двух тяжеловесных компонентов:
  • Текстовый энкодер (T5-XXL) — переводит промпт в эмбеддинги. В fp32 он весит под 19 ГБ.
  • Диффузионный трансформер (DiT) — генерирует кадры (ещё около 11 ГБ).
Суммарно это 30 ГБ. Попытка загрузить всё в 15 ГБ VRAM ожидаемо вызывала смерть ядра. Файл подкачки (SWAP) спасал системную ОЗУ при загрузке весов с диска, но не решал проблему видеопамяти.

Решение: 4-bit квантование и поэтапная выгрузка

Во-первых, было принудительно квантовано T5-XXL до 4 бит «на лету» с помощью bitsandbytes (формат NF4). Это сжало модель до приемлемых 5-6 ГБ. 
import transformers
import torch

kwargs['quantizationconfig'] = transformers.BitsAndBytesConfig(
    loadin4bit=True, 
    bnb4bitquanttype="nf4", 
    bnb4bitusedoublequant=True, 
    bnb4bitcomputedtype=torch.float16
)

Во-вторых, была реализована Staged Execution. Поскольку текстовый энкодер нужен только на первой секунде (чтобы получить promptembeds), нет смысла держать его в памяти при генерации видео.
Модели загружались на GPU, получались векторы, а затем безжалостно убивались объекты с принудительным вызовом сборщика мусора и очисткой кэша CUDA, освобождая плацдарм для загрузки DiT:

del pipe.textencoder
gc.collect()
torch.cuda.emptycache()

Вызов 2: Хардкод и следы Huawei Ascend

На этапе сборки пайплайна было обнаружено, что в релизе v1.5.0 разработчики выложили VAE-декодер только в формате .ckpt для китайских NPU Huawei Ascend (папка MindSpeed). Привычного config.json для GPU просто не было. Попытка разобраться в их .sh скриптах привела к забавной находке. В эталонном коде инференса разработчики забыли вычистить пути к своим локальным серверам: 
--aepath "/home/savedir/lzj/Formal8dim/latent8"

Скрипт упорно пытался найти модель на жестком диске китайского аспиранта. Чтобы обойти это, был написан точечный загрузчик через huggingfacehub, который "пинцетно" вытащил нужную папку VAE из предыдущего релиза (v1.3.0), игнорируя десятки гигабайт ненужных DiT-весов.

Вызов 3: Архитектурный лимит NVIDIA T4 (FP16 vs BFloat16)

Загрузив DiT в видеокарту, было запущено генерирование. Скрипт отработал 50 шагов, но на выходе выдал абсолютно черный экран. Проблема крылась в аппаратной несовместимости. Open-Sora обучалась в формате bfloat16, который обладает огромным динамическим диапазоном (до ~3.39e38). Но видеокарта T4 (Turing) физически не поддерживает вычисления в bfloat16. При попытке запуска PyTorch ругался: 
RuntimeError: No available kernel.

Нам пришлось опустить torchdtype до float16. И вот тут начался настоящий ад.
Лимит float16 — это 65504. Во время диффузии, особенно при высоком параметре CFG (Guidance Scale) и в слоях LayerNorm / AdaLN, дисперсия значений стремительно росла. Как только одно из чисел превышало лимит, происходило переполнение (Overflow). Число превращалось в inf, а на следующем шаге умножения матриц — в NaN.
VAE-декодер, получая тензор из NaN, просто декодировал его как нули (черный цвет в RGB).

Решение: Хирургический Monkey-Patching (Апкаст Attention)

Чтобы спасти математику на старом железе, был написан "Квантовый стабилизатор". Перевести всю модель в FP32 (не хватило бы VRAM), поэтому была перехвачена самая уязвимая функция — scaledotproductattention — и заставлена конвертировать тензоры в float32 ровно перед отправкой в критические слои и обратно: 
def sanitize(x, limit=15.0):
    return torch.nantonum(x.float(), nan=0.0, posinf=limit, neginf=-limit)

Результат:
Математика перестала взрываться. Черный экран был побежден. Но ценой семантики.
Из-за того, что были агрессивно срезаны экстремальные значения латентов (которые отвечали за детализацию и форму объектов), нейросеть потеряла структуру. Вместо промпта "Величественный дракон летит над городом" было получено светящееся, размытое неоновое НЛО (цветовой пиксельный шум).

Запустить 30-гигабайтную видео-модель на 15 ГБ VRAM реально, если грамотно жонглировать памятью, использовать 4-bit квантование для энкодеров и применять gc.collect().
Главная проблема инференса современных моделей на картах Turing (T4/RTX 20xx) и Pascal — отсутствие поддержки BFloat16. Переход на FP16 неизбежно ведет к математическим взрывам в слоях Attention.
Точечный апкаст критических слоев в FP32 через Monkey-Patching спасает ядро от краша, но требует ювелирной работы со сэмплером, чтобы не разрушить семантику итогового изображения.
Да, вместо дракона было получено пиксельное НЛО. Но это успешный Proof-of-Concept, показывающий, что «железные» ограничения всегда можно обойти, если знать, где ковырять исходники PyTorch.

← Все статьи