Несколько недель назад я закончил специализацию Deep Learning на курсере, это серия из 5 курсов, где рассказываются основы глубоких нейронных сетей, машинного зрения на их основе, а также обработки последовательностей.

Одной из первых тем, с которой я начал экспериментировать уже сам, стала Style Transfer. Быстро стало понятно, что чтобы нормально экспериментировать с такими вещами, нужны более мощный компьютер и видеокарта. Я решил купить GeForce RTX 3080. И вот, когда наступил ответственный момент и я решил прогнать бенчмарки, оказалось что текущая версия TensorFlow пока не поддерживает видеокарты на основе архитектуры Ampere. Что было дальше?

Что не работает?

Не работает не совсем корректное слово, на самом деле все работает просто сыпется куча предупреждений вида:

Your CUDA software stack is old. We fallback to the NVIDIA driver for some compilation. Update your CUDA version to get the best performance. The ptxas error was: ptxas fatal : Value 'sm_86' is not defined for option 'gpu-name'

При этом тренировка моделей проходит успешно и все остальное тоже работает корректно. Но почему на новой видеокарте какие-то ошибки, что вообще все это значит? Надо разобраться.

Compute Capability

Давайте начнем с того, что разберемся что такое sm_86. Оказывается, sm_86 — это значение так называемого compute capability видеокарты Nvidia. На вот этой странице можно найти соответствующие значение этого capability для каждой видеокарты.

Для RTX 3080 это значение — 8.6. Само по себе оно является просто некоторой версией, по которой можно сказать какие CUDA-фичи поддерживает ваша видеокарта. Чем больше, тем лучше, 8.6 на данный момент максимальное значение и это значит что видеокарты с архитектурой Ampere поддерживают несколько новых фишек, которые позволяют ускорять вычисления. Вот здесь можно найти описание этих фич с разбивкой по версии.

А причем тут TensorFlow? TF работает поверх CUDA, языка программирования, который позволяет выполнять вычисления на видеокартах. Но CUDA это язык высокого уровня, а значит он компилируется в некий машинный код. Именно для nvcc, компилятора CUDA, и нужен этот compute capability. Он позволяет компилятору понять, какие инструкции можно использовать, а значит и какие оптимизации можно или нельзя применять.

В TF при компиляции, вы указываете какие compute capability вы хотите поддерживать, и те части TF, которые написаны на CUDA компилируются сразу для всех необходимых вам версий. Таким образом одна и та же TF сборка может поддерживать несколько разных вариантов видеокарт. Но на момент написания текст, RTX 3080 еще достаточно свежая карта, поэтому текущие пакеты tensorflow и tensorflow-nightly собираются с compute capability < 8.6.

JIT компиляция

Второе непонятное слово в этой ошибки — ptxas. Мы уже говорили, что nvcc (компилятор CUDA) транслирует код на языке высокого уровня (по сути, диалекта C++) в некий машинный код. На самом деле процесс чуть сложнее.

Сначала nvcc превращает CUDA в PTX, аналог ассемблера для параллельных вычислений, которые заточены под выполнение на видеокартах. Информации про это не очень много, но если я правильно понимаю, то этот код не зависит от compute capability. При этом, если вы попросили nvcc сгенерить вам код с конкретным compute capability, то вы получите уже набор инструкций.

Но можно попросить nvcc отложить превращение ассемблера в инструкции. Тогда на выходе вы получите код, который нужно препроцессить утилитой ptxas (видимо, PTX assembler). Тогда вам не нужно заранее думать на каких видеокартах вы будете запускать сгенерированный код, но придется дополнительно ждать, пока ptxas сгенерирует машинный код конкретно под текущую видеокарту. В случае TF он это будет делать на каждом запуске интерпретатора.

Так что же случилось?

Теперь мы можем нормально интерпретировать ошибку. Все просто, последняя версия tensorflow скомпилирована без поддержки compute capability 8.6. А утилита ptxas еще не поддерживает этот compute capability, потому что tensorflow скомпилирован с библиотекой libcuda-11.0 (CUDA Toolkit). Поддержка архитектуры Ampere в этой библиотеке появилась только в версии 11.1.

Получается, что нам «просто» нужно собрать tensorflow из исходников. Заодно мне стало интересно, на какой ОС TF будет быстрее работать. Так что я решил собрать TF и прогнать бенчмарки в 3 средах: Windows 10, WSL Ubuntu 20.04 и «чистый» Ubuntu 20.04.

