NVIDIA-NeMo / NVIDIA-NeMo/RL

BUG: MoE model fails with CPU offload enabled (gradient/parameter device mismatch)

Open
#2,030 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

# MoE model fails with CPU offload enabled (gradient/parameter device mismatch)

## Summary

When running DPO (or other training) with an MoE model and the DTensor v2 backend with `policy.dtensor_cfg.cpu_offload=true`, the first training step fails during `loss.backward()` with:

```text
RuntimeError: attempting to assign a gradient with device type 'cuda' to a tensor with device type 'cpu'. Please ensure that the gradient and the tensor are on the same device
```

The failure occurs inside PyTorch FSDP’s `post_backward` when assigning the reduced gradient to `sharded_param.grad`: the gradient is on CUDA while the parameter is on CPU due to CPU offload.

## Environment

- **Script**: `test_assets/run_dpo_nanov3.sh`
- **Config**: DPO with DTensor v2, MoE model (e.g. Qwen3-30B-A3B), `policy.dtensor_cfg.cpu_offload=true`
- **Model**: MoE (e.g. `Qwen/Qwen3-30B-A3B-Instruct-2507`)

## Steps to Reproduce

1. Use DTensor v2 with an MoE model and enable CPU offload:
- `policy.dtensor_cfg._v2=true`
- `policy.dtensor_cfg.cpu_offload=true`
- MoE model, e.g. `policy.model_name=Qwen/Qwen3-30B-A3B-Instruct-2507`
2. Run DPO (or any training that uses the same path), e.g.:

```bash
uv run examples/run_dpo.py \
--config examples/configs/dpo.yaml \
policy.model_name=Qwen/Qwen3-30B-A3B-Instruct-2507 \
policy.dtensor_cfg.enabled=true \
++policy.dtensor_cfg._v2=true \
++policy.dtensor_cfg.expert_parallel_size=8 \
policy.dtensor_cfg.cpu_offload=true \
...
```

3. Training fails at Step 1 during the backward pass.

## Traceback

```text
========================= Step 1/20 =========================
▶ Taking a training step...
Traceback (most recent call last):
File ".../examples/run_dpo.py", line 104, in
main()
...
File ".../nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py", line 428, in train
mb_results = automodel_forward_backward(...)
File ".../nemo_rl/models/automodel/train.py", line 443, in automodel_forward_backward
loss.backward()
...
File ".../torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py", line 544, in post_backward
) = foreach_reduce(...)
File ".../torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py", line 614, in foreach_reduce
fsdp_param.sharded_param.grad = new_sharded_dtensor_grad
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: attempting to assign a gradient with device type 'cuda' to a tensor with device type 'cpu'. Please ensure that the gradient and the tensor are on the same device
```

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.