integrations.expert_parallel.plugin

integrations.expert_parallel.plugin

Expert-Parallel plugin for axolotl (DeepEP or torch all-to-all dispatch).

Classes

Name	Description
ExpertParallelPlugin	Plugin that swaps MoE dispatch/combine for expert-parallel token dispatch.

ExpertParallelPlugin

integrations.expert_parallel.plugin.ExpertParallelPlugin()

Plugin that swaps MoE dispatch/combine for expert-parallel token dispatch.

Methods

Name	Description
fully_shard_experts	Pre-wrap each Experts module (or the expert-LoRA ParamWrapper enclosing it) with
post_model_load	Propagate DDP-ignored params to the outermost model wrapper.

fully_shard_experts

integrations.expert_parallel.plugin.ExpertParallelPlugin.fully_shard_experts(
    model,
    dp_shard_mesh,
    fsdp2_kwargs,
)

Pre-wrap each Experts module (or the expert-LoRA ParamWrapper enclosing it) with FSDP on dp_shard_mesh: the non-ep data axes under composition, a per-rank size-1 mesh under pure EP.

Called from the patched fsdp2_prepare_model BEFORE the outer auto-wrap so experts become FSDPModules and the auto-wrap walker skips them. Inherits the outer wrap’s policy (mp, offload, reshard) so inner/outer collective dtypes line up; only mesh is overridden.

post_model_load

integrations.expert_parallel.plugin.ExpertParallelPlugin.post_model_load(
    cfg,
    model,
)

Propagate DDP-ignored params to the outermost model wrapper.

post_model_build set _ddp_params_and_buffers_to_ignore on the inner model. After PEFT wraps it (in PeftModel), DDP wraps PeftModel, but DDP looks for the attribute on the top-level module — which is now PeftModel, not our inner model. Mirror the list up.

Functions

Name	Description
expert_fsdp_mesh	The mesh the routed experts FSDP-shard on under EP composition, or None.
expert_shard_axis	The non-ep mesh axis the routed experts FSDP-shard on under EP composition, or None.
per_rank_expert_mesh	The mesh the routed experts FSDP-wrap on when EP has no shard axis (pure EP, or

expert_fsdp_mesh

integrations.expert_parallel.plugin.expert_fsdp_mesh(mesh)

The mesh the routed experts FSDP-shard on under EP composition, or None.

Every non-ep data axis (dp_shard and cp) is flattened into one shard dim: cp ranks hold different halves of the same sequences, so their expert grads must be reduced exactly like dp_shard’s. dp_replicate stays a separate outer dim (HSDP).

expert_shard_axis

integrations.expert_parallel.plugin.expert_shard_axis(mesh_dim_names)

The non-ep mesh axis the routed experts FSDP-shard on under EP composition, or None.

Prefers dp_shard (EP×dp_shard: experts shard on the data axis); falls back to cp (EP×cp, where the cp ranks of an ep-group hold the SAME experts since cp shards the sequence, not the experts, so FSDP-sharding them on cp keeps each rank from holding the full ep-group slice). Returns None for pure EP (no secondary axis) or when there is no ep axis to compose with — those paths don’t pre-wrap the experts here.

per_rank_expert_mesh

integrations.expert_parallel.plugin.per_rank_expert_mesh(mesh, device_type)

The mesh the routed experts FSDP-wrap on when EP has no shard axis (pure EP, or dp_replicate x ep): a fresh (dp_replicate?, ep, dp_shard=1) mesh whose size-1 dp_shard dim keeps each rank’s expert slice whole (the all-gather is a no-op) while making every param a DTensor; dp_replicate stays the outer dim so HSDP reduces the expert grads across replicas. Returns the (dp_replicate, dp_shard) or dp_shard sub-mesh to hand fully_shard_experts.