moe#

HuggingFace-coupled MoE expert-parallel (EP) helpers.

The DP=EP all-to-all expert transport (torch_a2a).

The all-to-all transport is plain torch.distributed.all_to_all_single (NCCL) wrapped around transformers’ _grouped_linear. register_torch_a2a() registers it into transformers’ ExpertsInterface lazily – never at import time, so importing this module has no transformers-ExpertsInterface side effect.

Functions

a2a_moe

Route one MoE layer's expert stage by plain torch.distributed.all_to_all_single (NCCL).

register_torch_a2a

Register everything the torch_a2a path needs into stock transformers (idempotent).

torch_a2a_experts_forward

All-to-all expert-parallel forward for PTQ calibration (DP=EP).
a2a_moe(group, hidden_states, topk_idx, topk_weights, num_local_experts, expert_fn)#

Route one MoE layer’s expert stage by plain torch.distributed.all_to_all_single (NCCL).

Works cross-node at any rank count, with explicit token permute/unpermute.

Flow: expand (token, expert) pairs -> route each to dest_rank = global_expert // num_local -> all_to_all the tokens (+ local-expert id) -> local grouped_mm -> all_to_all back -> weight by the routing weights and sum each token’s top_k results.

Tokens travel and the result returns in hidden_states’ dtype – the model’s, so its residual stream keeps one dtype. The routing weights never travel: the rank that holds a token weights its experts’ results in the weights’ own dtype and sums them as transformers’ grouped_mm forward does.

register_torch_a2a()#

Register everything the torch_a2a path needs into stock transformers (idempotent).

Done lazily from load_and_prepare_ep, not at import time, so importing this module has no transformers side effect:

  • the torch_a2a experts forward into ExpertsInterface (ALL_EXPERTS_FUNCTIONS);

  • two EP parallel styles into ALL_PARALLEL_STYLES:

    • torch_a2a_router – emit global, unmasked expert ids (skip RouterParallel’s per-EP-rank masking) so the a2a dispatch can route each (token, expert) to its owning rank;

    • torch_a2a_experts – inject the EP process_group into the experts forward (4th arg) and skip the all-reduce (the a2a combine does the reduction).

These two styles are functionally transformers’ megamoe_router / megamoe_experts (added in 5.10.4 for DeepGEMM Mega-MoE), reimplemented here so the transport works below 5.10.4 too. The base classes (RouterParallel, TensorParallelLayer) are stock wherever the EP infra exists (transformers >= 5.6, the framework floor enforced in load_and_prepare_ep), so this works on any such version; grouped_gemm (expert weight sharding) is stock too.

torch_a2a_experts_forward(self, hidden_states, top_k_index, top_k_weights, process_group=None)#

All-to-all expert-parallel forward for PTQ calibration (DP=EP).

Registered as the “torch_a2a” experts implementation on transformers’ ExpertsInterface.

Dense+router run data-parallel; experts are EP-sharded; tokens are routed to the expert-owning rank by the NCCL all-to-all above. Local experts run through transformers’ _grouped_linear on the tokens dispatched to them from every rank. process_group (the EP group) is injected by the torch_a2a_experts TP style’s _prepare_input_fn; the router must emit GLOBAL ids unmasked (torch_a2a_router).

self.num_experts is this rank’s share of the experts: the expert-parallel loader loads them through transformers’ expert-parallel sharding (the grouped_gemm / packed styles), which slices the experts’ weights along dim 0 and sets num_experts to the per-rank count.