[shoumikhin]

What is wrong today

RecomposeRmsNorm rebuilds a decomposed RMSNorm into one rms_norm node. To find the norm's input it walks back over casts, but it only knew aten.to.dtype, and edge graphs carry the cast as _to_copy or _to_dim_order_copy. For a Qwen- or Llama-style norm (upcast to float32, normalize, cast back), the rebuilt rms_norm read the float32 cast, and the cast back was lost.

bfloat16 models then failed later, for example:

masked_scatter: expected self and source to have same dtypes
RuntimeError: Missing out variants: {'aten::rms_norm'}

What this change does

The pass assumes QNN's RmsNorm takes one dtype throughout. So it now rebuilds the norm only when its input and its output share a dtype, and no cast in between rounds below that dtype:

• It skips casts that change only the dtype (to.dtype and type_as in captured graphs, _to_copy and _to_dim_order_copy in edge graphs): the output keeps the input's layout, device and dim order.
• It walks back to the nearest tensor in the output dtype, but not past a cast to a narrower dtype or one that changes the dim order.
• Only such casts may sit between the normalization and the weight. Before, a skipped norm let a later multiply be rebuilt as the norm with the wrong weight.
• Outside quantization capture, the dtype also needs a QNN type.
• The tensor it scales has to be the tensor it squares, past such casts (T5's norm squares a float32 copy), and the square has to be pow(2) or h * h. Before, the scaled tensor was not checked and could be dropped.

So a float16 Qwen- or Llama-style norm is rebuilt on its float16 input, and a bfloat16 norm stays decomposed on the CPU unless it is about to be quantized. A norm that rounds to float16 before a float32 weight also stays decomposed now. Before, it was rebuilt in float32, which dropped that rounding.

Tests

test_recompose_rms_norm_keeps_the_norm_dtype covers the Qwen, Llama and T5 forms in float16, float32 and bfloat16, with and without dim order, channels-last inputs, narrowing casts, a multiply after the norm, and quantization capture with to.dtype and type_as. Each case checks which tensor the norm reads and compares the output with eager. test_recompose_rms_norm_needs_the_norm_of_one_tensor checks that a different scaled tensor, a cube and a product of two tensors stay decomposed.

A second commit makes RemoveRedundancy, which runs before this pass, keep a _to_copy with memory_format that also changes the dtype, such as x.to(torch.float16, memory_format=torch.contiguous_format). Without dim order it dropped every such cast. A new test checks both cases.

The new tests pass on Linux x86_64 with torch 2.14 and 2.16 nightly, and each fails when its rule is removed. With the QNN SDK (SM8650), a float16 Qwen-style norm is delegated as one rms_norm, and a bfloat16 norm followed by masked_scatter and a small bfloat16 Llama-style block now lower. The rest of test_passes.py passes, except two tests that fail the same way on main in my environment.

cc @cbilgin @psiddh

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24428

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ 112 Pending, 1 Unrelated Failure

As of commit 22dc132 with merge base cec2e73 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / android / run-emulator (gh) (detected as infra flaky with no log or failing log classifier)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 25m 39s —— View job

---

PR Review: #24428

No blockers, but in rare mixed-dtype graphs the pass can still build an rms_norm that QNN rejects, or rebuild a later mul with the wrong weight; the float32-weight handling and the test tolerances also need another look.

Detailed review

Summary

The PR makes RecomposeRmsNorm walk past _to_copy/_to_dim_order_copy casts to the nearest input in the norm's output dtype. It leaves the norm decomposed when that dtype has no QNN type.

Nothing blocks the PR. The problems are:

• two gaps in mixed-dtype graphs, one of which gives silently wrong output;
• a design question about float32 weights;
• missing test coverage.

All findings come from reading the code. Python can't run in this environment, so none of them were executed.

Code Quality

1. A skipped norm lets a later mul be rebuilt with the wrong weight (backends/qualcomm/_passes/recompose_rms_norm.py:151-154).

   • When the new gate skips a norm, its weight mul y stays in the graph.
   • The "*" in the patterns spans up to three nodes (max_wildcard_life=3, _passes/utils.py:278), so it can cover y and the cast back.
   • A following z = y * s then matches with s as gamma, when s is a [D]-shaped placeholder, get_attr or dq.
   • Example: a bfloat16 Qwen-style norm followed by * s, with s in float32. z is float32, the walk stops at the float32 upcast and the gate passes. z becomes rms_norm(x.to(float32), [D], s), and the real weight is silently dropped.
   • Before this PR, y was always rebuilt first, so this graph failed loudly at lowering. The older skips (no gamma found, the shape checks) already had this hole; the new gate opens it for bfloat16 norms.
   • Suggested fix: recompose only when the nodes between last_mul_node and rms_norm_pattern[2] are dtype-only casts.
   • Fix this →

2. The new dtype gate does not check the weight (recompose_rms_norm.py:151-153).

   • The partitioner's has_unsupported_dtype (partition/utils.py:149-159, called at qnn_partitioner.py:177) rejects a node if any unquantized input has no QNN type. The pass checks only the output dtype.
   • Example: UpcastNorm(torch.bfloat16) with a float32 x has a float32 output, so the norm is rebuilt with a bfloat16 gamma. The partitioner keeps it on CPU, and lowering fails with the same Missing out variants: {'aten::rms_norm'} error this PR fixes.
   • A float16 x with a bfloat16 weight fails the same way, since it promotes to float32.
   • Suggested fix: also require gamma_node.meta["val"].dtype to have a QNN type.
   • Fix this →

3. With a float32 weight, the target dtype is taken after type promotion (recompose_rms_norm.py:146-150).

   • The comment says rms_norm computes in its input's dtype, which "has to be the pattern's output dtype". That doesn't match eager: aten.rms_norm upcasts float16/bfloat16 internally and applies the weight after casting back to the input dtype. Its output dtype is promote_types(input, weight).
   • Because the target is last_mul_node's dtype, a Qwen-style float16 x with a float32 weight is rebuilt on x.to(float32). That drops the model's h.to(float16) rounding before the weight multiply, which is the kind of rounding the comment at lines 103-105 is meant to preserve. The float32-weight test case only passes because of atol=rtol=1e-2 at test_passes.py:561.
   • A cast-free float16 norm with a float32 weight is now left decomposed: x is float16 and the target is float32. The old pass rebuilt it, and the PR description doesn't mention this change.
   • If QNN's RmsNorm needs a single dtype for input, weight and output, the comment should say so, and the weight then has to match too (see item 2).
   • Otherwise, target the dtype of the weight mul's other operand (the cast back). That rebuilds both cases on the float16 x and matches what eager computes for the original module.

4. The walk still skips narrowing casts (recompose_rms_norm.py:106-109).

   • The loop walks past every cast whose output is not the target dtype, including narrowing ones. So the guarantee in the comment above it only covers round trips that end in the target dtype.
   • Example: x32 → .to(float16) → norm → .to(float32) → w32 * is rebuilt as a float32 norm on x32, which drops the float16 rounding and compute.
   • Suggested fix: walk only past casts that don't narrow, or narrow the comment's claim.
   • Fix this →

5. _is_cast compares dim_order with the contiguous order instead of the input's (recompose_rms_norm.py:99).

   • MemoryFormatOpsPass gives a preserve-format cast the input's dim order (exir/passes/memory_format_ops_pass.py:114-117).
   • So a dtype-only cast of a channels-last tensor isn't treated as a cast, and its norm is not recomposed.
   • A cast that also turns a channels-last input contiguous is skipped, which contradicts "Only casts that change nothing but the dtype are skipped".
   • Suggested fix: compare with list(node.args[0].meta["val"].dim_order()).
   • Fix this →

Testing

• The test is too lenient to catch numeric bugs (backends/qualcomm/tests/test_passes.py:509, :561).
  • UpcastNorm's weight is torch.ones, and the comparison uses atol=rtol=1e-2. A wrong or dropped gamma (item 1) or a dropped rounding (item 3) still passes.
  • ScaledNorm above uses torch.randn and default tolerances.
  • Only the float32-weight case should need a looser bound. I haven't run the other cases to confirm they pass at default tolerances.
• Cases with no test:
  • the Llama-style pattern (x * x, x.float(), .type_as(x); examples/models/llama/norm.py:34-41), which the description says was checked on device;
  • the round trip described in the _get_input_node comment (float32 → float16 → float32 before the norm);
  • _is_cast rejecting a cast with extra kwargs or a non-contiguous dim_order. Dropping those two conditions (lines 98-99 of the pass) still passes every case.
• The new cases are in the wrong test. They are appended to test_recompose_rms_norm_only_takes_last_dim_with_weight, which is about shape gating. A separate test would make a failure point at the dtype logic.
• Fix this →

Recommendation

Needs Discussion

Nothing blocks. Item 3 needs an answer on whether QNN's RmsNorm requires one dtype throughout, because that decides how items 2 and 3 should be fixed. Items 1 and 2 are small enough to close in this PR.

---

• Gather context (diff, changed files, CLAUDE.md)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

Thanks for the review. The latest push changes this:

1. Wrong weight after a skipped norm: fixed. Only casts that change nothing but the dtype may now sit between the normalization and the weight multiply, so a later * s is no longer rebuilt as the norm. The new test has a bfloat16 norm followed by a float32 scale.
2. Weight dtype: not changed, because this case is never rebuilt. When the weight's dtype differs from the multiply's output, the weight reaches the multiply through a cast in the edge graph. The pass then finds no parameter or constant as gamma, so the norm stays decomposed. I checked float32 input with a bfloat16 weight, float16 input with a bfloat16 weight, and float32 input with a float16 weight: no rms_norm is built in any of them.
3. One dtype or mixed: I found nothing that says QNN's RmsNorm takes different dtypes for input, weight and output, so I went with one dtype throughout. The norm is rebuilt only when its input and its output share a dtype, and no cast in between rounds below that dtype. So a float16 Qwen-style norm with a float32 weight now stays decomposed. Before, it was rebuilt in float32 and lost the float16 rounding. The description now says this.
4. Narrowing casts: fixed. The walk stops at a cast to a dtype narrower than the output.
5. dim_order: fixed. A cast is compared with its input's dim order. The test has a channels-last input, once with a cast that keeps its dim order and once with a cast to contiguous, which stops the rebuild.

Tests: the dtype cases are now in their own test, test_recompose_rms_norm_keeps_the_norm_dtype, with random weights and default tolerances. It adds a Llama-style norm (type_as), float16 rounding before a float32 norm (rebuilt on the rounded tensor) and inside one (left decomposed), and the channels-last cases.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 36m 11s —— View job

---

PR Review: #24428

During quantization capture, the new check between the normalization and the weight stops a norm that casts back with type_as from being recomposed, which is a regression that needs a fix before merge; the other two items are older bugs next to the change.

Detailed review

Summary

The PR makes RecomposeRmsNorm rebuild a norm only on an input in the norm's output dtype. It finds that input by walking past casts that change only the dtype, and requires that only such casts sit between the normalization and the weight.

• Item 1 is a regression in the quantization flow, caused by the new check on the weight side.
• Items 2 and 3 are older bugs next to the change.

I covered both changed files. Python can't run in this environment, so nothing was executed; every finding comes from reading the code.

Code Quality

1. During quantization capture, a norm that casts back with type_as is no longer recomposed (backends/qualcomm/_passes/recompose_rms_norm.py:91, :133-140).

   • Why it misses:
     • In capture mode, skip_targets is only torch.ops.aten.to.dtype.
     • A captured graph keeps h.type_as(x) as aten.type_as.default, even when the dtype doesn't change.
       • Arm's annotation pipeline runs a pass that removes such no-op calls on parameters (backends/arm/_passes/remove_redundant_type_as_pass.py:24, arm_pass_manager.py:808).
       • _cast_source in backends/transforms/fuse_rms_norm.py:210 matches the same op.
     • QNN's capture-mode RemoveRedundancy removes only _assert_tensor_metadata, so the type_as node stays.
     • The new walk back from the weight multiply stops at that node, normalized is pattern[2] fails, and _get_norm_input returns None.
     • The old pass let the "*" take the type_as and built one rms_norm on x.
   • Trigger: the Meta reference Llama RMSNorm (self._norm(x.float()).type_as(x) * self.weight) or the PR's TypeAsNorm, in float32 or bfloat16, run through RecomposeRmsNorm(quantization_capture=True).
     • The norm stays decomposed through annotation, so its ops are quantized one by one instead of as one rms_norm.
     • For bfloat16, this contradicts the description, which says a bfloat16 norm stays decomposed unless it is about to be quantized.
   • Test gap: TypeAsNorm runs only on the edge path (backends/qualcomm/tests/test_passes.py:561). The capture rows (:567-568) use only UpcastNorm, which casts back with to.dtype.
   • Fix: add torch.ops.aten.type_as.default to the capture skip_targets. Add capture rows such as (TypeAsNorm(f32), f32, c, True, True, [(f32, True)]) and the bfloat16 version.
   • Related:
     • _is_cast reads layout, device and dim order only from node.kwargs (:102-111).
     • aten.to.dtype takes memory_format positionally. So in capture, a to.dtype that makes a channels-last input contiguous still counts as changing only the dtype. The values are unaffected.
     • Comparing node.args[0].meta["val"] with node.meta["val"] instead would apply one rule to all four cast ops. _cast_source already compares device and layout this way (fuse_rms_norm.py:194).
   • Basis: I confirmed the mechanism by reading the code. That captured graphs keep aten.type_as.default is inferred from those two passes, not from a run.
   • Fix this →

2. Nothing checks that the tensor being scaled is the norm's input (recompose_rms_norm.py:123, :140). This bug predates the PR.

   • Problem: the input is taken from pattern[6].args[0], the base of the square, and the new check verifies only the weight side. Nothing checks that:
     • the other operand of pattern[2] (the h in h * rsqrt(...)) is that same tensor;
     • both operands of the Llama-form h * h are the same tensor;
     • the pow exponent is 2.
   • Trigger: y * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * w is rebuilt as rms_norm(x, [D], w), and y is silently dropped.
   • Why fix it here: _match_rms_norm_core in backends/transforms/fuse_rms_norm.py:104-108 already checks the scaled tensor and the exponent. This PR rewrote how the input is chosen, and its test asserts which tensor the norm reads.
   • Fix this →

3. With _skip_dim_order=True, the dim-order rule never sees the cast it targets (recompose_rms_norm.py:102-111, test_passes.py:559). The cause predates the PR and is outside the diff.

   • Cause: RemoveRedundancy runs before this pass in the to_edge pipeline (qnn_pass_manager.py:344). Its _to_copy_op_condition (remove_redundancy.py:51-52) deletes every _to_copy that has a memory_format kwarg, including one that changes the dtype. It doesn't re-trace, so the metadata downstream stays stale.
   • Trigger: ContiguousUpcastNorm(f16) lowered with _skip_dim_order=True.
     • _to_copy(x, dtype=torch.float32, memory_format=torch.contiguous_format) is deleted.
     • The pow then reads the float16 x, while its metadata still says float32.
     • The norm is rebuilt on x, although under the PR's rule this cast keeps the norm decomposed.
   • Wider effect: any x.to(dtype, memory_format=...) in that flow loses its dtype change the same way.
   • Why the test misses it: it calls the pass on its own and checks the rule only with dim order on.
   • Fix: make _to_copy_op_condition also require an unchanged dtype, as _dim_order_op_condition does.
   • Fix this →

Recommendation

Request Changes

Item 1 makes the quantization flow worse for norms written with type_as. The fix is one skip_targets entry and two test rows. Items 2 and 3 predate the PR and can be follow-ups.

---

• Gather context (diff, changed files, CLAUDE.md)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

[shoumikhin]

Thanks. The latest push changes this:

1. type_as in quantization capture: fixed. aten.type_as now counts as a cast in captured graphs. _is_cast also compares the cast's output with its input (layout, device, dim order) instead of reading kwargs, so one rule covers every cast op, however it passes its arguments. The test runs TypeAsNorm in capture for float32 and bfloat16.
2. The scaled tensor and the exponent: fixed here. The norm is rebuilt only when the tensor it scales is the tensor it squares, and the square is pow(2) or h * h. The comparison looks past dtype-only casts, because T5's norm squares a float32 copy of the tensor it scales; the test covers T5's norm too. A new test checks that a different scaled tensor, a cube and a product of two tensors stay decomposed.
3. RemoveRedundancy dropping a dtype change: fixed in a second commit. A _to_copy with memory_format is now removed only when it keeps the dtype, as _to_dim_order_copy already was. A new test exports such casts without dim order: the layout-only one is removed, and the float16 one stays.

[shoumikhin]

@claude /pr-review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[claude]

PR Review: #24428

• Gather context (diff, changed files, CLAUDE.md)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

View job run