Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / moe_utils.py: 100%
40 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
20import torch.distributed as dist
21from torch import nn
23from hyper_parallel.core.activation_memory.wrapper import ActivationWrapper
25if TYPE_CHECKING:
26 from hyper_parallel.components.modules.moe import MoE
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 Model adapters may provide ``expert_bias_update_groups`` and
84 ``expert_bias_update_rate`` to override generic groups and update speed.
85 """
86 groups = getattr(moe, "expert_bias_update_groups", None)
87 if groups is None:
88 groups = (tp_group, cp_group, dp_group)
89 for group in groups:
90 if group is not None:
91 dist.all_reduce(moe.tokens_per_expert, group=getattr(group, "group", group))
93 moe.update_expert_bias(
94 lr=getattr(moe, "expert_bias_update_rate", lr),
95 num_recomputations=num_recomputations,
96 )
99class MoEMonitorCallback:
100 """Callback that automates expert bias updates across all MoE layers.
102 Provides a three-layer invocation architecture:
104 - **Layer 0**: :func:`sync_and_update_expert_bias` — per-module sync + update
105 - **Layer 1**: :meth:`on_step_end` — iterates all MoE layers, calling Layer 0
106 - **Layer 2**: :meth:`register` — registers optimizer post-hook for auto-trigger
108 Each layer can be used independently; Layer 0 is sufficient for basic usage,
109 while Layers 1–2 add convenience for production training loops.
111 Args:
112 model: The model containing MoE layers to monitor.
113 lr: Step size for expert bias updates. Defaults to ``1e-3``.
114 tp_group: Tensor parallel process group. Only needed for TP+SP
115 (sequence-sharded activations across TP ranks). Standard TP
116 does not need this — pass ``None``.
117 cp_group: Context parallel process group. Sync across sequence shards.
118 dp_group: Data parallel process group. Sync across different samples.
119 num_recomputations: Number of forward executions per optimizer step.
120 Default ``1``. Set to ``2`` when activation checkpoint is enabled.
122 Example:
123 >>> # Layer 1: manual trigger in training loop
124 >>> callback = MoEMonitorCallback(
125 ... model, lr=1e-3, dp_group=dp_group_info,
126 ... )
127 >>> for batch in dataloader:
128 ... loss = model(batch)
129 ... loss.backward()
130 ... optimizer.step()
131 ... callback.on_step_end() # auto-updates all MoE layers
132 ... optimizer.zero_grad()
133 >>>
134 >>> # Layer 2: auto-trigger via optimizer post-hook
135 >>> callback = MoEMonitorCallback(model, lr=1e-3, dp_group=dp_group_info)
136 >>> callback.register(optimizer)
137 >>> # optimizer.step() now auto-triggers callback.on_step_end()
139 Note:
140 Sync dimensions use ``dp_group``/``tp_group``/``cp_group`` (not ``ep_group``).
141 EP does not need synchronization because all EP ranks have identical
142 routing data (``tokens_per_expert`` shape = ``[total_experts]``).
143 """
145 def __init__(
146 self,
147 model: "nn.Module",
148 lr: float = 1e-3,
149 tp_group: Optional[Any] = None,
150 cp_group: Optional[Any] = None,
151 dp_group: Optional[Any] = None,
152 num_recomputations: int = 1,
153 ) -> None:
154 """Initialize MoEMonitorCallback."""
155 self.model = model
156 self.lr = lr
157 self.tp_group = tp_group
158 self.cp_group = cp_group
159 self.dp_group = dp_group
160 self.num_recomputations = num_recomputations
161 self._hook_handle = None
163 @property
164 def last_mean_aux_loss(self) -> Optional[torch.Tensor]:
165 """Return the mean router metric as an on-device scalar Tensor.
167 Keep the mean on-device so graph capture never reads a Python scalar.
168 Logging callbacks convert the result only when formatting output.
169 Writers detach ``last_aux_loss`` at the source, so the result carries
170 no autograd graph in production. Auxiliary loss remains observable
171 when expert-bias updates are disabled.
172 """
173 # Transparent wrappers expose the same router; its child owns the metric.
174 aux_losses = [
175 module.last_aux_loss
176 for module in self.model.modules()
177 if not isinstance(module, ActivationWrapper) and getattr(module, "last_aux_loss", None) is not None
178 ]
179 return torch.stack(aux_losses).mean() if aux_losses else None
181 def on_step_end(self) -> None:
182 """Update expert bias for all MoE layers in the model.
184 Iterates over all :class:`MoE` sub-modules with
185 ``enable_expert_bias=True`` and calls
186 :func:`sync_and_update_expert_bias` for each one.
188 Auxiliary loss is read separately through ``last_mean_aux_loss`` so
189 bias updates do not synchronize device scalars for logging each step.
190 """
191 for moe in self.model.modules():
192 if isinstance(moe, ActivationWrapper) or not getattr(moe, "enable_expert_bias", False):
193 continue
194 sync_and_update_expert_bias(
195 moe,
196 lr=self.lr,
197 tp_group=self.tp_group,
198 cp_group=self.cp_group,
199 dp_group=self.dp_group,
200 num_recomputations=self.num_recomputations,
201 )
203 def register(self, optimizer: Any) -> None:
204 """Register an optimizer post-hook to auto-trigger :meth:`on_step_end`.
206 On PyTorch 2.0+, uses ``optimizer.register_step_post_hook``.
207 The hook fires after each ``optimizer.step()`` call, so that
208 expert bias is updated once parameters have been updated. Repeated
209 calls replace the existing hook registration.
211 Args:
212 optimizer: The optimizer instance to register the hook on.
214 Raises:
215 RuntimeError: If the optimizer does not support step post-hooks.
216 """
217 if not hasattr(optimizer, "register_step_post_hook"):
218 raise RuntimeError(
219 "Optimizer does not support register_step_post_hook. "
220 "Requires PyTorch 2.0+. Call on_step_end() manually instead."
221 )
222 self.remove()
223 self._hook_handle = optimizer.register_step_post_hook(
224 lambda *_: self.on_step_end(),
225 )
227 def remove(self) -> None:
228 """Remove the optimizer post-hook if registered."""
229 if self._hook_handle is not None:
230 self._hook_handle.remove()
231 self._hook_handle = None