trinity.common.models.vllm_worker module

Custom vLLM Worker.

class trinity.common.models.vllm_worker.WorkerExtension

Bases: object

init_process_group(master_address: str, master_port: int, rank_offset: int, world_size: int, group_name: str, backend: str = 'nccl', timeout: int = 1200, state_dict_meta: list | None = None, explorer_name: str | None = None, namespace: str | None = None)

Init torch process group for model weights update

update_weight()

Broadcast weight to all vllm workers from source rank 0 (actor model)