ptyrad.runtime.device

ptyrad.runtime.device#

Runtime environment and hardware configuration.

This module handles the initialization of PyTorch device settings and HuggingFace Accelerate environments, enabling seamless transitions between single-GPU, multi-GPU, CPU, and Apple Silicon (MPS) runtimes.

Functions

set_accelerator()

Initializes the HuggingFace Accelerator for distributed training.

set_gpu_device([gpuid])

Sets the default PyTorch computation device.
ptyrad.runtime.device.set_accelerator()[source]#

Initializes the HuggingFace Accelerator for distributed training.

This function attempts to load accelerate and configure it for PtyRAD’s specific needs. It enables split_batches=True for data loaders and sets find_unused_parameters=True for DistributedDataParallel (DDP) to prevent crashes when toggling gradients dynamically between iterations. It also verifies that the requested number of processes does not exceed the available hardware.

Returns:

The initialized Accelerator object if the library is available and correctly configured, otherwise None.

Return type:

accelerate.Accelerator or None

Raises:

ValueError – If the configured number of accelerate processes exceeds the actual number of available GPUs on the system.

ptyrad.runtime.device.set_gpu_device(gpuid=0)[source]#

Sets the default PyTorch computation device.

If ‘acc’ is passed, device management is explicitly deferred to the HuggingFace Accelerator. Otherwise, it configures the global PyTorch default device to the requested CUDA GPU, Apple Silicon (MPS), or CPU, handling fallbacks automatically if the requested hardware is unavailable.

Parameters:

gpuid (str or int, optional) –

The device identifier to use. Can be:

  • ”acc”: Defer device assignment to accelerate.

  • ”cpu”: Force execution on the CPU.

  • An integer (or string representation of an integer) for a specific NVIDIA GPU ID. Defaults to 0.

Returns:

The selected PyTorch device object, or None if assignment was deferred to accelerate.

Return type:

torch.device or None

Raises:

ValueError – If gpuid cannot be parsed as ‘acc’, ‘cpu’, or an integer.

Contents