Add CoreML-stable RMSNorm for llama eager paths (#19523) by telgamal-1 · Pull Request #19523 · pytorch/executorch

telgamal-1 · 2026-05-12T21:33:31Z

Summary:

The standard RMSNorm formulation x * rsqrt(mean(x²)) * weight is numerically unstable on CoreML/ANE because the explicit FP32 cast around the mean reduction is silently stripped from the lowered graph, leaving the squared sum to overflow in FP16. The ANE PTE then diverges from the eager reference even on checkpoints fine-tuned in BF16/FP16.

This diff introduces RMSNormCoreML in examples/models/llama/norm.py. The module expresses the normalization as x * sqrt(d) / vector_norm(x, dim=-1) — torch.linalg.vector_norm keeps the reduction in a single op that survives CoreML lowering, so FP16 inference remains stable.

To avoid 0 / 0 = NaN on zero-padded positions (chunked prefill in StaticAttentionIOManager pads each chunk to input_len with zeros), the denominator is floored with sqrt(dim * eps). This matches standard RMSNorm's rsqrt(mean(x²) + eps) semantics on a zero input and is large enough to survive fp16 — a plain 1e-6 underflows. Real (non-zero) tokens satisfy vector_norm(x) >> sqrt(dim * eps), so the floor is a no-op on real positions.

A new use_coreml_norm: bool = False field on ModelArgs opts into the new norm without disturbing existing models. When True, every llama-side norm site constructs RMSNormCoreML:

llama_transformer.py: attention_norm, ffn_norm, the final self.norm on Transformer.
attention.py: q_norm_fn / k_norm_fn in the affine QK-norm path, AND the else branch of _init_qk_norms (the scaleless / non-affine QK-norm path that the original landing missed).
static_attention.py: q_norm / k_norm in the scaleless path, propagated through from_attention_mha by detecting rms_norm_class is RMSNormCoreML.

The QNN/HTP export path is untouched and continues to use torch.nn.RMSNorm.

Differential Revision: D104862210

pytorch-bot · 2026-05-12T21:33:35Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Run pull request jobs on OSDC runners in shadow mode

✅ No Failures

As of commit fc26253 with merge base d8e4ffd ():
💚 Looks good so far! There are no failures yet. 💚

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

meta-codesync · 2026-05-12T21:33:44Z

@telgamal-1 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D104862210.

github-actions · 2026-05-12T21:34:27Z

This PR needs a `release notes:` label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

metascroy · 2026-05-12T21:39:52Z

+            eps (float, optional): Stored for API compatibility; ignored in the math.
+
+        Attributes:
+            eps (float): Stored for API compatibility; not consumed by `_norm`.


Can we assert eps is 0 rather than silently drop it?

Added an explicit assert of eps is set to 0

Summary: The standard `RMSNorm` formulation `x * rsqrt(mean(x²)) * weight` is numerically unstable on CoreML/ANE because the explicit FP32 cast around the mean reduction is silently stripped from the lowered graph, leaving the squared sum to overflow in FP16. The ANE PTE then diverges from the eager reference even on checkpoints fine-tuned in BF16/FP16. This diff introduces `RMSNormCoreML` in `examples/models/llama/norm.py`. The module expresses the normalization as `x * sqrt(d) / vector_norm(x, dim=-1)` — `torch.linalg.vector_norm` keeps the reduction in a single op that survives CoreML lowering, so FP16 inference remains stable. To avoid `0 / 0 = NaN` on zero-padded positions (chunked prefill in `StaticAttentionIOManager` pads each chunk to `input_len` with zeros), the denominator is floored with `sqrt(dim * eps)`. This matches standard RMSNorm's `rsqrt(mean(x²) + eps)` semantics on a zero input and is large enough to survive fp16 — a plain `1e-6` underflows. Real (non-zero) tokens satisfy `vector_norm(x) >> sqrt(dim * eps)`, so the floor is a no-op on real positions. A new `use_coreml_norm: bool = False` field on `ModelArgs` opts into the new norm without disturbing existing models. When True, every llama-side norm site constructs `RMSNormCoreML`: - `llama_transformer.py`: `attention_norm`, `ffn_norm`, the final `self.norm` on `Transformer`. - `attention.py`: `q_norm_fn` / `k_norm_fn` in the affine QK-norm path, AND the `else` branch of `_init_qk_norms` (the scaleless / non-affine QK-norm path that the original landing missed). - `static_attention.py`: `q_norm` / `k_norm` in the scaleless path, propagated through `from_attention_mha` by detecting `rms_norm_class is RMSNormCoreML`. The QNN/HTP export path is untouched and continues to use `torch.nn.RMSNorm`. Differential Revision: D104862210

metascroy · 2026-05-14T18:54:32Z

        self.weight.requires_grad = False


+class RMSNormCoreML(torch.nn.Module):


How does this differ from:

executorch/examples/apple/coreml/llama/llama_transformer.py

Line 112 in 7cd209d

class CoreMLRMSNorm(torch.nn.Module):

Can we consolidate? Putting it here is fine, but then import this version into examples/apple/coreml/llama/llama_transformer.py.

I imported the new version in examples/apple/coreml/llama/llama_transformer.py because it was tested to not produce NaN in QAT

metascroy · 2026-05-14T18:58:51Z

-                )
-                if self.has_kv_weights:
-                    self.k_norm_fn = RMSNorm(
+                if args.use_coreml_norm:


Does this have to be integrated so far down?

Could we not leave llama_transformer/static attention as is, introduce the new norm in norm.py, and then do a module swap from RMSNorm -> CoreMLRMsNorm at export time?

addresed and now using the strategy to replace_rms_norm_for_coreml_

Summary: The standard `RMSNorm` formulation `x * rsqrt(mean(x²)) * weight` is numerically unstable on CoreML/ANE because the explicit FP32 cast around the mean reduction is silently stripped from the lowered graph, leaving the squared sum to overflow in FP16. The ANE PTE then diverges from the eager reference even on checkpoints fine-tuned in BF16/FP16. This diff introduces `RMSNormCoreML` in `examples/models/llama/norm.py`. The module expresses the normalization as `x * sqrt(d) / vector_norm(x, dim=-1)` — `torch.linalg.vector_norm` keeps the reduction in a single op that survives CoreML lowering, so FP16 inference remains stable. To avoid `0 / 0 = NaN` on zero-padded positions (chunked prefill in `StaticAttentionIOManager` pads each chunk to `input_len` with zeros), the denominator is floored with `sqrt(dim * eps)`. This matches standard RMSNorm's `rsqrt(mean(x²) + eps)` semantics on a zero input and is large enough to survive fp16 — a plain `1e-6` underflows. Real (non-zero) tokens satisfy `vector_norm(x) >> sqrt(dim * eps)`, so the floor is a no-op on real positions. A new `use_coreml_norm: bool = False` field on `ModelArgs` opts into the new norm without disturbing existing models. When True, every llama-side norm site constructs `RMSNormCoreML`: - `llama_transformer.py`: `attention_norm`, `ffn_norm`, the final `self.norm` on `Transformer`. - `attention.py`: `q_norm_fn` / `k_norm_fn` in the affine QK-norm path, AND the `else` branch of `_init_qk_norms` (the scaleless / non-affine QK-norm path that the original landing missed). - `static_attention.py`: `q_norm` / `k_norm` in the scaleless path, propagated through `from_attention_mha` by detecting `rms_norm_class is RMSNormCoreML`. The QNN/HTP export path is untouched and continues to use `torch.nn.RMSNorm`. Differential Revision: D104862210

metascroy · 2026-05-15T00:24:35Z

LGTM! You need to run the lintrunner, though

meta-codesync · 2026-05-15T00:40:05Z

@telgamal-1 has imported this pull request. If you are a Meta employee, you can view this in D104862210.

This PR was reopened (likely due to being reverted), so your approval was removed. Please request another review.

telgamal-1 requested a review from lucylq as a code owner May 12, 2026 21:33

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 12, 2026

meta-codesync Bot added fb-exported meta-exported labels May 12, 2026

metascroy reviewed May 12, 2026

View reviewed changes

meta-codesync Bot changed the title ~~Add CoreML-stable RMSNorm for llama eager paths~~ Add CoreML-stable RMSNorm for llama eager paths (#19523) May 14, 2026

telgamal-1 force-pushed the export-D104862210 branch from cda18f8 to b2acb39 Compare May 14, 2026 17:43

telgamal-1 force-pushed the export-D104862210 branch from b2acb39 to b5af889 Compare May 14, 2026 17:53

metascroy reviewed May 14, 2026

View reviewed changes

telgamal-1 force-pushed the export-D104862210 branch from b5af889 to ae1926c Compare May 14, 2026 20:30

telgamal-1 force-pushed the export-D104862210 branch from ae1926c to fc26253 Compare May 15, 2026 00:15

metascroy previously approved these changes May 15, 2026

View reviewed changes

telgamal-1 closed this May 15, 2026

telgamal-1 reopened this May 15, 2026

metascroy approved these changes May 15, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add CoreML-stable RMSNorm for llama eager paths (#19523)#19523

Add CoreML-stable RMSNorm for llama eager paths (#19523)#19523
telgamal-1 wants to merge 1 commit into
pytorch:mainfrom
telgamal-1:export-D104862210

telgamal-1 commented May 12, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

pytorch-bot Bot commented May 12, 2026 •

edited

Loading

Uh oh!

meta-codesync Bot commented May 12, 2026

Uh oh!

github-actions Bot commented May 12, 2026

Uh oh!

metascroy May 12, 2026

Uh oh!

telgamal-1 May 14, 2026

Uh oh!

metascroy May 14, 2026

Uh oh!

telgamal-1 May 15, 2026

Uh oh!

metascroy May 14, 2026

Uh oh!

telgamal-1 May 15, 2026

Uh oh!

metascroy commented May 15, 2026

Uh oh!

meta-codesync Bot commented May 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		self.weight.requires_grad = False


		class RMSNormCoreML(torch.nn.Module):

Conversation

telgamal-1 commented May 12, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented May 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❗ 1 Active SEVs

✅ No Failures

Uh oh!

meta-codesync Bot commented May 12, 2026

Uh oh!

github-actions Bot commented May 12, 2026

This PR needs a release notes: label

Uh oh!

metascroy May 12, 2026

Choose a reason for hiding this comment

Uh oh!

telgamal-1 May 14, 2026

Choose a reason for hiding this comment

Uh oh!

metascroy May 14, 2026

Choose a reason for hiding this comment

Uh oh!

telgamal-1 May 15, 2026

Choose a reason for hiding this comment

Uh oh!

metascroy May 14, 2026

Choose a reason for hiding this comment

Uh oh!

telgamal-1 May 15, 2026

Choose a reason for hiding this comment

Uh oh!

metascroy commented May 15, 2026

Uh oh!

meta-codesync Bot commented May 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

telgamal-1 commented May 12, 2026 •

edited by meta-codesync Bot

Loading

pytorch-bot Bot commented May 12, 2026 •

edited

Loading

This PR needs a `release notes:` label