Inconsistent DTensorPolicyWorker logits with TP>1 and bf16
Open
bug
t-pytdensor
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the bug**
When using TP>1 and bf16 policy.precision, we observed inconsistent logits computation during the forward pass of DTensorPolicyWorker, impacting RM training.
- Identical rows yield different outputs depending on their position in the microbatch
- Consecutive forward passes produce different outputs
Using fp32 for policy.precision seems to mitigate this issue.
Contributor guide
Assessment
This issue has not been assessed yet.