Norm Samples updates for B,S,H style tensor inputs#432
Conversation
📝 WalkthroughWalkthroughLayerNorm and RMSNorm tensor layouts were changed from flattened 4D forms to explicit sequence-oriented shapes. Backward property inference now copies dimensions and strides from input and scale tensors, with corresponding C++ and Python validation updates. ChangesNormalization layout updates
Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
samples/python/29_rmsnorm.ipynb (1)
318-318: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRun
black-jupyteron the modified notebook.The trailing semicolons are not Black-formatted. As per coding guidelines, Python tutorial notebooks must use
black-jupyterformatting.Also applies to: 436-436
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@samples/python/29_rmsnorm.ipynb` at line 318, Run black-jupyter on the modified notebook, including the cells containing inv_var.set_name(...) and the corresponding code at the other referenced location, so Black removes the trailing semicolons and applies standard formatting throughout the notebook.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@samples/python/29_rmsnorm.ipynb`:
- Line 18: Update the Open in Colab badge URL in the notebook’s introductory
markdown cell to reference 29_rmsnorm.ipynb instead of 28_rmsnorm.ipynb, leaving
the badge text and repository path unchanged.
- Around line 338-345: Update the Python Binding API setup before the
variant_pack definition to allocate local out_gpu and inv_var_gpu tensors,
rather than relying on variables created by the earlier wrapper flow. Ensure the
existing out and inv_var bindings continue to reference these locally defined
outputs when running the flow independently.
In `@test/python/test_layernorm.py`:
- Line 58: Update every reference torch.var call used by the LayerNorm
expected-value calculations to pass correction=0: test/python/test_layernorm.py
lines 58-58, samples/python/20_layernorm_forward.ipynb lines 168-170 and
258-260, and samples/python/21_layernorm_backward.ipynb lines 174-176 and
337-339. Preserve the existing dimensions and keepdim arguments.
---
Nitpick comments:
In `@samples/python/29_rmsnorm.ipynb`:
- Line 318: Run black-jupyter on the modified notebook, including the cells
containing inv_var.set_name(...) and the corresponding code at the other
referenced location, so Black removes the trailing semicolons and applies
standard formatting throughout the notebook.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3b1aa3cf-7be5-4702-9738-7390a6624071
📒 Files selected for processing (11)
include/cudnn_frontend/node/dln.hinclude/cudnn_frontend/node/rmsnorm.hsamples/cpp/norm/batchnorm.cppsamples/cpp/norm/layernorm.cppsamples/cpp/norm/rmsnorm.cppsamples/python/20_layernorm_forward.ipynbsamples/python/21_layernorm_backward.ipynbsamples/python/22_layernorm_inference.ipynbsamples/python/29_rmsnorm.ipynbtest/python/test_layernorm.pytest/python/test_rmsnorm.py
| "cell_type": "markdown", | ||
| "metadata": {}, | ||
| "source": [ | ||
| "[](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/28_rmsnorm.ipynb)" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Point the Colab badge to this notebook.
The badge links to 28_rmsnorm.ipynb; it should target 29_rmsnorm.ipynb.
Proposed fix
-[](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/28_rmsnorm.ipynb)
+[](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/29_rmsnorm.ipynb)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| "[](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/28_rmsnorm.ipynb)" | |
| "[](https://colab.research.google.com/github/NVIDIA/cudnn-frontend/blob/main/samples/python/29_rmsnorm.ipynb)" |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@samples/python/29_rmsnorm.ipynb` at line 18, Update the Open in Colab badge
URL in the notebook’s introductory markdown cell to reference 29_rmsnorm.ipynb
instead of 28_rmsnorm.ipynb, leaving the badge text and repository path
unchanged.
| "variant_pack = {\n", | ||
| " x: x_gpu.detach(),\n", | ||
| " scale: scale_gpu.detach(),\n", | ||
| " bias: bias_gpu.detach(),\n", | ||
| " epsilon: eps_cpu,\n", | ||
| " out: out_gpu,\n", | ||
| " inv_var: inv_var_gpu,\n", | ||
| "}\n", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Allocate direct-binding outputs locally.
out_gpu and inv_var_gpu come from the earlier wrapper section. Running only the Python Binding API flow after setup fails with NameError.
Proposed fix
graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
+out_gpu = torch.empty_like(x_gpu)
+inv_var_gpu = torch.empty((*x_gpu.shape[:-1], 1), device=x_gpu.device, dtype=torch.float32)
+
# Mapping of (handles -> memory)
variant_pack = {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| "variant_pack = {\n", | |
| " x: x_gpu.detach(),\n", | |
| " scale: scale_gpu.detach(),\n", | |
| " bias: bias_gpu.detach(),\n", | |
| " epsilon: eps_cpu,\n", | |
| " out: out_gpu,\n", | |
| " inv_var: inv_var_gpu,\n", | |
| "}\n", | |
| out_gpu = torch.empty_like(x_gpu) | |
| inv_var_gpu = torch.empty((*x_gpu.shape[:-1], 1), device=x_gpu.device, dtype=torch.float32) | |
| "variant_pack = {\n", | |
| " x: x_gpu.detach(),\n", | |
| " scale: scale_gpu.detach(),\n", | |
| " bias: bias_gpu.detach(),\n", | |
| " epsilon: eps_cpu,\n", | |
| " out: out_gpu,\n", | |
| " inv_var: inv_var_gpu,\n", | |
| "}\n", |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@samples/python/29_rmsnorm.ipynb` around lines 338 - 345, Update the Python
Binding API setup before the variant_pack definition to allocate local out_gpu
and inv_var_gpu tensors, rather than relying on variables created by the earlier
wrapper flow. Ensure the existing out and inv_var bindings continue to reference
these locally defined outputs when running the flow independently.
| mean_expected = x_gpu.to(torch.float32).mean(dim=(1, 2, 3), keepdim=True) | ||
| inv_var_expected = torch.rsqrt(torch.var(x_gpu.to(torch.float32), dim=(1, 2, 3), keepdim=True) + epsilon_value) | ||
| mean_expected = x_gpu.to(torch.float32).mean(dim=-1, keepdim=True) | ||
| inv_var_expected = torch.rsqrt(torch.var(x_gpu.to(torch.float32), dim=-1, keepdim=True) + epsilon_value) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
torch.var default correction unbiased dim keepdim
💡 Result:
The torch.var function calculates the variance of an input tensor over specified dimensions. Its behavior regarding correction, dimensions, and output shape is as follows: Correction: The correction parameter represents the difference between the sample size and the sample degrees of freedom [1][2]. It defaults to 1, which corresponds to Bessel's correction for an unbiased estimate of the population variance [1][2]. Before PyTorch 2.0, this parameter was known as unbiased (a boolean where True equaled correction=1 and False equaled correction=0) [1][3]. Dim: The dim parameter specifies the dimension or dimensions along which to calculate the variance [1][2]. It accepts an int or a tuple of ints [1][3]. If it is set to None (the default), the variance is calculated over all dimensions of the input tensor [1][2]. Keepdim: The keepdim parameter is a boolean that determines whether the output tensor retains the reduced dimensions [1][2]. - If keepdim is True, the output tensor maintains the same number of dimensions as the input, with the reduced dimensions having a size of 1 [1][2]. - If keepdim is False (default), the reduced dimensions are removed (squeezed), resulting in an output tensor with fewer dimensions [1][4].
Citations:
- 1: https://docs.pytorch.org/docs/2.12/generated/torch.var.html
- 2: https://docs.pytorch.org/docs/stable/generated/torch.var.md
- 3: https://docs.pytorch.org/docs/2.3/generated/torch.var.html
- 4: https://docs.pytorch.org/docs/2.11/generated/torch.var.html
🏁 Script executed:
python3 - <<'PY'
from pathlib import Path
files = [
"test/python/test_layernorm.py",
"samples/python/20_layernorm_forward.ipynb",
"samples/python/21_layernorm_backward.ipynb",
]
for p in files:
path = Path(p)
print(f"\n=== {p} ({sum(1 for _ in path.open('r', encoding='utf-8'))} lines) ===")
text = path.read_text(encoding='utf-8')
for needle in ["torch.var(x_gpu.to(torch.float32), dim=-1, keepdim=True", "correction=", "unbiased="]:
if needle in text:
print(f"-- occurrences of {needle!r}:")
for i, line in enumerate(text.splitlines(), 1):
if needle in line:
print(f"{i}: {line}")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 1061
🏁 Script executed:
python3 - <<'PY'
from pathlib import Path
targets = {
"test/python/test_layernorm.py": range(1, 120),
"samples/python/20_layernorm_forward.ipynb": range(140, 280),
"samples/python/21_layernorm_backward.ipynb": range(150, 360),
}
for file, wanted in targets.items():
path = Path(file)
print(f"\n=== {file} ===")
lines = path.read_text(encoding="utf-8").splitlines()
for i in wanted:
if i <= len(lines):
line = lines[i - 1]
if any(k in line for k in ["layernorm", "var(", "mean(", "rsqrt(", "epsilon", "normalized_shape", "F.layer_norm", "native_layer_norm"]):
print(f"{i}: {line}")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 2511
🌐 Web query:
PyTorch layer_norm variance biased population variance unbiased=False docs
💡 Result:
PyTorch's LayerNorm implementation calculates variance using a biased estimator [1][2][3]. This is explicitly documented as being equivalent to using torch.var(input, correction=0) [1][2][3]. There is no parameter named unbiased or correction in the torch.nn.LayerNorm constructor to change this behavior [4][5][3]. The layer's variance calculation is fixed to the biased estimator [1].
Citations:
- 1: https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/normalization.py
- 2: https://docs.pytorch.org/docs/stable/generated/torch.nn.modules.normalization.LayerNorm.md
- 3: https://docs.pytorch.org/docs/main/generated/torch.nn.modules.normalization.LayerNorm.html
- 4: https://docs.pytorch.org/docs/2.3/generated/torch.nn.LayerNorm.html
- 5: https://docs.pytorch.org/docs/2.2/generated/torch.nn.LayerNorm.html
Use biased variance in the LayerNorm references
torch.var(...) defaults to correction=1, but LayerNorm uses the biased estimator (correction=0). Add correction=0 to each reference torch.var call in test/python/test_layernorm.py, samples/python/20_layernorm_forward.ipynb, and samples/python/21_layernorm_backward.ipynb.
📍 Affects 3 files
test/python/test_layernorm.py#L58-L58(this comment)samples/python/20_layernorm_forward.ipynb#L168-L170samples/python/20_layernorm_forward.ipynb#L258-L260samples/python/21_layernorm_backward.ipynb#L174-L176samples/python/21_layernorm_backward.ipynb#L337-L339
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/test_layernorm.py` at line 58, Update every reference torch.var
call used by the LayerNorm expected-value calculations to pass correction=0:
test/python/test_layernorm.py lines 58-58,
samples/python/20_layernorm_forward.ipynb lines 168-170 and 258-260, and
samples/python/21_layernorm_backward.ipynb lines 174-176 and 337-339. Preserve
the existing dimensions and keepdim arguments.
1 participant
Before submitting
pre-commit runand committed any formatting changes.Affected area
Documentation or samples
Summary
Layout and RMS Norm samples are updated to accept direct B,S,H tensor shapes
Why
cudnn Backend supports passing X tensor as (B,S,H) for LN/RMS norm. All FE samples transformed the input to (B*S, H ,1,1) which is not really necessary
API and compatibility impact
None
Testing
./bin/samples "LayerNorm Backward"
Filters: "LayerNorm Backward"
Randomness seeded to: 1937091254
All tests passed (7 assertions in 1 test case)
./bin/samples "LayerNorm Inference"
Filters: "LayerNorm Inference"
Randomness seeded to: 3002993422
All tests passed (9 assertions in 1 test case)
./bin/samples "LayerNorm Training"
Filters: "LayerNorm Training"
Randomness seeded to: 4011977597
All tests passed (7 assertions in 1 test case)
/bin/samples "RmsNorm Backward"
Filters: "RmsNorm Backward"
Randomness seeded to: 572308681
All tests passed (8 assertions in 1 test case)
./bin/samples "RmsNorm Inference"
Filters: "RmsNorm Inference"
Randomness seeded to: 4203307452
All tests passed (8 assertions in 1 test case)
Summary by CodeRabbit
Improvements
(batch, sequence, hidden)tensor layouts, including backward and inference workflows.Examples & Tests