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).
apply_distributed_support	Prepare hf_config so model_type can be loaded expert-parallel, from its model spec.
check_distributed_support	Refuse model_type's expert parallelism if its spec does; enforce its transformers floor.
expert_parallel_mesh	The mesh model's routed experts are sharded over, or None if its plan shards none.
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.

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)

check_distributed_support(model_type)

Refuse model_type’s expert parallelism if its spec does; enforce its transformers floor.

apply_distributed_support() checks this before it prepares a config; a config already prepared with its plan needs only the check. An EP degree left to auto picks none for a model this refuses (_spec_refuses_ep()), so only an EP degree asked for explicitly, or an expert-parallel plan handed in, meets these errors; both name the way around expert parallelism.

Parameters:

model_type (str)

Return type:

None

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)

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.