torch.compiler.cudagraph_mark_step_begin

torch.compiler.cudagraph_mark_step_begin() [исходный код]

Указывает, что скоро начнётся новая итерация инференса или обучения.

Графы CUDA освобождают тензоры предыдущей итерации. При каждом вызове torch.compile начинается новая итерация, если только не ожидается обратный проход, который ещё не был вызван.

Если эта эвристика не срабатывает, например в следующем примере, отметьте начало итерации вручную с помощью этого API.

@torch.compile(mode="reduce-overhead")
def rand_foo():
    return torch.rand([4], device="cuda")


for _ in range(5):
    torch.compiler.cudagraph_mark_step_begin()
    rand_foo() + rand_foo()

Подробнее см. в разделе torch.compiler_cudagraph_trees # noqa: B950