NVIDIA-NeMo / NVIDIA-NeMo/RL

BUG: DTensor backend fails with expert_parallel_size=1 (moe_mesh is None)

Open
#2,028 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

# DTensor backend fails with expert_parallel_size=1 (moe_mesh is None)

## Summary

When running GRPO with the DTensor v2 backend and `expert_parallel_size=1`, the process crashes during `setup()` with:

```text
TypeError: 'NoneType' object is not subscriptable
```

The error occurs in the MoE parallelizer when it tries to use `moe_mesh[ep_shard_axis_names]` while `moe_mesh` is `None` (which is expected when expert parallelism size is 1).

## Environment

- **Script**:
```bash
NRL_FORCE_REBUILD_VENVS=true uv run python examples/run_grpo.py \
--config=examples/configs/grpo_math_1B.yaml \
policy.model_name="Qwen/Qwen3-30B-A3B-Instruct-2507" \
grpo.max_num_steps=3 \
logger.wandb_enabled=false \
logger.wandb.project=ruit_personal_debug \
policy.dtensor_cfg._v2=true \
++policy.dtensor_cfg.expert_parallel_size=1 \
++policy.dtensor_cfg.automodel_kwargs.use_liger_kernel=false \
++policy.dtensor_cfg.automodel_kwargs.backend._target_="nemo_automodel.components.moe.utils.BackendConfig" \
++policy.dtensor_cfg.automodel_kwargs.backend.attn=te \
++policy.dtensor_cfg.automodel_kwargs.backend.linear=te \
++policy.dtensor_cfg.automodel_kwargs.backend.rms_norm=te \
++policy.dtensor_cfg.automodel_kwargs.backend.enable_deepep=true \
++policy.dtensor_cfg.automodel_kwargs.backend.fake_balanced_gate=false \
++policy.dtensor_cfg.automodel_kwargs.backend.enable_hf_state_dict_adapter=true \
++policy.dtensor_cfg.automodel_kwargs.backend.enable_fsdp_optimizations=false \
++policy.dtensor_cfg.automodel_kwargs.backend.gate_precision="float64" \
policy.generation.vllm_cfg.gpu_memory_utilization=0.8 \
checkpointing.save_period=1 \
checkpointing.enabled=false \
checkpointing.model_save_format="safetensors" \
checkpointing.save_consolidated=false \
cluster.gpus_per_node=8 \
cluster.num_nodes=2
```
- **Config**: e.g. `examples/configs/grpo_math_1B.yaml` with `++policy.dtensor_cfg.expert_parallel_size=1`
- **Model**: MoE model (e.g. Qwen3-30B-A3B)

## Traceback

```text
Traceback (most recent call last):
File ".../examples/run_grpo.py", line 178, in
main()
File ".../examples/run_grpo.py", line 110, in main
) = setup(config, tokenizer, dataset, val_dataset)
File ".../nemo_rl/algorithms/grpo.py", line 681, in setup
policy.print_node_ip_and_gpu_id()
File ".../nemo_rl/models/policy/lm_policy.py", line 928, in print_node_ip_and_gpu_id
results = ray.get(...)
...
ray.exceptions.ActorDiedError: The actor died because of an error raised in its creation task, ray::lm_policy-0-0:DTensorPolicyWorkerV2.__init__()
...
File ".../nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py", line 278, in __init__
model_and_optimizer_state = setup_model_and_optimizer(...)
File ".../nemo_rl/models/automodel/setup.py", line 469, in setup_model_and_optimizer
moe_parallelize_model(
File ".../3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/moe/parallelizer.py", line 279, in parallelize_model
ep_shard_mesh = moe_mesh[ep_shard_axis_names]
~~~~~~~~^^^^^^^^^^^^^^^^^^^^^
TypeError: 'NoneType' object is not subscriptable
```

## Root Cause

- When `expert_parallel_size=1`, the distributed manager sets `moe_mesh = None` (see e.g. `nemo_automodel/components/distributed/fsdp2.py`: `self.moe_mesh = self._get_moe_mesh() if self.ep_size > 1 else None`).
- `nemo_rl/models/automodel/setup.py` always passes `ep_shard_axis_names=("ep_shard",)` into `moe_parallelize_model()`.
- In `nemo_automodel/components/moe/parallelizer.py`, `ep_enabled` correctly guards on `moe_mesh is not None` (line 266), but the `ep_shard_mesh` assignment (lines 277–279) only checks `ep_shard_axis_names is not None` and then does `ep_shard_mesh = moe_mesh[ep_shard_axis_names]`, so when `moe_mesh` is `None` we get the `TypeError`.

## Expected Behavior

With `expert_parallel_size=1`, MoE parallelization should run without using a MoE mesh (no EP, no ep_shard). The parallelizer should treat `ep_shard_axis_names` as unused when `moe_mesh` is `None` and set `ep_shard_mesh = None` instead of subscripting `moe_mesh`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.