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 all-to-all (NCCL).
apply_distributed_support	Prepare hf_config so model_type can be loaded expert-parallel, from its model spec.
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_global_experts, num_local_experts, expert_fn)

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

apply_distributed_support(model_type, hf_config, *, ckpt_path=None, trust_remote_code=False)

Prepare hf_config so model_type can be loaded expert-parallel, from its model spec.

Reads the model’s EPSpec and applies it: enforce the transformers floor this model’s expert parallelism needs, then attach the base_model_ep_plan it requires to whichever config carries its MoE – the root for a plain MoE, the nested text/LLM config for a VLM whose experts live in its language sub-model, since the plan’s globs are written against that module tree and a plan on the wrong config matches nothing at all rather than failing.

An existing plan always wins: stock transformers ships one for some models, and overwriting it would substitute our idea of the layout for the model’s own.

Returns the config to load, unchanged for a model type with no spec or no EP section. The generic all-to-all transport is not wired here – it self-registers in load_and_prepare_ep(), so nothing about it is per-model.

Parameters:

• model_type (str)

• ckpt_path (str | None)

• trust_remote_code (bool)

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

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