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