Сборка под Ubuntu

В принципе, весь процесс достаточно подробно описан в официальной документации: https://www.tensorflow.org/install/source. Важно не пропустить и чтобы у вас заработала поддержка GPU пройти и часть про GPU Support: https://www.tensorflow.org/install/gpu.

Если вы хотите запускать все это под WSL, то сначала надо обязательно пройти вот этот гайд и удостовериться, что CUDA видит вашу видеокарту. Например, с помощью утилиты deviceQuery. И не забудьте, что следуя инструкциям TensoFlow про поддержку GPU, не нужно устанавливать драйвер видоекарты внутри WSL, он уже есть в Windows.

Скажу прямо, развлечение на несколько часов, поэтому если хотите сэкономить время, можете воспользоваться моим докер-файлом: https://gist.github.com/Termina1/91f8f1122e3b7516408444b660704e8f. Предварительно нужно установить самый свежий драйвер для вашей системы.

По-моему опыту, часто бывает что релизные тэги не собираюсь, а main собирается. И наоборот, так что если один из вариантов не получился, попробуйте второй.

Сборка под Windows

Тут, к сожалению, схалтурить не получится. Есть официальный гайд: https://www.tensorflow.org/install/source_windows и отдельная секция в гайде про поддержку GPU: https://www.tensorflow.org/install/gpu#windows_setup.

Занимает это уйму времени, больше всего проблем было с установкой Microsoft Build Tools, почему-то последняя на тот момент версия не работала и выдавала какую-то странную ошибку при компиляции, причем в одной из стандартных библиотек. Помог даунгрейд на версию ниже. Также советую просто поставить MS Visual Studio с нужными библиотеками, а не отдельно Build Tools.

Важно! Обязательно установите первым делом MS Visual Studio в самом начале! Иначе часть cuda библиотек не установится.

Если при компиляции вы уберете флаг --config=opt, то компиляция будет сильно быстрее, но работать TF будет примерно в 9 раз медленнее.

Также при конфигурации билда для bazel есть вопрос про то, нужно ли отключить некий eigen strong inline для ускорения сборки. Советую ответить да, если у вас нет лишних 60-120 гигов оперативки. С включенной опцией у меня в процессе сборки съело все 32 гига, система ушла в своп и по сути компиляция зависла. Это можно либо обойти настройкой ограничения максимального объема памяти, либо количеством одновременно запущенных джоб в bazel.

В итоге, я запустил компиляцию без ограничений, большая часть проекта скомпилилась. Потом включил максимальное кол-во параллельных джоб на 4 и сутки ждал результат. При этом на бенчмарках прирост производительности где-то 1%, не уверен, что оно стоит того.

Кстати, если вы знаете что-такое этот eigen strong inline, ответьте мне на SO: https://stackoverflow.com/questions/65710930/what-is-eigen-strong-inline-optimization-when-building-tensorflow.

Результаты

Бенчмарки проводил с помощью утилиты ai-benchmark на системе с RTX 3080, AMD Ryzen 5800x, 32 GB RAM.

WSL 2 Ubuntu 20.04

Device Inference Score: 11785
Device Training Score: 9365
Device AI Score: 21150

Windows 10

Device Inference Score: 18114
Device Training Score: 18262
Device AI Score: 36376

Ubuntu 20.04

Device Inference Score: 20053
Device Training Score: 21121
Device AI Score: 41174

Результат достаточно предсказуемый. WSL 2 это просто виртуальная машина, поэтому ожидаемо, что результат там значительно хуже.

Что касается Windows 10, я скорее удивлен результатом. Но причина, вероятно, в том, что Windows 10 резервирует примерно 2GB оперативной памяти видеокарты под свои нужды и TF по сути видит 8 GB GPU RAM, вместо 10 GB.

Поэтому лучших результатов можно достичь только на Linux, при этом Windows 10 отстает не так сильно. К сожалению, разработка даже Jupyter под Windows 10 не самое приятное занятия, а WSL срезает перформанс новой видеокарты примерно до уровня RTX 2080 (и даже ниже).

Надеюсь кому-то этот текст будет полезным, если у вас есть вопросы, то пишите в твиттер: https://twitter.com/thought_sync или на GH: https://github.com/Termina1/blog/issues.