Initialize a distributed worker and task queue for single-node multi-GPU setup. Args: fd_config (FDConfig): Arguments related to inference, containing attributes such as weight_dtype, act_dtype, mp_size, hidden_size, head_dim, num_attentio
(self, fd_config: FDConfig, ranks: int = 1, local_rank: int = 0)
| 148 | """ |
| 149 | |
| 150 | def __init__(self, fd_config: FDConfig, ranks: int = 1, local_rank: int = 0) -> None: |
| 151 | """ |
| 152 | Initialize a distributed worker and task queue for single-node multi-GPU setup. |
| 153 | Args: |
| 154 | fd_config (FDConfig): Arguments related to inference, containing |
| 155 | attributes such as weight_dtype, act_dtype, mp_size, hidden_size, head_dim, |
| 156 | num_attention_heads, and ffn_hidden_size. |
| 157 | """ |
| 158 | self.ranks = ranks |
| 159 | self.local_rank = local_rank |
| 160 | self.fd_config = fd_config |
| 161 | self.parallel_config = fd_config.parallel_config |
| 162 | self.cache_config = fd_config.cache_config |
| 163 | self.scheduler_config = fd_config.scheduler_config |
| 164 | self.eplb_config = fd_config.eplb_config |
| 165 | |
| 166 | # TODO(gongshaotian): Use worker factory to get worker |
| 167 | self.worker = get_worker(fd_config=fd_config, local_rank=self.local_rank, rank=self.ranks) |
| 168 | |
| 169 | self.max_chips_per_node = 16 if current_platform.is_iluvatar() else 8 |
| 170 | self.speculative_decoding = fd_config.speculative_config.method is not None |
| 171 | self.enable_overlap_schedule = self.scheduler_config.enable_overlap_schedule and ( |
| 172 | not self.speculative_decoding |
| 173 | ) |
| 174 | self.cached_control_reqs = [] |
| 175 | |
| 176 | def init_control(self): |
| 177 | engine_worker_queue_port = self.parallel_config.local_engine_worker_queue_port |
nothing calls this directly
no test coverage detected