BUG: MoE model fails with CPU offload enabled (gradient/parameter device mismatch)
- 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
Assessment
This issue has not been assessed yet.