moe

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

Finding a model’s routed experts, choosing the EP degree they are sharded over, the DP=EP all-to-all expert transport (torch_a2a). The expert-parallel model loader is load_and_prepare_ep().

The all-to-all transport is plain torch.distributed.all_to_all_single (NCCL) wrapped around transformers’ quantizable _grouped_linear. register_torch_a2a() registers it into transformers’ ExpertsInterface (lazily, from load_and_prepare_ep’s a2a path – 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).
choose_ep_dp_degrees	Pick (expert_parallel, expert_data_parallel) degrees with EP * EDP == world_size.
expert_parallel_mesh	The mesh model's routed experts are sharded over, or None if its plan shards none.
num_experts_from_config	Return the MoE expert count from an HF config (0 if the model is not MoE).
register_torch_a2a	Register everything the torch_a2a path needs into stock transformers (idempotent).
require_transformers	Fail loudly if the installed transformers is older than min_version (the floor for feature).
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 (quantizable) -> 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.

choose_ep_dp_degrees(world_size, num_experts)

Pick (expert_parallel, expert_data_parallel) degrees with EP * EDP == world_size.

EP must divide both num_experts (whole experts per rank) and world_size (a valid mesh factorization); we take the largest such EP, and each expert shard is then replicated over the remaining EDP = world_size // EP ranks. Falls back to (1, world_size) when the model is not MoE or no valid EP degree exists. Heuristic – callers may override.

Parameters:

• world_size (int)

• num_experts (int)

Return type:

tuple[int, int]

expert_parallel_mesh(model)

The mesh model’s routed experts are sharded over, or None if its plan shards none.

load_and_prepare_ep() loads the model over its expert-parallel submesh, which transformers keeps on the model as _device_mesh. Calibration takes the mesh to re-install the all-to-all forwards (distributed_forward_loop()).

Parameters:

model (Module)

num_experts_from_config(hf_config)

Return the MoE expert count from an HF config (0 if the model is not MoE).

Recurses into .text_config / .llm_config so VLMs whose nested text submodel is the MoE (Kimi-K2.7, Llama-4 multimodal, NemotronH-Omni, …) report the LM’s expert count.

Return type:

int

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.

require_transformers(min_version, feature)

Fail loudly if the installed transformers is older than min_version (the floor for feature).

Callers enforce the version where the needed APIs first exist, so a too-old transformers errors clearly instead of silently mis-behaving (e.g. transformers 5.5 silently shards nothing for EP). See each call site for the empirically-determined floor.

Parameters:

• min_version (str)

• feature (str)

Return type:

None

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 the ModelOpt-quantizable _grouped_linear so per-expert weight + shared input quantizers attach over the dispatched (global) token pool. 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.