moe#

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

Finding a model’s routed experts, choosing the EP degree they are sharded over, and 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’ quantizable _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 all-to-all (NCCL).

choose_ep_dp_degrees

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

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).

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_global_experts, num_local_experts, expert_fn)#

Route one MoE layer’s expert stage by all-to-all (NCCL).

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]

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 the stock megamoe_router / megamoe_experts (which were added for DeepGEMM Mega-MoE in transformers 5.10.4), reimplemented here – ~5 lines each – so the a2a path OWNS its EP styles and does NOT depend on megamoe (i.e. it stays portable below 5.10.4). 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 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).