fairseq2.device

This module provides abstractions for managing PyTorch devices and handling CUDA contexts.

Interfaces

class fairseq2.device.CudaContext

Bases: ABC

Represents an interface for interacting with CUDA runtime and device information.

abstract is_available() → bool

abstract device_count() → int

abstract get_device_properties(device: device) → Any

abstract memory_stats(device: device) → dict[str, Any]

abstract reset_peak_memory_stats() → None

class fairseq2.device.SupportsDeviceTransfer

Bases: ABC

Represents an object that can be transferred between devices.

abstract to(device: device, *, non_blocking: bool = False) → None

Transfers this object to the specified device.

If non_blocking is True, the transfer will be performed asynchronously if possible.

Classes

final class fairseq2.device.StandardCudaContext

Bases: CudaContext

Represents the standard implementation of CudaContext.

is_available() → bool

device_count() → int

get_device_properties(device: device) → Any

memory_stats(device: device) → dict[str, Any]

reset_peak_memory_stats() → None

Functions

fairseq2.device.get_default_device() → device

Returns the default device of this process.

The default device is determined by the following precedence:

1. If FAIRSEQ2_DEVICE environment variable is set, the specified device will be used.

2. If CUDA is enabled and CUDA_VISIBLE_DEVICES environment variable contains a single device, the specified device will be used.

3. If CUDA is enabled and LOCAL_RANK environment variable is set, the CUDA device at the specified index will be used.

4. CPU will be used.

Raises:

LocalRankOutOfRangeError – If LOCAL_RANK environment variable is less than zero or exceeds the number of available devices.

fairseq2.device.get_current_device() → device

Returns the current device of the calling thread.

The current device of a thread can be changed by using a device as a context manager. See here for more information.

Note

PyTorch does not currently expose a public API to retrieve the current device of a thread. If such an API becomes available in the future, this function will act as an alias for it.

Warning

This function might impose a slight performance cost. Avoid calling it in hot code paths.

import torch

from fairseq2.device get_current_device

# Default device used by PyTorch. Typically CPU.
default_device = torch.get_default_device()

assert get_current_device() == default_device

device = torch.device("cuda:0")

# Instruct PyTorch to use the specified device instead of the default
# device for tensor factory operations.
with device:
    assert get_current_device() == device

Exceptions

class fairseq2.device.LocalRankOutOfRangeError(local_rank: int, num_devices: int, device_type: str)

Bases: Exception

Raised by get_default_device() when LOCAL_RANK environment variable is less than zero or exceeds the number of available devices.