kempnerforge.distributed.expert_parallel

Expert Parallelism: partition MoE experts across an EP process group.

Each EP rank holds num_experts // ep_size experts. Tokens are shuffled between ranks via all-to-all so every token reaches its assigned expert, then results are returned to the originating rank.

When ep=1 (default), this module is a no-op and the model runs with all experts replicated on every rank (the pre-EP behavior).

Functions

apply_expert_parallel(model, device_mesh)

Partition MoE experts across the EP dimension of the DeviceMesh.

ep_dispatch_and_compute(x, weights, indices, ...)

All-to-all dispatch, local expert compute, all-to-all combine.

expert_parallel_params(model)

The routed-expert parameters that EP splits across ranks, and the EP group.
kempnerforge.distributed.expert_parallel.ep_dispatch_and_compute(x, weights, indices, moe, ep_group, local_expert_start, num_local_experts, ep_world_size, gradient_scale=False)[source]

All-to-all dispatch, local expert compute, all-to-all combine.

When the EP ranks hold the same batch, as they do in training, the backward pass gives the local experts the mean of the ranks’ gradients, which is one process’s, and gives x and weights the same gradients as without EP.

Parameters:

  • x (torch.Tensor) – (num_tokens, dim) flattened token representations.

  • weights (torch.Tensor) – (num_tokens, top_k) routing weights.

  • indices (torch.Tensor) – (num_tokens, top_k) global expert indices.

  • moe (MoEMLP) – The MoEMLP module. Holds either an nn.ModuleList of local experts (unpacked) or packed expert tensors already sliced to the local range.

  • ep_group (torch.distributed.ProcessGroup) – EP process group.

  • local_expert_start (int) – First global expert index on this rank.

  • num_local_experts (int) – Number of experts on this rank.

  • ep_world_size (int) – Number of ranks in the EP group.

  • gradient_scale (bool) – When moe is in training mode, scale each local expert’s output by the average load over all experts divided by its own load, as MoEMLP._local_forward does.

Returns:

(num_tokens, dim) weighted combination of expert outputs.

Return type:

output

kempnerforge.distributed.expert_parallel.apply_expert_parallel(model, device_mesh)[source]

Partition MoE experts across the EP dimension of the DeviceMesh.

For each MoEMLP in the model:

  • Prunes experts to the local subset for this EP rank

  • Stores EP metadata (group, mesh, world_size, local expert range)

The EP mesh lets a checkpoint place each rank’s experts among all of them (see kempnerforge.checkpoint.expert_layout).

Must be called AFTER tensor parallelism and BEFORE FSDP2.

No-op when ep is not in the mesh or has size 1.

Parameters:
Return type:

None

kempnerforge.distributed.expert_parallel.expert_parallel_params(model)[source]

The routed-expert parameters that EP splits across ranks, and the EP group.

Each EP rank holds its own slice of these and a copy of every other parameter. Returns ([], None) when no MoEMLP has ep_world_size > 1.

Raises:

ValueError – Two MoEMLP layers record different EP groups, so no single group covers every routed expert.

Parameters:

model (torch.nn.Module)

Return type:

tuple[list[torch.nn.Parameter], dist.ProcessGroup | None]