JAX незаметно забирает массив с GPU: как поймать перенос через transfer guard

Модель работает на GPU, но отладочная печать массива может незаметно вернуть данные на CPU. В JAX такие переносы можно обнаружить через transfer guard: сначала увидеть их в журнале, затем запретить нежелательные операции на выбранном участке.

Где прячется перенос

По документации JAX, guard различает явные вызовы device_put/device_get и неявные перемещения данных. Режим log сообщает о неявных переносах, disallow запрещает их с RuntimeError. Явные операции оба режима разрешают. Поэтому отсутствие ошибки ещё не означает, что данные вообще никуда не копировались.

Причина неочевидности — асинхронное выполнение. Python может получить объект массива раньше завершения вычисления. Его форму можно узнать без чтения значений, а печать значений или преобразование в NumPy заставляют дождаться результата. В журналировании обучающего цикла это принципиально разные действия: сообщение о размере тензора и вывод его содержимого требуют разной работы.

Проверка на одном участке

Возьмите короткий воспроизводимый шаг своей модели. Сохраните исходный вариант и включите наблюдение только вокруг подозрительного кода. Так проще сопоставить сообщение с действием, чем разбирать журнал всего приложения.

Учебный пример для среды с GPU или TPU: сначала разместите массив на ускорителе, затем проверьте чтение его значений. Здесь нет измерений скорости или обещанного ускорения.

import jax
import numpy as np

x = jax.device_put(np.arange(8, dtype=np.float32))
x.block_until_ready()

with jax.transfer_guard("log"):
    print(x)  # потенциальный неявный перенос на хост

Для следующего прогона создайте свежий массив и замените log на disallow. Уже полученные на хост значения могут не потребовать нового переноса. На CPU получение буфера разрешено независимо от уровня guard — такой запуск не проверит запрет чтения с ускорителя.

Что делать с найденным

Запишите три вещи: какая строка вызвала перенос, нужны ли значения именно на CPU и как часто выполняется эта строка. Если вывод нужен только для диагностики, попробуйте сократить его частоту. Если потребителю действительно нужен массив на хосте, обозначьте границу явным device_get и оцените её стоимость отдельно.

После разбора можно оставить disallow в проверке критичного участка. Guard обнаруживает переносы, но не доказывает ускорение. Сравните время всего нужного шага до и после изменения и проверьте результат вычисления. Контекст guard действует в текущем потоке; новый поток использует глобальную настройку. Если приложение многопоточное, учитывайте это при выборе места проверки.