torch.accelerator

Создано: 27 окт. 2024 | Последнее обновление: 19 июн. 2026

Ленивая инициализация и безопасность при fork

Среды выполнения ускорителей (CUDA, XPU, MPS и т. д.) инициализируются лениво — только при выполнении первой операции, обращающейся к устройству. Это гарантирует, что import torch и запросы возможностей не нарушат последующие вызовы fork(). Дополнительную информацию см. в разделе Отравление fork при многопроцессорной обработке.

Некоторые API должны быть доступны до создания fork (например, device_count(), is_available()), а некоторые бэкенды предоставляют механизмы, включаемые по желанию, для обеспечения безопасности при fork (например, CUDA через NVML, XPU через Level Zero Sysman). Чтобы обеспечить согласованность поведения и состояния среды выполнения между torch.accelerator и модулями отдельных бэкендов, следующие API делегируются каждому бэкенду:

is_available() / device_count() в идеале должны возвращать ответ, не запуская среду выполнения, поскольку DataLoader и подобные инструменты полагаются на их вызов до создания fork. Возможность этого зависит от бэкенда, поэтому torch.accelerator перенаправляет вызов соответствующей реализации бэкенда.
_lazy_call() используется для отложенного управления ГСЧ. Вызов manual_seed() до создания fork не должен принудительно инициировать среду выполнения. torch.accelerator оборачивает функцию обратного вызова для задания начального значения с помощью _lazy_call(), которая перенаправляет вызов в собственную очередь обратных вызовов бэкенда (CUDA, XPU, MTIA и т. д.). Каждый бэкенд управляет собственным флагом инициализации и очередью обратных вызовов. Если бэкенд не предоставляет _lazy_call (например, MPS), функция обратного вызова выполняется немедленно.

Этот пакет добавляет поддержку текущего ускорителя в Python.

`device_count`	Возвращает количество доступных текущих ускорителей.
`is_available`	Проверяет, доступен ли текущий ускоритель во время выполнения: он был собран, все необходимые драйверы доступны и видно хотя бы одно устройство.
`current_accelerator`	Возвращает устройство ускорителя, доступного во время компиляции.
`set_device_index`	Задает индекс текущего устройства.
`set_device_idx`	(Устарело) Задает индекс текущего устройства.
`current_device_index`	Возвращает индекс выбранного в данный момент устройства для текущего ускорителя.
`current_device_idx`	(Устарело) Возвращает индекс выбранного в данный момент устройства для текущего ускорителя.
`get_device_capability`	Возвращает возможности выбранного в данный момент устройства.
`set_stream`	Устанавливает текущий поток в указанный поток.
`current_stream`	Возвращает текущий выбранный поток для указанного устройства.
`synchronize`	Ожидает завершения всех ядер во всех потоках на указанном устройстве.
`device_index`	Менеджер контекста для задания индекса текущего устройства для текущего ускорителя.

Графы

Graph

Обертка для графа ускорителя, поддерживающая захват и воспроизведение.

Управление памятью

`empty_cache`	Освобождает всю неиспользуемую кэшированную память, занятую в данный момент кэширующим аллокатором, чтобы ее могли использовать другие приложения.
`empty_host_cache`	Освобождает всю неиспользуемую кэшированную память хоста (закрепленную), занятую в данный момент кэширующим аллокатором хоста, чтобы ее могли использовать другие приложения.
`get_memory_info`	Возвращает текущую информацию о памяти устройства для указанного индекса устройства.
`max_memory_allocated`	Возвращает максимальный объем памяти устройства текущего ускорителя, занятый тензорами, в байтах для указанного индекса устройства.
`max_memory_reserved`	Возвращает максимальный объем памяти устройства текущего ускорителя, управляемый кэширующим аллокатором, в байтах для указанного индекса устройства.
`memory_allocated`	Возвращает текущий объем памяти устройства ускорителя, занятый тензорами, в байтах для указанного индекса устройства.
`memory_reserved`	Возвращает текущий объем памяти устройства ускорителя, управляемый кэширующим аллокатором, в байтах для указанного индекса устройства.
`memory_stats`	Возвращает словарь статистики распределителя памяти устройства ускорителя для указанного индекса устройства.
`reset_accumulated_memory_stats`	Сбрасывает «накопительную» (историческую) статистику, отслеживаемую распределителем памяти текущего ускорителя, для указанного индекса устройства.
`reset_peak_memory_stats`	Сбрасывает статистику «пиковых значений», отслеживаемую распределителем памяти текущего ускорителя, для указанного индекса устройства.

Генератор случайных чисел

`get_rng_state`	Возвращает состояние ГСЧ генератора по умолчанию `torch.Generator` для текущего ускорителя в виде `torch.Tensor` типа `torch.uint8` для указанного устройства ускорителя.
`get_rng_state_all`	Возвращает список `torch.Tensor` типа `torch.uint8`, представляющих состояния ГСЧ всех устройств текущего ускорителя.
`initial_seed`	Возвращает начальное значение генератора по умолчанию `torch.Generator` для текущего ускорителя на указанном устройстве.

© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/accelerator.html