trinity.utils.device module

Device detection and abstraction layer.

Unifies the differences among NPU / GPU / CPU devices for trinity modules.

class trinity.utils.device.DeviceType(*values)

Bases: str, Enum

Device type enum. Inherits str so it can be passed directly to APIs that require “npu”/”cuda” strings.

NPU = 'npu'

CUDA = 'cuda'

CPU = 'cpu'

trinity.utils.device.get_device_type() → DeviceType

Detect the currently available device type, with process-level caching.

Returns:

DeviceType.NPU / DeviceType.CUDA / DeviceType.CPU

trinity.utils.device.is_npu() → bool

Whether the current process is running in an NPU environment.

trinity.utils.device.is_cuda() → bool

Whether the current process is running in a CUDA environment.

trinity.utils.device.is_cpu() → bool

Whether the current process is running in a CPU environment.

trinity.utils.device.get_ray_resource_key() → str

Accelerator key name in the Ray cluster Resources dict.

NPU nodes report as “NPU”, GPU nodes report as “GPU”.

trinity.utils.device.get_collective_backend() → str

Collective communication backend name. NPU uses hccl, GPU uses nccl.

trinity.utils.device.get_device_capability() → int

Get major device capability version (device-agnostic).

Used to decide whether to enable meta tensor initialization for FSDP2. - NPU: returns 10 (supports meta tensor init, equivalent to sm90+) - CUDA: returns the actual major compute capability from torch.cuda - CPU: returns 0 (meta tensor not beneficial)