Triton RMSNorm example uses epsilon 1e-6 but definition requires 1e-5
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 296
- Forks
- 34
- PR merge metrics
- No merged PRs in 30d
Description
Summary
The Triton RMSNorm example does not implement its own problem definition.
Observed mismatch for x = BF16(1e-3) and weight = 1: expected 0.30078125 vs. candidate output 0.70703125.
definition.jsonfixesEPS = 1e-5in the reference and says “Epsilon is fixed at 1e-5.”kernel.pylaunches_rmsnorm_fwd_kernelwith1e-6.- The embedded
kernel.pycontent insolution_triton.jsonalso contains1e-6.
This is present on current main at commit a9fa0804c793d438e70850c33fe34426e66d53dd.
Minimal reproducer
Use a valid official-shape BF16 row with small finite values and unit weights:
import torch
x = torch.full((1, 4096), 1e-3, dtype=torch.bfloat16, device="cuda")
weight = torch.ones(4096, dtype=torch.bfloat16, device="cuda")
def rmsnorm(x, eps):
x32 = x.float()
return (
x32 * torch.rsqrt(x32.square().mean(-1, keepdim=True) + eps)
* weight.float()
).bfloat16()
print(x[0, 0].item())
print(rmsnorm(x, 1e-5)[0, 0].item())
print(rmsnorm(x, 1e-6)[0, 0].item())
Input: 0.00099945068359375
Expected (reference, epsilon = 1e-5): 0.30078125
Observed (Triton candidate, epsilon = 1e-6): 0.70703125
So the concise comparison is: expected 0.30078125 vs. observed 0.70703125.
A randomized (randn * 1e-3).bfloat16() input with batch size 7 produced 26,631 mismatching elements out of 28,672.
This is not reduction-order noise: for a constant row, the output is analytically x / sqrt(x² + eps), and the observed ratio agrees with changing epsilon from 1e-5 to 1e-6.
Expected behavior
The bundled example solution should use the epsilon required by rmsnorm_h4096 for all valid BF16 inputs.
Suggested fix
- Change the launch argument in
examples/triton/rmsnorm/kernel.pyfrom1e-6to1e-5. - Make the same change in the embedded source inside
examples/triton/rmsnorm/solution_triton.json. - Add a small-magnitude BF16 correctness case; unit-scale inputs make epsilon comparatively negligible and can hide this mismatch.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with examples/triton/rmsnorm/definition.json and compare its fixed epsilon with the launch in examples/triton/rmsnorm/kernel.py. Update the launch and the embedded source in examples/triton/rmsnorm/solution_triton.json, then add a small-magnitude BF16 correctness case; done means the bundled example matches the reference for valid BF16 inputs.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, testing-qa
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100