`core.models.hybrid.shortcut_block`#

Implement shortcut-connected MoE blocks based on the ScMoE architecture introduced in Shortcut-connected Expert Parallelism for Accelerating Mixture of Experts <https://arxiv.org/abs/2404.05019>_.

ScMoE routes the preceding layer’s representation through sparse experts while the paired layer computes attention and the shared expert, breaking the usual sequential dependency so expert-parallel All-to-All communication can overlap with computation. Our variant adds two LayerNorms to the paper’s design: shortcut_pre_mlp_layernorm normalizes the shortcut representation before routing and expert dispatch, while shortcut_post_norm normalizes the merged routed- and shared-expert result before the MoE residual/BDA step. Also supports selective activation recomputation and activation offloading on shortcut_pre_mlp_layernorm and shortcut_post_norm, respectively, to reduce the additional device-memory pressure from enabling Shortcut-MoE.

Module Contents#

Classes#

ShortcutMoEBlock

Own and execute one attention-layer/shortcut-MoE pair.

Functions#

`_get_offloading_interface`	Get the fine-grained activation offloading interface lazily.
`group_layers_into_shortcut_blocks`	Group physical layers into their registered shortcut-block hierarchy.

Data#

SUPPORTED_SHORTCUT_PREDECESSORS

API#

core.models.hybrid.shortcut_block.SUPPORTED_SHORTCUT_PREDECESSORS#: ‘frozenset(…)’

core.models.hybrid.shortcut_block._get_offloading_interface()#: Get the fine-grained activation offloading interface lazily.

core.models.hybrid.shortcut_block.group_layers_into_shortcut_blocks( layers: torch.nn.ModuleList, layer_type_list: Sequence[str], config: megatron.core.transformer.transformer_config.TransformerConfig, pp_layer_offset: int = 0, ) → torch.nn.ModuleList#

Group physical layers into their registered shortcut-block hierarchy.

Grouping updates the layer names through the returned ModuleList hierarchy. Layers not followed by an MoE remain direct children of the returned ModuleList.

Parameters:

layers – Physical layers in execution order.
layer_type_list – Physical layer symbols in execution order.
config – Transformer configuration controlling shortcut scheduling.
pp_layer_offset – Global offset of this pipeline stage’s first physical layer.

Returns:

The registered logical layers.

Raises:

ValueError – If an MoE has an unsupported predecessor or its shortcut pair crosses a pipeline-stage boundary.

class core.models.hybrid.shortcut_block.ShortcutMoEBlock( attn_layer, moe_layer, overlap_a2a: bool, attn_local_idx: int | None = None, moe_local_idx: int | None = None, )#

Bases: megatron.core.transformer.module.MegatronModule

Own and execute one attention-layer/shortcut-MoE pair.

Initialization

_parallel_stream: torch.cuda.Stream | None#: None

classmethod _get_a2a_overlap_stream() → torch.cuda.Stream#: Return the process-wide high-priority shortcut stream.

_moe_router_preprocess( shortcut_hidden, padding_mask=None, packed_seq_params=None, )#: Run shortcut normalization, routing, and dispatch preprocessing.

_moe_shared_experts( hidden_states, padding_mask=None, packed_seq_params=None, )#

Run the paired MoE layer’s pre-MLP norm and shared experts.

Returns:: (shared_expert_output, moe_unflatten_mbs, residual, mlp_state).

_postprocess( residual, combined_output, shared_expert_output, packed_seq_params=None, moe_unflatten_mbs=None, mlp_state=(), )#: Join routed/shared output, apply shortcut post-norm, and finish residual/BDA.

_launch_dispatch( hidden_states: torch.Tensor, probs: torch.Tensor, async_op: bool = False, ) → tuple[torch.Tensor, torch.Tensor]#: Launch dispatch on the current stream or the shortcut side stream.

_wait_dispatch( dispatched_input: torch.Tensor, dispatched_probs: torch.Tensor, ) → tuple[torch.Tensor, torch.Tensor]#: Wait for dispatch and return its outputs on the main stream.

_launch_combine( output: torch.Tensor, async_op: bool = False, ) → torch.Tensor#: Launch combine on the current stream or the shortcut side stream.

_wait_combine(combined_output: torch.Tensor) → torch.Tensor#: Wait for the asynchronous combine and return its output on the main stream.

forward( hidden_states, attention_mask, inference_context, rotary_pos_emb, sequence_len_offset, packed_seq_params, padding_mask, quant_context_factory, cp_layout_state=None, packed_sequence_cp_metadata=None, )#: Run the eager schedule with each physical layer’s quantization context.

core.models.hybrid.shortcut_block#