Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / moe_utils.py: 100%
38 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 )
96class MoEMonitorCallback:
97 """Callback that automates expert bias updates across all MoE layers.
99 Provides a three-layer invocation architecture:
101 - **Layer 0**: :func:`sync_and_update_expert_bias` — per-module sync + update
102 - **Layer 1**: :meth:`on_step_end` — iterates all MoE layers, calling Layer 0
103 - **Layer 2**: :meth:`register` — registers optimizer post-hook for auto-trigger
105 Each layer can be used independently; Layer 0 is sufficient for basic usage,
106 while Layers 1–2 add convenience for production training loops.
108 Args:
109 model: The model containing MoE layers to monitor.
110 lr: Step size for expert bias updates. Defaults to ``1e-3``.
111 tp_group: Tensor parallel process group. Only needed for TP+SP
112 (sequence-sharded activations across TP ranks). Standard TP
113 does not need this — pass ``None``.
114 cp_group: Context parallel process group. Sync across sequence shards.
115 dp_group: Data parallel process group. Sync across different samples.
116 num_recomputations: Number of forward executions per optimizer step.
117 Default ``1``. Set to ``2`` when activation checkpoint is enabled.
119 Example:
120 >>> # Layer 1: manual trigger in training loop
121 >>> callback = MoEMonitorCallback(
122 ... model, lr=1e-3, dp_group=dp_group_info,
123 ... )
124 >>> for batch in dataloader:
125 ... loss = model(batch)
126 ... loss.backward()
127 ... optimizer.step()
128 ... callback.on_step_end() # auto-updates all MoE layers
129 ... optimizer.zero_grad()
130 >>>
131 >>> # Layer 2: auto-trigger via optimizer post-hook
132 >>> callback = MoEMonitorCallback(model, lr=1e-3, dp_group=dp_group_info)
133 >>> callback.register(optimizer)
134 >>> # optimizer.step() now auto-triggers callback.on_step_end()
136 Note:
137 Sync dimensions use ``dp_group``/``tp_group``/``cp_group`` (not ``ep_group``).
138 EP does not need synchronization because all EP ranks have identical
139 routing data (``tokens_per_expert`` shape = ``[total_experts]``).
140 """
142 def __init__(
143 self,
144 model: "nn.Module",
145 lr: float = 1e-3,
146 tp_group: Optional[Any] = None,
147 cp_group: Optional[Any] = None,
148 dp_group: Optional[Any] = None,
149 num_recomputations: int = 1,
150 ) -> None:
151 """Initialize MoEMonitorCallback."""
152 self.model = model
153 self.lr = lr
154 self.tp_group = tp_group
155 self.cp_group = cp_group
156 self.dp_group = dp_group
157 self.num_recomputations = num_recomputations
158 self._hook_handle = None
160 @property
161 def last_mean_aux_loss(self) -> Optional[float]:
162 """Read the mean raw router loss from the latest local microbatch.
164 Read only when logging: scalar conversion may synchronize the device.
165 Auxiliary loss is independent of whether expert-bias updates are enabled.
166 """
167 aux_losses = [
168 module.last_aux_loss.item()
169 for module in self.model.modules()
170 if getattr(module, "last_aux_loss", None) is not None
171 ]
172 return sum(aux_losses) / len(aux_losses) if aux_losses else None
174 def on_step_end(self) -> None:
175 """Update expert bias for all MoE layers in the model.
177 Iterates over all :class:`MoE` sub-modules with
178 ``enable_expert_bias=True`` and calls
179 :func:`sync_and_update_expert_bias` for each one.
181 Auxiliary loss is read separately through ``last_mean_aux_loss`` so
182 bias updates do not synchronize device scalars for logging each step.
183 """
184 for moe in self.model.modules():
185 if not getattr(moe, "enable_expert_bias", False):
186 continue
187 sync_and_update_expert_bias(
188 moe,
189 lr=self.lr,
190 tp_group=self.tp_group,
191 cp_group=self.cp_group,
192 dp_group=self.dp_group,
193 num_recomputations=self.num_recomputations,
194 )
196 def register(self, optimizer: Any) -> None:
197 """Register an optimizer post-hook to auto-trigger :meth:`on_step_end`.
199 On PyTorch 2.0+, uses ``optimizer.register_step_post_hook``.
200 The hook fires after each ``optimizer.step()`` call, so that
201 expert bias is updated once parameters have been updated. Repeated
202 calls replace the existing hook registration.
204 Args:
205 optimizer: The optimizer instance to register the hook on.
207 Raises:
208 RuntimeError: If the optimizer does not support step post-hooks.
209 """
210 if not hasattr(optimizer, "register_step_post_hook"):
211 raise RuntimeError(
212 "Optimizer does not support register_step_post_hook. "
213 "Requires PyTorch 2.0+. Call on_step_end() manually instead."
214 )
215 self.remove()
216 self._hook_handle = optimizer.register_step_post_hook(
217 lambda *_: self.on_step_end(),
218 )
220 def remove(self) -> None:
221 """Remove the optimizer post-hook if registered."""
222 if self._hook_handle is not None:
223 self._hook_handle.remove()
224 self._hook_handle = None