torch.Tensor.record_stream

Tensor.record_stream(stream)

Отмечает, что тензор использовался этим потоком. При освобождении тензора необходимо убедиться, что его память не будет повторно использована для другого тензора, пока не завершится вся работа, поставленная в очередь в stream на момент освобождения.

Примечание

Кэш-аллокатор учитывает только поток, в котором был выделен тензор. Благодаря этому он корректно управляет жизненным циклом тензоров, используемых только в одном потоке. Однако если тензор используется в потоке, отличном от исходного, аллокатор может неожиданно повторно использовать память. Вызов этого метода сообщает аллокатору, в каких потоках использовался тензор.

Предупреждение

Этот метод лучше всего подходит для случаев, когда вы предоставляете функцию, создающую тензор во вспомогательном потоке, и хотите, чтобы пользователи могли работать с этим тензором, не задумываясь о безопасности потоков. Такие гарантии безопасности сопряжены с некоторыми затратами на производительность и предсказуемость (аналогично компромиссу между GC и ручным управлением памятью), поэтому, если вы самостоятельно управляете полным жизненным циклом тензоров, можно вместо этого вручную управлять событиями CUDA, чтобы не вызывать этот метод. В частности, при вызове этого метода во время последующих выделений аллокатор будет опрашивать зарегистрированный поток, чтобы проверить, завершились ли все операции; при этом возможна гонка с вычислениями во вспомогательном потоке, из-за которой память для выделения будет повторно использована или не использована повторно недетерминированным образом.

Можно безопасно использовать тензоры, выделенные во вспомогательных потоках, без record_stream(); необходимо вручную гарантировать, что все операции с тензором в потоках, отличных от потока создания, синхронизированы с потоком создания до освобождения тензора. Поскольку кэш-аллокатор CUDA гарантирует, что память будет повторно использована только в том же потоке создания, этого достаточно, чтобы гарантировать, что запись в память при последующих выделениях будет отложена до завершения операций в других потоках. (Как ни парадоксально, можно заметить, что на стороне CPU тензор уже повторно выделен, хотя ядра CUDA для старого тензора всё ещё выполняются. Это допустимо: операции CUDA с новым тензором будут корректно дожидаться завершения старых операций, поскольку все они выполняются в одном потоке.)

На практике это выглядит так:

with torch.cuda.stream(s0):
    x = torch.zeros(N)

s1.wait_stream(s0)
with torch.cuda.stream(s1):
    y = some_comm_op(x)

... some compute on s0 ...

# synchronize creation stream s0 to side stream s1
# before deallocating x
s0.wait_stream(s1)
del x

Обратите внимание, что при выборе момента для выполнения s0.wait_stream(s1) требуется определённое усмотрение. В частности, если бы мы ожидали сразу после some_comm_op, во вспомогательном потоке не было бы смысла; это было бы эквивалентно выполнению some_comm_op в s0. Вместо этого синхронизацию следует разместить в подходящий более поздний момент, когда, как ожидается, вспомогательный поток s1 завершит работу. Это место обычно определяют с помощью профилирования, например трассировок Chrome, создаваемых torch.autograd.profiler.profile.export_chrome_trace(). Если установить ожидание слишком рано, работа в s0 будет заблокирована до завершения s1, что помешает дальнейшему перекрытию обмена данными и вычислений. Если установить ожидание слишком поздно, будет использоваться больше памяти, чем строго необходимо (поскольку x будет дольше оставаться активным). Конкретный пример применения этих рекомендаций на практике см. в публикации: FSDP и CUDACachingAllocator.