Масштабирование и запуск JAX-моделей на графических процессорах NVIDIA.

1. Что лучше всего проверить в первую очередь перед отладкой производительности JAX GPU?

   Увеличьте размер партии.

   Запустите jax.devices() и убедитесь, что JAX видит графический процессор.

   Преобразовать все массивы JAX в NumPy.

   Отключить jax.jit

2. Если jax.default_backend() выводит cpu , а jax.devices() показывает только CpuDevice , какова наиболее вероятная причина?

   Модель слишком велика для графического процессора.

   Проблема может быть связана с библиотекой jaxlib, работающей только с ЦП, несоответствием драйверов или запуском контейнера без доступа к графическому процессору.

   Размер партии установлен слишком большим.

   jax.jit еще не был вызван

3. Почему первый вызов функции jax.jit может быть намного медленнее, чем последующие?

   JAX всегда выполняет первый вызов на центральном процессоре.

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

   JAX отключает память графического процессора при первом вызове.

   Градиенты вычисляются только при первом вызове.

4. Какова распространённая причина случайной перекомпиляции в JAX?

   Повторное использование той же формы и типа данных входных данных.

   Хранение пакетов данных на графическом процессоре.

   Изменение формы входных данных между вызовами

   Вызов функции block_until_ready() перед измерением времени.

5. Почему внутри функции @jax.jit условие if x > 0: на отслеживаемом массиве вызывает ошибку?

   JAX не поддерживает оператор >

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

   Для сравнения следует использовать jnp.greater вместо >

   jit навсегда отключает все элементы управления потоком выполнения Python.

6. Почему необходимо вызывать block_until_ready() перед остановкой таймера во время работы с графическим процессором?

   Это освобождает память графического процессора перед следующим шагом.

   JAX распределяет задачи на графическом процессоре асинхронно, поэтому без блокировки вы можете измерить только время, затраченное Python на добавление задачи в очередь.

   Функция перекомпилируется для обеспечения точного измерения времени.

   Результат передается обратно в центральный процессор.

7. Вызов функции float(loss) на каждом шаге внутри цикла обучения может замедлить его, потому что это:

   Функция потерь перекомпилируется на каждом шаге.

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

   Изменяет тип данных функции потерь на float64.

   Отключает кэш компиляции JAX.

8. Почему в процессе обучения данные преобразуются в пакеты фиксированного размера, а неполный итоговый пакет отбрасывается?

   Для более эффективного перемешивания данных

   Чтобы сохранить постоянную форму входных данных, этап компиляции с использованием JIT-преобразования не будет перекомпилироваться для меньшего последнего пакета.

   Поскольку JAX не может обрабатывать пакеты размером более 512 символов,

   Для уменьшения количества параметров модели

9. Как обновляются параметры на этапе обучения JAX?

   Весовые коэффициенты изменяются непосредственно внутри модели.

   Чистая функция принимает на вход параметры, состояние и пакет данных и возвращает новые параметры и состояние на выходе — без изменения параметров на месте.

   jax.jit автоматически обновляет веса.

   Градиенты применяются непосредственно драйвером графического процессора.

10. Что требуется для установки параметра implementation="cudnn" в jax.nn.dot_product_attention ?

    Резервный режим работы ЦП и входные данные типа float64

    Входные данные типа bfloat16 или float16 и вычислительные возможности >= 8.0

    По крайней мере две видеокарты.

    Отключение причинно-следственной маскировки

11. Почему наивный механизм внимания использует больше памяти, чем ядро ​​cuDNN с интегрированным модулем на длинных последовательностях?

    Он хранит веса модели дважды.

    Это позволяет реализовать полную матрицу внимания (B, N, T, S) , которая растет пропорционально квадрату длины последовательности.

    По умолчанию он работает в формате float64.

    Это позволяет хранить данные на центральном процессоре.

12. Что же на самом деле создает параллелизм при параллельном обучении с использованием JAX?

    Переписывание этапа обучения с явным использованием циклов по графическим процессорам.

    Способ размещения массивов — сегментированные данные и реплицированные параметры — при этом код этапа обучения остается неизменным.

    Вызов jax.jit дважды.

    Увеличение скорости обучения на каждый графический процессор.

13. Что означает PartitionSpec('data', None) для двумерного массива на сетке с осью 'data' ?

    Продублируйте весь массив на каждом графическом процессоре.

    Разметьте нулевое измерение по оси 'data' и продублируйте измерение 1.

    Разделите оба измерения по оси 'data'

    Переместите массив на центральный процессор.

14. Как преобразователь из Урока 7 получает причинно-следственное (только обратное) внимание?

    На каждом этапе вручную создается треугольная матрица маски.

    В качестве attention_fn объекта nnx.MultiHeadAttention передается causal_sdpa (которая вызывает dot_product_attention с is_causal=True ).

    Он сортирует токены в обратном порядке.

    Он использует отдельную библиотеку декодера.

15. Что сохраняется, а что нет при сохранении модели NNX с помощью Orbax StandardCheckpointer ?

    Весь Python-класс, включая методы, сериализуется с помощью сериализации (pickle).

    Параметры ( nnx.state(model, nnx.Param) ) сохраняются; структура графа — нет, она берется из класса Python при восстановлении.

    Сохраняется только состояние оптимизатора.

    Ничего не сохраняется до окончания обучения.

16. Какую проблему решает AOT-компиляция ( lower() → compile() ) при обслуживании?

    Это уменьшает количество параметров модели.

    Это разделяет компиляцию и выполнение, поэтому компиляция выполняется один раз при запуске, а запросы обрабатываются без задержки холодного запуска.

    Он автоматически преобразует модель в формат float16.

    Модель распределяется между несколькими графическими процессорами.

17. Какой путь развертывания создает переносимый артефакт, который работает в любой среде выполнения JAX без исходного кода модели?

    jax.jit в процессе

    jax.export — сериализованные байты StableHLO

    Вызов модели внутри цикла Python.

    AOT-компиляция с использованием lower().compile()

Отправить ответы