deepinfer-eplb-migration

deepinfer/deepinfer/model_executor/eplb/__init__.py · a679bb8 · L195–278 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

195    def _initialize_shuffle(self) -> None:
196        """首次 rearrange：释放缓存显存防 OOM；单机时建 symm_mem 镜像走 IPC shuffle。"""
197        ep_group = self._ep_group
198        torch.cuda.empty_cache()
199        # TODO(NPU): NPU doesn't support symm_mem yet.
200        self._intranode_shuffle = PG.is_within_nvl_domain(ep_group) and PLATFORM.is_gpu
201        if not self._intranode_shuffle:
202            return
203
204        logger.info("[EPLB] shuffle intra node")
205        # 镜像只需容纳一层权重，shuffle 时逐层刷新复用。
206        self._symmetric_weight_buffers = [
207            symm_mem.empty(weight.shape, dtype=weight.dtype, device="cuda") for weight in self._expert_weights[0]
208        ]
209        self._symmetric_memory_handles = [symm_mem.rendezvous(buffer, ep_group) for buffer in self._symmetric_weight_buffers]
210        self._ipc_weight_buffers = [
211            [
212                handle.get_buffer(rank, buffer.shape, buffer.dtype)
213                for buffer, handle in zip(
214                    self._symmetric_weight_buffers,
215                    self._symmetric_memory_handles,
216                    strict=True,
217                )
218            ]
219            for rank in range(ep_group.size())
220        ]
221
222    @allow_cuda_sync()
223    def _rearrange_experts(self) -> None:
224        ep_group = self._ep_group
225        should_log_metrics = ep_group.rank() == 0
226        if should_log_metrics:
227            torch.cuda.synchronize()
228            start_time = time.perf_counter()
229
230        if not self._shuffle_initialized:
231            self._initialize_shuffle()
232            self._shuffle_initialized = True
233
234        # 窗口内各 logical expert 的 token 数：本地 physical 槽位聚合后全局 all-reduce
235        window_size, num_moe_layers, num_local_physical_experts = self._expert_count_window.shape
236        num_physical_experts = self.physical_to_logical_map.shape[1]
237        num_logical_experts = self.logical_count.shape[1]
238        local_physical_to_logical_map = self.physical_to_logical_map[
239            :,
240            ep_group.rank() * num_local_physical_experts : (ep_group.rank() + 1) * num_local_physical_experts,
241        ]
242        logical_expert_token_counts = torch.zeros(
243            window_size,
244            num_moe_layers,
245            num_logical_experts,
246            dtype=torch.int,
247            device=local_physical_to_logical_map.device,
248        )
249        logical_expert_token_counts.scatter_add_(
250            -1,
251            local_physical_to_logical_map.to(torch.int64).expand_as(self._expert_count_window),
252            self._expert_count_window,
253        )
254        dist.all_reduce(logical_expert_token_counts, group=ep_group)
255
256        next_physical_to_logical_map, next_logical_to_physical_map, next_logical_count = rebalance_experts_global(
257            logical_expert_token_counts,
258            num_physical_experts,
259            num_local_physical_experts,
260        )
261        if self._intranode_shuffle:
262            expert_model_parallel_shuffle_inplace_intranode(
263                self._ipc_weight_buffers,
264                self.physical_to_logical_map,
265                self._expert_weights,
266                next_physical_to_logical_map,
267                ep_group,
268            )
269        else:
270            expert_model_parallel_shuffle_inplace(
271                self.physical_to_logical_map,
272                self._expert_weights,
273                next_physical_to_logical_map,
274                ep_group,
275            )
276        self.physical_to_logical_map.copy_(next_physical_to_logical_map)
277        self.logical_to_physical_map.copy_(next_logical_to_physical_map)
278        self.logical_count.copy_(next_logical_count)