test: pin AdaLNContinuous scale/shift order and gradients (PRPUNDIT-17) - #1100
Conversation
Closes PRPUNDIT-17. Plain-ops oracle plus fused-path CUDA check for NeMo chunk order, plus a backward-flow assertion.
Copilot review flagged that the fused-forward test only checked forward numerics, leaving the "pin gradients" goal covered only by the CPU plain-ops path (which doesn't exercise the primus::fused_ln_modulate custom autograd op). Add a CUDA test that runs backward() through the fused path and asserts finite, nonzero gradients reach both x and cond, mirroring the existing CPU backward test.
… tests Copilot flagged that torch.manual_seed(0) ran after AdaLNContinuous(...) construction in test_fused_forward_matches_plain_ops_formula and test_forward_matches_manual_layernorm_and_modulate, so the modulation Linear's parameter init wasn't covered by the seed. Move the seed call before construction in both so the whole test (module init + inputs) is reproducible.
The previous snippet-based commit (899c316) miscounted the replaced line range and left a duplicated `cond = torch.randn(...)` line. Drop the stray duplicate.
The previous snippet-based fix (bf5ed50) removed a duplicated `cond = torch.randn(...)` line but missed a duplicate `output = adaln(x, cond)` call left in test_forward_matches_manual_layernorm_and_modulate by the same line-range mistake. Drop the redundant second call.
My prior snippet-based edit inserted a bare literal "213" line and still left the duplicate `output = adaln(x, cond)` call in place, which would be a SyntaxError. Remove both, leaving a single call.
…tions Copilot review (PR #1087) flagged that `.sum().backward()` can yield an exactly-zero gradient under LayerNorm's zero-mean-per-token property if modulation ever becomes zero/constant (e.g. a future zero-init), making the gradient-flow assertions fragile for the wrong reason. Switch both backward tests to `.pow(2).sum().backward()`, which cannot cancel to zero unless every output element is exactly zero.
Addresses Copilot review feedback (review #5105410572): the fused forward/backward tests validated numerics but never confirmed the module actually dispatched to primus::fused_ln_modulate, so a future default flip to the plain-ops branch could silently stop exercising the fused kernel while these tests kept passing.
There was a problem hiding this comment.
🟢 Approval recommended
Only a minor test-naming inconsistency was found; the functional intent of the added coverage looks sound.
Pull request overview
Adds targeted unit tests to lock in AdaLNContinuous’s NeMo-aligned (scale, shift) chunk order and to validate both forward numerics and gradient flow for the fused CUDA path, while also covering a CPU plain-ops fallback path.
Changes:
- Add CUDA tests asserting fused
AdaLNContinuousmatches the reference formulanorm(x) * (1 + scale) + shift. - Add CUDA/CPU backward tests asserting gradients flow to both the input tensor and conditioning embedding.
- Add CPU-only plain-ops tests (via
config.adaln_plain_ops=True) to validate formula equivalence, zero-modulation reduction to LayerNorm, and chunk ordering.
File summaries
| File | Description |
|---|---|
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py |
Extends AdaLNContinuous test coverage for fused CUDA forward/backward behavior and adds CPU plain-ops oracle tests for formula/order invariants. |
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
TestIntel PR Steward: Checked the The failed Since this PR was opened, a fresh |
…1100) The test only asserts gradients on x and cond, not on the adaLN_modulation parameters, so "reaches_input_and_modulation" was misleading. Rename to reaches_input_and_conditioning to match the sibling fused-path test's naming.
|
TestIntel PR Steward: Follow-up on this poll (
No other open threads. |
There was a problem hiding this comment.
🟢 Approval recommended
The PR is test-only and the added checks align with the stated contract; remaining feedback is limited to minor test-oracle/convention improvements.
Review details
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:163
- For numeric tensor comparisons, the diffusion unit tests elsewhere commonly use
torch.testing.assert_close(...)(gives better diagnostics thanassert torch.allclose(...)). Consider switching this assertion (and the other newallcloseassertions below) toassert_close.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:221 - This test claims to validate a manual LayerNorm+modulation oracle, but
expectedis currently computed viaadaln.norm(x), which couples the oracle to the module under test. Usingtorch.nn.functional.layer_normhere would better match the test intent and PR summary (independent plain-ops formula).
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
|
TestIntel PR Steward: Follow-up on this poll (new
No push needed this round. |
There was a problem hiding this comment.
🟢 Approval recommended
The changes are isolated to test additions that directly exercise the intended AdaLNContinuous contracts (formula, chunk order, and gradient flow) without introducing production-code risk.
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
Same-repo refile of #1087 so Primus-CI-TAS runs against
AMD-AGI/Primus(fork PRs do not receive Docker Hub credentials, sobuild-docker/ torch unit tests never ran).This PR was created by dougljia via Test Gap Resolver.
Summary
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.pywith CPUTestAdaLNContinuousForwardPlainOps: independentnorm(x) * (1 + scale) + shiftoracle, zero-modulation LayerNorm reduction, and NeMo chunk order (first half = scale, second half = shift).xandcond.Closes test gap PRPUNDIT-17.
Test plan
pytest tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py(plain-ops CPU class + fused CUDA class; needsmegatron.core)black --check/isort --profile blackon the edited file