NVIDIA-NeMo / NVIDIA-NeMo/RL

feat: support context parallelism for reward model training on DTensor backend

Open
#2,482 0 comments 0 reactions 1 assignee Claimed by @ashors1 View on GitHub
enhancement from-slack
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

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.