CUDA не хватает памяти при использовании Llama3.1-8B для вывода ⇐ Python

[Anonymous]

Я написал простой скрипт Python, который использует библиотеку преобразователей HuggingFace вместе с torch для запуска Llama3.1-8B-instruct исключительно для вывода после ввода некоторых длинных данных. немного текста (около 10–20 тысяч токенов). Он отлично работает на моем ноутбуке, у которого есть графический процессор с 12 ГБ ОЗУ, но он также может получить доступ к общей памяти до 28 ГБ (я полагаю, из основной оперативной памяти системы?)
Однако, когда я запускаю код на виртуальной машине Windows «Standard NC4as T4 v3» с одним графическим процессором Tesla T4 с 16 ГБ ОЗУ, он очень быстро выдает эту ошибку: CUDA не хватает памяти. Пытался выделить XXX ГиБ. GPU
По моим расчетам, этот код должен работать нормально, учитывая доступную оперативную память. Тем не менее, я попытался сделать сценарий более эффективным с точки зрения использования памяти:

• Изменение механизма внимания — путем установки attn_implementation при создании экземпляра модели — сначала «sdpa». Затем попытался использовать Flash-внимание, но обнаружил, что невозможно установить пакет flashattn в Windows
• Использование xformers для запуска Enable_xformers_memory_efficient_attention - думаю, это также о механизме внимания, но не смог его установить/запустить.
• Используя torch.inference_mode()
• Установив такие флаги, как low_cpu_mem_usage = True во время создания экземпляра модели.
• Явное задание torch_dtype значения torch.float16 или torch.bfloat16 во время создания экземпляра модели.< /li>
  Использование BitsAndBytesConfig для запуска 8-битного и 4-битного квантования.

Последний шаг оказался единственным, который имел какой-либо эффект: модель успевает сгенерировать один или два ответа, но все равно не может завершить цикл и по-прежнему дает сбой при самом первом вывод. Если я правильно понял документацию, 4-битное квантование должно радикально снизить требования к памяти и сделать прямой проход с таким размером модели с помощью этого графического процессора почти тривиальным.
Итак, я Я с подозрением отношусь к дальнейшей оптимизации кода. Следующим шагом, очевидно, будет вручную сопоставить все слои модели с графическим процессором и процессором соответственно, определив device_map.
Итак, мой вопросы:

• Вероятно, это действительно проблема с памятью или это отвлекающий маневр?
• Если это память, что я могу сделать, кроме того, что уже пробовал?Это действительно проблема с памятью или это отвлекающий маневр?
• Если это память, что я могу сделать, помимо того, что уже пробовал? li>
  Если это не так, есть ли очевидные вещи, которые следует проверить? (Я подозреваю, что это связано с настройкой системы или установкой пакетов...)

Подробнее здесь: https://stackoverflow.com/questions/792 ... -inference