nemo_automodel.components.moe.tp_plan_validation

Routed-expert ownership validation for custom-MoE tensor-parallel plans.

Both the dedicated MoE parallelizer (components.moe.parallelizer) and the generic dense parallelizer (components.distributed.parallelizer, for the EP=1 path) must reject TP plans that shard routed-expert or router parameters. The validator lives in this dependency-free module so the two parallelizers can share it without importing each other.

Module Contents

Functions

Name	Description
`_is_unsafe_moe_tp_module_path`	-
`_validate_moe_tp_plan`	Validate that a TP plan cannot shard routed-expert parameters.

API

nemo_automodel.components.moe.tp_plan_validation._is_unsafe_moe_tp_module_path(
    module_path: str
) -> bool

nemo_automodel.components.moe.tp_plan_validation._validate_moe_tp_plan(
    plan: object,
    model: torch.nn.Module | None = None
) -> dict[str, torch.distributed.tensor.parallel.ParallelStyle]

Validate that a TP plan cannot shard routed-expert parameters.

Routed experts are owned by expert parallelism and, when configured, the independent ep_shard FSDP mesh. Applying tensor parallelism to the same parameters would compose unrelated DTensor placements and either corrupt ownership or fail during the first grouped GEMM. Shared experts are regular dense MLPs and are intentionally allowed.