test: pin AdaLNContinuous scale/shift order and gradients (PRPUNDIT-17) - #1087
test: pin AdaLNContinuous scale/shift order and gradients (PRPUNDIT-17)#1087jiagaoxiang wants to merge 8 commits into
Conversation
Closes PRPUNDIT-17. Plain-ops oracle plus fused-path CUDA check for NeMo chunk order, plus a backward-flow assertion.
There was a problem hiding this comment.
🟡 Changes recommended
The new fused CUDA coverage validates forward numerics but does not exercise the fused custom op’s backward path, leaving the “pin gradients” goal only partially covered.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds targeted unit tests to lock in AdaLNContinuous’s modulation semantics (scale/shift chunk order and the norm(x) * (1 + scale) + shift formula) and to expand coverage to a CPU “plain ops” execution path in the Flux/Megatron diffusion stack.
Changes:
- Add a CUDA test asserting the fused AdaLNContinuous forward path matches the plain-ops formula (NeMo chunk order).
- Add a new CPU-only test class that validates the plain-ops branch numerics (manual oracle, zero-modulation reduction to LayerNorm, chunk order).
- Add a backward-flow assertion for gradients reaching both
xandcondin the plain-ops branch.
File summaries
| File | Description |
|---|---|
| tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py | Extends AdaLNContinuous unit tests with fused-forward formula checking and new CPU plain-ops correctness/gradient checks. |
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.
| torch.manual_seed(0) | ||
| x = torch.randn(ATTENTION_SEQ_LEN, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda() | ||
| cond = torch.randn(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda() | ||
|
|
||
| output = adaln(x, cond) | ||
| emb = adaln.adaLN_modulation(cond) | ||
| scale, shift = torch.chunk(emb, 2, dim=1) | ||
| expected = adaln.norm(x) * (1 + scale) + shift | ||
| assert torch.allclose(output, expected, atol=1e-5, rtol=1e-5) |
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.
jiagaoxiang
left a comment
jiagaoxiang
left a comment
There was a problem hiding this comment.
Addressed the Copilot finding inline; see the thread reply for details.
| emb = adaln.adaLN_modulation(cond) | ||
| scale, shift = torch.chunk(emb, 2, dim=1) | ||
| expected = adaln.norm(x) * (1 + scale) + shift | ||
| assert torch.allclose(output, expected, atol=1e-5, rtol=1e-5) |
There was a problem hiding this comment.
TestIntel PR Steward: Good catch — taken. Added test_fused_forward_backward_reaches_input_and_conditioning to TestAdaLNContinuous (pushed in a4b9621), which runs .sum().backward() through the fused CUDA primus::fused_ln_modulate path and asserts finite, nonzero gradients reach both x and cond, mirroring the existing CPU plain-ops backward test.
|
TestIntel PR Steward: Follow-up on this poll:
|
There was a problem hiding this comment.
🟢 Approval recommended
The PR is test-only and the added coverage aligns with the stated goal, with only minor test determinism/oracle-strength nits noted.
Review details
Suppressed comments (3)
Previously missed (2) — in code that hasn't changed since the last review.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:155
torch.manual_seed(0)is called after constructingAdaLNContinuous(...), so the module parameter initialization isn’t seeded and the test can be less deterministic than intended (especially if fused-vs-plain numeric deltas depend on modulation weight distribution). Move the seed call before creatingadalnso both parameters and inputs are reproducible.
This issue also appears on line 206 of the same file.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:217
- The test name suggests the LayerNorm portion is computed “manually”, but the expected value currently uses
adaln.norm(x)(the module under test). Usingtorch.nn.functional.layer_normhere makes the oracle independent ofadaln.normimplementation details and aligns better with the test intent.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:211
- Same determinism issue as the fused test:
torch.manual_seed(0)happens afterAdaLNContinuous(...)is constructed, so the module initialization isn’t seeded. If the goal is fully reproducible CPU plain-ops behavior, seed before instantiatingadaln.
config = self._make_config()
adaln = AdaLNContinuous(config, conditioning_embedding_dim=HIDDEN_DIM_FLUX, modulation_bias=False)
torch.manual_seed(0)
x = torch.randn(ATTENTION_SEQ_LEN, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX)
cond = torch.randn(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX)
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
… 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.
There was a problem hiding this comment.
🟡 Changes recommended
There is an accidental duplicate output = adaln(x, cond) call in the new CPU test that should be removed to avoid confusion and unnecessary work.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Lite
| output = adaln(x, cond) | ||
| output = adaln(x, cond) |
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.
|
TestIntel PR Steward: Follow-up on this poll:
Still holding off on merge: |
There was a problem hiding this comment.
🔵 Needs a closer look
The new backward tests use a loss (output.sum()) that can yield exactly-zero gradients under LayerNorm in common edge conditions (e.g., if modulation is zero/constant), making the gradient-flow assertions fragile.
Review details
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:177
- Using
.sum().backward()can produce exactly-zero gradients when the modulation happens to be zero/constant (e.g., if AdaLNContinuous ever adopts AdaLN-style zero-init), becauseLayerNorm(x)has zero-mean per token soLayerNorm(x).sum()is a constant. This makes the test fragile and can fail for the wrong reason; prefer a loss that can’t cancel under LayerNorm and explicitly force nonzero modulation weights socondmust receive gradient.
This issue also appears on line 253 of the same file.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:258
- This backward test also uses
.sum().backward(), which can yield a constant-zero loss (and thus exactly-zero gradients) if the modulation becomes zero/constant (e.g., future zero-init). To make the test robust and ensurecondgenuinely influences the output, force a nonzero modulation weight and use a non-canceling loss likeoutput.pow(2).sum().
config = self._make_config()
adaln = AdaLNContinuous(config, conditioning_embedding_dim=HIDDEN_DIM_FLUX, modulation_bias=False)
x = torch.randn(ATTENTION_SEQ_LEN, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX, requires_grad=True)
cond = torch.randn(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX, requires_grad=True)
adaln(x, cond).sum().backward()
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
…tions Copilot review (PR AMD-AGI#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.
|
TestIntel PR Steward: Follow-up on this poll:
|
There was a problem hiding this comment.
🔵 Needs a closer look
The new “fused” CUDA tests do not currently assert that the fused dispatch path is actually being exercised, so they can pass while leaving fused forward/backward behavior untested.
Review details
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:153
- These tests are intended to validate the fused CUDA path, but they don’t currently assert that the module is actually dispatching to the fused custom op. If a future change flips the default to the non-fused (or plain-ops) branch, this test would still pass and stop protecting the fused kernel behavior. Add an explicit assertion that fused dispatch is enabled.
tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py:171
- Same as the forward test: this gradient-flow test should assert it is exercising the fused CUDA path. Otherwise it can silently become a plain-ops / non-fused test while still passing, leaving the fused backward untested.
adaln = AdaLNContinuous(config, conditioning_embedding_dim=HIDDEN_DIM_FLUX).cuda()
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
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.
|
TestIntel PR Steward: Follow-up on this poll:
|
There was a problem hiding this comment.
🟢 Approval recommended
The changes are isolated to unit tests and the added assertions align with the implementation contract for AdaLNContinuous’s formula, chunk ordering, and gradient propagation.
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
|
Superseded by same-repo PR #1100 (head |
…7) (#1100) Same-repo refile of #1087 so Primus-CI-TAS runs against `AMD-AGI/Primus` (fork PRs do not receive Docker Hub credentials, so `build-docker` / torch unit tests never ran). This PR was created by dougljia via Test Gap Resolver. ## Summary - Extends `tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py` with CPU `TestAdaLNContinuousForwardPlainOps`: independent `norm(x) * (1 + scale) + shift` oracle, zero-modulation LayerNorm reduction, and NeMo chunk order (first half = scale, second half = shift). - Adds a fused-path CUDA check against the same formula, plus a backward-flow assertion on `x` and `cond`. - Existing CUDA shape / invalid-norm tests are unchanged. Closes test gap **PRPUNDIT-17**. ## Test plan - [ ] GPU CI: `pytest tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py` (plain-ops CPU class + fused CUDA class; needs `megatron.core`) - [x] `black --check` / `isort --profile black` on the edited file --------- Co-authored-by: Xiaoming-AMD <Xiaoming.Peng@amd.com>
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