trinity.utils.distributed module

For distributed training with multiple process groups.

trinity.utils.distributed.is_ipv6_address(ip_str: str) → bool

trinity.utils.distributed.get_available_port() → int

trinity.utils.distributed.get_endpoint(host: str, port: int) → str

trinity.utils.distributed.is_port_available(port: int, host='127.0.0.1') → bool

trinity.utils.distributed.init_process_group(host: str, port: int, group_name: str, backend: str | Backend = 'nccl', timeout: float | None = None, world_size: int = -1, rank: int = -1, pg_options: Any | None = None, device_id: device | None = None)

This function is used to initialize the process group. It requires torch >= 2.6.0

class trinity.utils.distributed.WeightTransferEngine

Bases: object

abstractmethod sync_weight(iterator)

Perform the weight sync.

abstractmethod teardown()

Tear down the weight sync group.

static create(engine_type: str, master_address: str, master_port: int, world_size: int, group_name: str)

Factory method to create the appropriate weight transfer engine based on the rollout engine type.

class trinity.utils.distributed.VLLMWeightTransferEngine(master_address: str, master_port: int, world_size: int, group_name: str)

Bases: WeightTransferEngine

A helper class to manage weight synchronization using vLLM’s API.

Device-aware: uses HCCL engine on NPU and NCCL engine on GPU. Both engines expose the same trainer_init / trainer_send_weights API surface, so the only difference is the engine class and args dataclass used.

__init__(master_address: str, master_port: int, world_size: int, group_name: str)

Initialize the process group for weight sync with vLLM’s API.

sync_weight(iterator)

Perform the weight sync using vLLM’s API (HCCL on NPU, NCCL on GPU).

teardown()

Tear down the weight sync group.

class trinity.utils.distributed.SGLangWeightTransferEngine(master_address: str, master_port: int, world_size: int, group_name: str)

Bases: WeightTransferEngine

A helper class to manage NCCL weight synchronization using SGLang’s API.

__init__(master_address: str, master_port: int, world_size: int, group_name: str)

Initialize the NCCL process group for weight sync with SGLang’s API.

sync_weight(iterator)

Perform the NCCL weight sync using SGLang’s API.

teardown()

Tear down the weight sync group.