Unused context_parallel_group parameter in dtensor_policy_worker.py despite being required in loss functions
Open
bug
t-pytdensor
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
All loss functions in `nemo_rl/algorithms/loss_function.py` contain a `context_parallel_group` parameter that is passed to them.
However, in the training loop `nemo_rl/models/policy/dtensor_policy_worker.py`(lines 790-795), the `context_parallel_group` is not being passed. It appears that cp information is then retrieved again through the `get_logprobs_from_vocab_parallel_logits` within the loss functions, is that a design choice?
Contributor guide
Assessment
This issue has not been assessed yet.