Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / moe_utils.py: 98%
55 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""MoE utilities for distributed training."""
17from types import SimpleNamespace
18from typing import TYPE_CHECKING, Optional, Any
20from hyper_parallel.core.fully_shard.hsdp_utils import GroupInfo
21from hyper_parallel.platform import get_platform
23if TYPE_CHECKING:
24 from hyper_parallel.platform.torch.common.moe import MoE
26platform = get_platform()
29def sync_and_update_expert_bias(
30 moe: "MoE",
31 lr: float = 1e-3,
32 tp_group: Optional[Any] = None,
33 cp_group: Optional[Any] = None,
34 dp_group: Optional[Any] = None,
35 num_recomputations: int = 1,
36) -> None:
37 """Synchronize tokens_per_expert across distributed ranks, then update expert bias.
39 This function handles distributed synchronization for expert bias updates in
40 MoE (Mixture of Experts) models. It should be called after optimizer.step() in
41 the training loop.
43 The synchronization is needed because:
44 - **DP**: Different ranks process different samples, need to aggregate global batch stats
45 - **TP+SP**: Standard TP does not need sync (all ranks see same input).
46 Only needed when TP is combined with sequence parallelism (TP+SP),
47 where activations are sequence-sharded across TP ranks.
48 - **CP**: Sequence dimension is sharded across ranks
50 For pure EP scenario (no DP/TP/CP), call moe.update_expert_bias() directly.
52 Args:
53 moe: The MoE module whose expert bias should be updated.
54 lr: Step size for the bias update. Defaults to ``1e-3``.
55 tp_group: Tensor parallel process group (ProcessGroup or GroupInfo).
56 Only needed when TP is combined with sequence parallelism (TP+SP),
57 where activations are sequence-sharded across TP ranks.
58 Standard TP (weight-only sharding) does not need this.
59 cp_group: Context parallel process group (ProcessGroup or GroupInfo).
60 Sync across sequence shards dimension.
61 dp_group: Data parallel process group (ProcessGroup or GroupInfo).
62 Sync across different samples to get global batch statistics.
63 num_recomputations: Number of forward executions per optimizer step.
64 Default ``1`` (normal training). Set to ``2`` when activation
65 checkpoint is enabled.
67 Example:
68 >>> # DP scenario
69 >>> sync_and_update_expert_bias(moe_layer, lr=1e-3, dp_group=dp_group_info)
70 >>>
71 >>> # TP×CP×DP scenario (reference: Megatron-LM)
72 >>> sync_and_update_expert_bias(
73 ... moe_layer,
74 ... lr=1e-3,
75 ... tp_group=tp_group_info,
76 ... cp_group=cp_group_info,
77 ... dp_group=dp_group_info,
78 ... )
80 Note:
81 Reference implementation: Megatron-LM megatron/core/transformer/moe/moe_utils.py
82 uses TP×CP×DP group synchronization for global batch statistics.
83 """
84 need_sync = tp_group is not None or cp_group is not None or dp_group is not None
86 if need_sync:
87 if tp_group is not None:
88 group_info = _ensure_group_info(tp_group)
89 platform.all_reduce(moe.tokens_per_expert, group_info)
90 if cp_group is not None:
91 group_info = _ensure_group_info(cp_group)
92 platform.all_reduce(moe.tokens_per_expert, group_info)
93 if dp_group is not None:
94 group_info = _ensure_group_info(dp_group)
95 platform.all_reduce(moe.tokens_per_expert, group_info)
97 moe.update_expert_bias(lr=lr, num_recomputations=num_recomputations)
100def _ensure_group_info(group: Any) -> GroupInfo:
101 """Convert ProcessGroup to GroupInfo if needed.
103 Args:
104 group: Either a ProcessGroup or a GroupInfo object.
106 Returns:
107 GroupInfo object.
108 """
109 if isinstance(group, GroupInfo):
110 return group
112 if hasattr(group, "group"):
113 return group
115 return SimpleNamespace(group=group)
118def _get_moe_layers(model: "nn.Module") -> list:
119 """Collect all MoE sub-modules with ``enable_expert_bias=True``.
121 Uses duck-typing (``getattr(m, 'enable_expert_bias', False)``) instead
122 of ``isinstance`` to avoid importing platform-specific MoE classes in
123 platform-agnostic ``core/`` code.
125 Args:
126 model: Root module to search.
128 Returns:
129 List of MoE module instances.
130 """
131 return [m for m in model.modules() if getattr(m, 'enable_expert_bias', False)]
134class MoEMonitorCallback:
135 """Callback that automates expert bias updates across all MoE layers.
137 Provides a three-layer invocation architecture:
139 - **Layer 0**: :func:`sync_and_update_expert_bias` — per-module sync + update
140 - **Layer 1**: :meth:`on_step_end` — iterates all MoE layers, calling Layer 0
141 - **Layer 2**: :meth:`register` — registers optimizer post-hook for auto-trigger
143 Each layer can be used independently; Layer 0 is sufficient for basic usage,
144 while Layers 1–2 add convenience for production training loops.
146 Args:
147 model: The model containing MoE layers to monitor.
148 lr: Step size for expert bias updates. Defaults to ``1e-3``.
149 tp_group: Tensor parallel process group. Only needed for TP+SP
150 (sequence-sharded activations across TP ranks). Standard TP
151 does not need this — pass ``None``.
152 cp_group: Context parallel process group. Sync across sequence shards.
153 dp_group: Data parallel process group. Sync across different samples.
154 num_recomputations: Number of forward executions per optimizer step.
155 Default ``1``. Set to ``2`` when activation checkpoint is enabled.
157 Example:
158 >>> # Layer 1: manual trigger in training loop
159 >>> callback = MoEMonitorCallback(
160 ... model, lr=1e-3, dp_group=dp_group_info,
161 ... )
162 >>> for batch in dataloader:
163 ... loss = model(batch)
164 ... loss.backward()
165 ... optimizer.step()
166 ... callback.on_step_end() # auto-updates all MoE layers
167 ... optimizer.zero_grad()
168 >>>
169 >>> # Layer 2: auto-trigger via optimizer post-hook
170 >>> callback = MoEMonitorCallback(model, lr=1e-3, dp_group=dp_group_info)
171 >>> callback.register(optimizer)
172 >>> # optimizer.step() now auto-triggers callback.on_step_end()
174 Note:
175 Sync dimensions use ``dp_group``/``tp_group``/``cp_group`` (not ``ep_group``).
176 EP does not need synchronization because all EP ranks have identical
177 routing data (``tokens_per_expert`` shape = ``[total_experts]``).
178 """
180 def __init__(
181 self,
182 model: "nn.Module",
183 lr: float = 1e-3,
184 tp_group: Optional[Any] = None,
185 cp_group: Optional[Any] = None,
186 dp_group: Optional[Any] = None,
187 num_recomputations: int = 1,
188 ) -> None:
189 """Initialize MoEMonitorCallback."""
190 self.model = model
191 self.lr = lr
192 self.tp_group = tp_group
193 self.cp_group = cp_group
194 self.dp_group = dp_group
195 self.num_recomputations = num_recomputations
196 self._hook_handle = None
197 self.last_mean_aux_loss: Optional[float] = None
199 def on_step_end(self) -> None:
200 """Update expert bias for all MoE layers in the model.
202 Iterates over all :class:`MoE` sub-modules with
203 ``enable_expert_bias=True`` and calls
204 :func:`sync_and_update_expert_bias` for each one.
206 When ``load_balance_coeff > 0`` on any MoE layer, collects
207 ``last_aux_loss`` values and logs the mean.
208 """
209 moe_layers = _get_moe_layers(self.model)
210 if not moe_layers:
211 return
213 for moe in moe_layers:
214 sync_and_update_expert_bias(
215 moe,
216 lr=self.lr,
217 tp_group=self.tp_group,
218 cp_group=self.cp_group,
219 dp_group=self.dp_group,
220 num_recomputations=self.num_recomputations,
221 )
223 # Log aux_loss if any MoE layer has it enabled.
224 aux_losses = [
225 moe.last_aux_loss.item()
226 for moe in moe_layers
227 if moe.last_aux_loss is not None
228 ]
229 if aux_losses:
230 mean_aux_loss = sum(aux_losses) / len(aux_losses)
231 self.last_mean_aux_loss = mean_aux_loss
232 else:
233 self.last_mean_aux_loss = None
235 def register(self, optimizer: Any) -> None:
236 """Register an optimizer post-hook to auto-trigger :meth:`on_step_end`.
238 On PyTorch 2.0+, uses ``optimizer.register_step_post_hook``.
239 The hook fires after each ``optimizer.step()`` call, so that
240 expert bias is updated once parameters have been updated.
242 Args:
243 optimizer: The optimizer instance to register the hook on.
245 Raises:
246 RuntimeError: If the optimizer does not support step post-hooks.
247 """
248 if not hasattr(optimizer, "register_step_post_hook"):
249 raise RuntimeError(
250 "Optimizer does not support register_step_post_hook. "
251 "Requires PyTorch 2.0+. Call on_step_end() manually instead."
252 )
253 self._hook_handle = optimizer.register_step_post_hook(
254 lambda *_: self.on_step_end(),
255 )
257 def remove(self) -> None:
258 """Remove the optimizer post-hook if registered."""
259 if self._hook_handle is not None:
260 self._hook_handle.remove()
261 self._hook_handle = None