feat: support context parallelism for reward model training on DTensor backend
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
Context parallelism (CP) is not currently supported for reward model (RM) training when using the DTensor backend. Enabling `context_parallel_size > 1` in the DTensor config results in errors, and there is no validation to surface this limitation clearly to users.
## Details
When running reward model training with context parallelism enabled on the DTensor backend, two distinct failures occur depending on the parallelism configuration:
1. **CP + TP**: Setting `context_parallel_size=2` with `tensor_parallel_size=2` fails with:
```
ValueError: Unknown parallel style: local_rowwise
```
2. **CP only**: Setting `context_parallel_size=2` without additional TP fails with:
```
NotImplementedError: Operator aten.log_sigmoid_forward.default does not have a sharding strategy registered
```
The `log_sigmoid` operator used in the RM loss does not have a DTensor sharding strategy registered for CP meshes. A naive workaround of redistributing the rewards tensor to `Replicate` placement produces numerically incorrect results (e.g., accuracy > 1.0), suggesting the fix requires deeper integration with the CP mesh throughout the RM forward pass and loss computation.
## Reproduction
```bash
# Works (no CP):
uv run examples/run_rm.py
# Fails — CP + TP:
uv run examples/run_rm.py \
policy.dtensor_cfg.context_parallel_size=2 \
cluster.gpus_per_node=4 \
policy.make_sequence_length_divisible_by=2 \
policy.dtensor_cfg.tensor_parallel_size=2
# Error: ValueError: Unknown parallel style: local_rowwise
# Fails — CP only:
uv run examples/run_rm.py \
policy.dtensor_cfg.context_parallel_size=2 \
cluster.gpus_per_node=2 \
policy.make_sequence_length_divisible_by=2
# Error: NotImplementedError: Operator aten.log_sigmoid_forward.default does not have a sharding strategy registered
```
## Suggested Action
At minimum, raise a clear error when `context_parallel_size > 1` is set for RM training on the DTensor backend, so users are not silently hitting unsupported configurations. Full CP support for RM would be the ideal resolution.
Contributor guide
Assessment
This issue has not been assessed yet.