`core.transformer.moe.router_diagnostics`#

Compact observations used to diagnose MoE router load balancing.

Module Contents#

Classes#

RouterDiagnosticChannel

Channels in the compact per-sequence router diagnostic tensor.

Functions#

build_router_diagnostics

Build compact score, load, and bias summaries for each local sequence.

Data#

ROUTER_DIAGNOSTIC_CHANNEL_COUNT

API#

class core.transformer.moe.router_diagnostics.RouterDiagnosticChannel#

Bases: enum.IntEnum

Channels in the compact per-sequence router diagnostic tensor.

Initialization

Initialize self. See help(type(self)) for accurate signature.

MEAN_SCORE#: 0

AUX_LOAD#: 1

ACTUAL_LOAD#: 2

EXPERT_BIAS#: 3

AUX_ACTUAL_OVERLAP#: 4

VALID_TOKEN_COUNT#: 5

TOPK_BOUNDARY_RELATIVE_MARGIN#: 6

core.transformer.moe.router_diagnostics.ROUTER_DIAGNOSTIC_CHANNEL_COUNT#: ‘len(…)’

core.transformer.moe.router_diagnostics.build_router_diagnostics( scores_for_aux_loss: torch.Tensor, routing_map_for_aux_loss: torch.Tensor, actual_routing_map: torch.Tensor, expert_bias: torch.Tensor | None, seq_length: int, batch_size: int, padding_mask: torch.Tensor | None = None, ) → torch.Tensor#

Build compact score, load, and bias summaries for each local sequence.

The returned tensor has shape [batch_size, channels, num_experts]. Distribution-valued channels are normalized over experts. Scalar channels use element zero and leave the remaining expert elements as zero.

Parameters:

scores_for_aux_loss – Normalized all-expert scores with shape [tokens, num_experts].
routing_map_for_aux_loss – Unbiased top-k assignments with the same shape.
actual_routing_map – Assignments used for token dispatch with the same shape.
expert_bias – Additive expert-selection correction, or None when disabled. Positive values favor an expert. Ordinary bias routing uses score units; quantile balancing uses -qb_beta in logit units.
seq_length – Local sequence length before the token dimension was flattened.
batch_size – Local micro-batch size.
padding_mask – Flattened mask where True marks padding.

Returns:

Float32 diagnostic tensor with one compact observation per sequence.

core.transformer.moe.router_diagnostics#