A helper class to manage NCCL weight synchronization using SGLang's API.
| 171 | |
| 172 | |
| 173 | class SGLangWeightTransferEngine(WeightTransferEngine): |
| 174 | """A helper class to manage NCCL weight synchronization using SGLang's API.""" |
| 175 | |
| 176 | def __init__(self, master_address: str, master_port: int, world_size: int, group_name: str): |
| 177 | """Initialize the NCCL process group for weight sync with SGLang's API.""" |
| 178 | self._model_update_group = init_process_group( |
| 179 | host=master_address, |
| 180 | port=master_port, |
| 181 | group_name=group_name, |
| 182 | backend="nccl", |
| 183 | world_size=world_size, |
| 184 | rank=0, |
| 185 | ) |
| 186 | |
| 187 | def sync_weight(self, iterator): |
| 188 | """Perform the NCCL weight sync using SGLang's API.""" |
| 189 | for _, param in iterator: |
| 190 | torch.distributed.broadcast(param, src=0, group=self._model_update_group) |
| 191 | |
| 192 | def teardown(self): |
| 193 | torch.distributed.destroy_process_group(self._model_update_group) |