Class VLLMWeightTransferEngine

trinity/utils/distributed.py:136–170

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

class VLLMWeightTransferEngine(WeightTransferEngine):
    """A helper class to manage NCCL weight synchronization using vLLM's API."""

    def __init__(self, master_address: str, master_port: int, world_size: int, group_name: str):
        """Initialize the NCCL process group for weight sync with vLLM's API."""
        from vllm.distributed.weight_transfer.nccl_engine import (
            NCCLWeightTransferEngine,
        )

        del group_name  # vLLM's NCCL engine does not require a group name
        self._model_update_group = NCCLWeightTransferEngine.trainer_init(
            dict(
                master_address=master_address,
                master_port=master_port,
                world_size=world_size,
            )
        )

    def sync_weight(self, iterator):
        """Perform the NCCL weight sync using vLLM's API."""
        from vllm.distributed.weight_transfer.nccl_engine import (
            NCCLTrainerSendWeightsArgs,
            NCCLWeightTransferEngine,
        )

        NCCLWeightTransferEngine.trainer_send_weights(
            iterator=iterator,
            trainer_args=NCCLTrainerSendWeightsArgs(
                group=self._model_update_group,
                packed=True,
            ),
        )

    def teardown(self):
        self._model_update_group.destroy()


class SGLangWeightTransferEngine(WeightTransferEngine):