Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / utils / moe_utils.py: 100%
44 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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 typing import TYPE_CHECKING, Optional, Any
19import torch.distributed as dist
20from torch import nn
22from hyper_parallel.core.activation_memory.wrapper import ActivationWrapper
24if TYPE_CHECKING:
25 from hyper_parallel.components.modules.moe import MoE
28def sync_and_update_expert_bias(
29 moe: "MoE",
30 lr: float = 1e-3,
31 tp_group: Optional[Any] = None,
32 cp_group: Optional[Any] = None,
33 dp_group: Optional[Any] = None,
34 num_recomputations: int = 1,
35) -> None:
36 """Synchronize tokens_per_expert across distributed ranks, then update expert bias.
38 This function handles distributed synchronization for expert bias updates in
39 MoE (Mixture of Experts) models. It should be called after optimizer.step() in
40 the training loop.
42 The synchronization is needed because:
43 - **DP**: Different ranks process different samples, need to aggregate global batch stats
44 - **TP+SP**: Standard TP does not need sync (all ranks see same input).
45 Only needed when TP is combined with sequence parallelism (TP+SP),
46 where activations are sequence-sharded across TP ranks.
47 - **CP**: Sequence dimension is sharded across ranks
49 For pure EP scenario (no DP/TP/CP), call moe.update_expert_bias() directly.
51 Args:
52 moe: The MoE module whose expert bias should be updated.
53 lr: Step size for the bias update. Defaults to ``1e-3``.
54 tp_group: Tensor parallel process group (ProcessGroup or GroupInfo).
55 Only needed when TP is combined with sequence parallelism (TP+SP),
56 where activations are sequence-sharded across TP ranks.
57 Standard TP (weight-only sharding) does not need this.
58 cp_group: Context parallel process group (ProcessGroup or GroupInfo).
59 Sync across sequence shards dimension.
60 dp_group: Data parallel process group (ProcessGroup or GroupInfo).
61 Sync across different samples to get global batch statistics.
62 num_recomputations: Number of forward executions per optimizer step.
63 Default ``1`` (normal training). Set to ``2`` when activation
64 checkpoint is enabled.
66 Example:
67 >>> # DP scenario
68 >>> sync_and_update_expert_bias(moe_layer, lr=1e-3, dp_group=dp_group_info)
69 >>>
70 >>> # TP×CP×DP scenario (reference: Megatron-LM)
71 >>> sync_and_update_expert_bias(
72 ... moe_layer,
73 ... lr=1e-3,
74 ... tp_group=tp_group_info,
75 ... cp_group=cp_group_info,
76 ... dp_group=dp_group_info,
77 ... )
79 Note:
80 Reference implementation: Megatron-LM megatron/core/transformer/moe/moe_utils.py
81 uses TP×CP×DP group synchronization for global batch statistics.
82 Model adapters may provide ``expert_bias_update_groups`` and
83 ``expert_bias_update_rate`` to override generic groups and update speed.
84 """
85 groups = getattr(moe, "expert_bias_update_groups", None)
86 if groups is None:
87 groups = (tp_group, cp_group, dp_group)
88 for group in groups:
89 if group is not None:
90 dist.all_reduce(moe.tokens_per_expert, group=getattr(group, "group", group))
92 moe.update_expert_bias(
93 lr=getattr(moe, "expert_bias_update_rate", lr),
94 num_recomputations=num_recomputations,
95 )
98def _get_moe_layers(model: "nn.Module") -> list:
99 """Collect all MoE sub-modules with ``enable_expert_bias=True``.
101 Uses duck-typing (``getattr(m, 'enable_expert_bias', False)``) instead
102 of ``isinstance`` to avoid importing platform-specific MoE classes in
103 platform-agnostic ``core/`` code.
105 Args:
106 model: Root module to search.
108 Returns:
109 List of MoE module instances.
110 """
111 # Transparent wrappers forward the child's attributes but own no router state.
112 return [
113 module for module in model.modules()
114 if not isinstance(module, ActivationWrapper) and getattr(module, 'enable_expert_bias', False)
115 ]
118class MoEMonitorCallback:
119 """Callback that automates expert bias updates across all MoE layers.
121 Provides a three-layer invocation architecture:
123 - **Layer 0**: :func:`sync_and_update_expert_bias` — per-module sync + update
124 - **Layer 1**: :meth:`on_step_end` — iterates all MoE layers, calling Layer 0
125 - **Layer 2**: :meth:`register` — registers optimizer post-hook for auto-trigger
127 Each layer can be used independently; Layer 0 is sufficient for basic usage,
128 while Layers 1–2 add convenience for production training loops.
130 Args:
131 model: The model containing MoE layers to monitor.
132 lr: Step size for expert bias updates. Defaults to ``1e-3``.
133 tp_group: Tensor parallel process group. Only needed for TP+SP
134 (sequence-sharded activations across TP ranks). Standard TP
135 does not need this — pass ``None``.
136 cp_group: Context parallel process group. Sync across sequence shards.
137 dp_group: Data parallel process group. Sync across different samples.
138 num_recomputations: Number of forward executions per optimizer step.
139 Default ``1``. Set to ``2`` when activation checkpoint is enabled.
141 Example:
142 >>> # Layer 1: manual trigger in training loop
143 >>> callback = MoEMonitorCallback(
144 ... model, lr=1e-3, dp_group=dp_group_info,
145 ... )
146 >>> for batch in dataloader:
147 ... loss = model(batch)
148 ... loss.backward()
149 ... optimizer.step()
150 ... callback.on_step_end() # auto-updates all MoE layers
151 ... optimizer.zero_grad()
152 >>>
153 >>> # Layer 2: auto-trigger via optimizer post-hook
154 >>> callback = MoEMonitorCallback(model, lr=1e-3, dp_group=dp_group_info)
155 >>> callback.register(optimizer)
156 >>> # optimizer.step() now auto-triggers callback.on_step_end()
158 Note:
159 Sync dimensions use ``dp_group``/``tp_group``/``cp_group`` (not ``ep_group``).
160 EP does not need synchronization because all EP ranks have identical
161 routing data (``tokens_per_expert`` shape = ``[total_experts]``).
162 """
164 def __init__(
165 self,
166 model: "nn.Module",
167 lr: float = 1e-3,
168 tp_group: Optional[Any] = None,
169 cp_group: Optional[Any] = None,
170 dp_group: Optional[Any] = None,
171 num_recomputations: int = 1,
172 ) -> None:
173 """Initialize MoEMonitorCallback."""
174 self.model = model
175 self.lr = lr
176 self.tp_group = tp_group
177 self.cp_group = cp_group
178 self.dp_group = dp_group
179 self.num_recomputations = num_recomputations
180 self._hook_handle = None
181 self.last_mean_aux_loss: Optional[float] = None
183 def on_step_end(self) -> None:
184 """Update expert bias for all MoE layers in the model.
186 Iterates over all :class:`MoE` sub-modules with
187 ``enable_expert_bias=True`` and calls
188 :func:`sync_and_update_expert_bias` for each one.
190 When ``load_balance_coeff > 0`` on any MoE layer, collects
191 ``last_aux_loss`` values and logs the mean.
192 """
193 moe_layers = _get_moe_layers(self.model)
194 if not moe_layers:
195 return
197 for moe in moe_layers:
198 sync_and_update_expert_bias(
199 moe,
200 lr=self.lr,
201 tp_group=self.tp_group,
202 cp_group=self.cp_group,
203 dp_group=self.dp_group,
204 num_recomputations=self.num_recomputations,
205 )
207 # Log aux_loss if any MoE layer has it enabled.
208 aux_losses = [
209 moe.last_aux_loss.item()
210 for moe in moe_layers
211 if moe.last_aux_loss is not None
212 ]
213 if aux_losses:
214 mean_aux_loss = sum(aux_losses) / len(aux_losses)
215 self.last_mean_aux_loss = mean_aux_loss
216 else:
217 self.last_mean_aux_loss = None
219 def register(self, optimizer: Any) -> None:
220 """Register an optimizer post-hook to auto-trigger :meth:`on_step_end`.
222 On PyTorch 2.0+, uses ``optimizer.register_step_post_hook``.
223 The hook fires after each ``optimizer.step()`` call, so that
224 expert bias is updated once parameters have been updated. Repeated
225 calls replace the existing hook registration.
227 Args:
228 optimizer: The optimizer instance to register the hook on.
230 Raises:
231 RuntimeError: If the optimizer does not support step post-hooks.
232 """
233 if not hasattr(optimizer, "register_step_post_hook"):
234 raise RuntimeError(
235 "Optimizer does not support register_step_post_hook. "
236 "Requires PyTorch 2.0+. Call on_step_end() manually instead."
237 )
238 self.remove()
239 self._hook_handle = optimizer.register_step_post_hook(
240 lambda *_: self.on_step_end(),
241 )
243 def remove(self) -> None:
244 """Remove the optimizer post-hook if registered."""
245 if self._hook_handle is not None:
246 self._hook_handle.remove()
247 self._hook_handle = None