> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.moe.state_dict_utils

## Module Contents

### Functions

| Name                                                                                                                         | Description                                                             |
| ---------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`_get_expert_mesh_dim_index`](#nemo_automodel-components-moe-state_dict_utils-_get_expert_mesh_dim_index)                   | Find the device-mesh dimension that partitions a grouped expert tensor. |
| [`create_dtensor_from_local`](#nemo_automodel-components-moe-state_dict_utils-create_dtensor_from_local)                     | Create a DTensor from a local tensor for expert parallelism.            |
| [`get_expert_range_for_rank_from_mesh`](#nemo_automodel-components-moe-state_dict_utils-get_expert_range_for_rank_from_mesh) | Get the range of experts that should be loaded for the current rank.    |
| [`get_expert_slice_for_rank`](#nemo_automodel-components-moe-state_dict_utils-get_expert_slice_for_rank)                     | Get the slice of experts present on the current rank for a DTensor.     |
| [`get_submesh`](#nemo_automodel-components-moe-state_dict_utils-get_submesh)                                                 | Access a submesh by dim names from the given mesh.                      |
| [`is_dtensor`](#nemo_automodel-components-moe-state_dict_utils-is_dtensor)                                                   | Check if a tensor is a DTensor.                                         |
| [`should_load_expert_for_rank`](#nemo_automodel-components-moe-state_dict_utils-should_load_expert_for_rank)                 | Check if a specific expert should be loaded on the current rank.        |
| [`split_experts_weights_dtensor_aware`](#nemo_automodel-components-moe-state_dict_utils-split_experts_weights_dtensor_aware) | Split expert weights, handling both regular tensors and DTensors.       |
| [`validate_dtensor_expert_sharding`](#nemo_automodel-components-moe-state_dict_utils-validate_dtensor_expert_sharding)       | Validate that a DTensor is properly sharded for expert parallelism.     |

### API

```python
nemo_automodel.components.moe.state_dict_utils._get_expert_mesh_dim_index(
    dtensor: torch.distributed._tensor.DTensor
) -> int | None
```

Find the device-mesh dimension that partitions a grouped expert tensor.

**Parameters:**

**`dtensor`** `DTensor`

DTensor with global shape \[experts, ...] and arbitrary trailing dimensions. A named `ep` mesh
dimension owns the expert partition when present; otherwise a `Shard(0)` placement identifies an
EP-free FSDP partition of the experts axis.

---

**Returns:** `int | None`

Index of the expert-partitioning mesh dimension, or `None` when every rank retains all experts.

```python
nemo_automodel.components.moe.state_dict_utils.create_dtensor_from_local(
    local_tensor: torch.Tensor,
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh],
    rank: int | None = None
) -> torch.Tensor
```

Create a DTensor from a local tensor for expert parallelism.

**Parameters:**

**`local_tensor`** `torch.Tensor`

Local portion of the tensor on this rank

---

**`device_mesh`** `Optional[DeviceMesh]`

Device mesh for DTensor creation

---

**`rank`** `int | None` — default: None

Current rank (for device placement)

---

**Returns:** `torch.Tensor`

DTensor if device\_mesh is provided and DTensor is available, otherwise local\_tensor

```python
nemo_automodel.components.moe.state_dict_utils.get_expert_range_for_rank_from_mesh(
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh],
    n_experts: int
) -> tuple[int, int]
```

Get the range of experts that should be loaded for the current rank.

**Parameters:**

**`device_mesh`** `Optional[DeviceMesh]`

Device mesh for expert parallelism

---

**`n_experts`** `int`

Total number of experts

---

**Returns:** `tuple[int, int]`

Tuple of (start\_expert\_id, end\_expert\_id) for this rank

```python
nemo_automodel.components.moe.state_dict_utils.get_expert_slice_for_rank(
    experts_tensor: torch.Tensor,
    n_experts: int
) -> tuple[torch.Tensor, int, int]
```

Get the slice of experts present on the current rank for a DTensor.

For non-DTensors, returns the full tensor with start\_expert=0, end\_expert=n\_experts.
For DTensors sharded along the expert dimension (dim=0), returns only the local experts. The mesh dimension
may be the explicit `ep` dimension or an EP-free FSDP dimension such as `dp_shard_cp`.

**Parameters:**

**`experts_tensor`** `torch.Tensor`

Tensor with global shape \[experts, ...] and arbitrary trailing dimensions. For a DTensor,
the per-rank local shape is \[local\_experts, ...] when a mesh dimension uses `Shard(0)`; other
placements retain `experts` on the first local axis.

---

**`n_experts`** `int`

Total number of experts across all ranks.

---

**Returns:** `torch.Tensor`

Tuple containing the local tensor of shape \[local\_experts, ...], the inclusive global ID of its first

```python
nemo_automodel.components.moe.state_dict_utils.get_submesh(
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    dims: tuple[str, ...]
) -> torch.distributed.device_mesh.DeviceMesh
```

Access a submesh by dim names from the given mesh.

```python
nemo_automodel.components.moe.state_dict_utils.is_dtensor(
    tensor: torch.Tensor
) -> bool
```

Check if a tensor is a DTensor.

```python
nemo_automodel.components.moe.state_dict_utils.should_load_expert_for_rank(
    expert_id: int,
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh],
    n_experts: int
) -> bool
```

Check if a specific expert should be loaded on the current rank.

**Parameters:**

**`expert_id`** `int`

The expert ID to check

---

**`device_mesh`** `Optional[DeviceMesh]`

Device mesh for expert parallelism

---

**`n_experts`** `int`

Total number of experts

---

**Returns:** `bool`

True if this expert should be loaded on the current rank

```python
nemo_automodel.components.moe.state_dict_utils.split_experts_weights_dtensor_aware(
    weight: torch.Tensor,
    n_experts: int
) -> tuple[list[torch.Tensor], list[int]]
```

Split expert weights, handling both regular tensors and DTensors.

For DTensors sharded on the expert axis, only splits the experts present on the current rank. Other DTensor
placements retain all experts and are adjusted after removing the expert tensor dimension.

**Parameters:**

**`weight`** `torch.Tensor`

Tensor with global shape \[experts, ...] and arbitrary trailing dimensions. A `Shard(0)` DTensor
has local shape \[local\_experts, ...]; inner-axis shards retain shape \[experts, ...] locally.

---

**`n_experts`** `int`

Total number of experts across all ranks.

---

**Returns:** `list[torch.Tensor]`

Tuple containing per-expert tensors of shape \[...] and their global expert IDs. Each output aliases the

```python
nemo_automodel.components.moe.state_dict_utils.validate_dtensor_expert_sharding(
    tensor: torch.Tensor,
    expected_experts: int,
    tensor_name: str = 'tensor'
) -> bool
```

Validate that a DTensor is properly sharded for expert parallelism.

**Parameters:**

**`tensor`** `torch.Tensor`

Tensor to validate

---

**`expected_experts`** `int`

Expected total number of experts

---

**`tensor_name`** `str` — default: 'tensor'

Name for error messages

---

**Returns:** `bool`

True if valid, raises ValueError if invalid