Add rlformers forward-pass features to ExecuTorch backbone for on-device export parity (#19096)

[ifed-ucsd]

Summary:

The 730M dense model checkpoint uses several features that the ExecuTorch XNNPACK export path did not implement. Without these, the exported model produces numerically incorrect output.

This diff adds support for 8 missing features:

1. normalize_tok_embeddings — scaleless RMSNorm after embedding lookup
2. qk_norm_before_rope — conversion from GenAI args (attention code already supported it)
3. scale_query_by — custom scalar multiplier on Q after QK norm
4. use_attn_o_gate — sigmoid gate on attention output using a learned linear projection of the layer input
5. use_attn_o_norm — scaleless per-head RMSNorm on attention output (applied before o_gate)
6. use_residual_gate — NormPreservingResidualConnection with learned per-dim gates for both attention and FFN residual connections
7. use_ffn_learnable_scales — RMSNormWithInputScale replacing standard post-FFN norm, computing rms_norm(gamma * x) instead of gamma * rms_norm(x)
8. output_soft_cap_temp — tanh(logits/temp) * temp soft capping on output logits

Additionally, this diff fixes a QK norm checkpoint compatibility issue: some checkpoints contain learned QK norm weights even though their params.json has qk_norm_affine=False (due to default changes after training). The ET model was creating ScalelessRMSNorm (no weight parameter) based on params.json, silently discarding the checkpoint's trained QK norm weights. The rlformers reference model loaded them correctly, causing ~53-67 dB SNR divergence. The fix peeks at the checkpoint state dict before model construction — if QK norm weights are present, qk_norm_affine is overridden to True so the ET model creates affine QK norms that load those weights.

All features are off by default (backward compatible). They activate when the corresponding fields are set in the checkpoint's params.json and propagated through model_args_conversion.

Weight key mappings added for: attention.og.weight, add_attn.gate, add_ffn.gate, post_ffn_norm.weight.

Reviewed By: chinnadhurai, digantdesai

Differential Revision: D102030169

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

✅ You can merge normally! (2 Unrelated Failures)

As of commit 3d5f0d6 with merge base d9688da ():

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

• pull / unittest-editable / macos / macos-job (gh) (similar failure)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / macos / macos-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[meta-codesync]

@ifed-ucsd has exported this pull request. If you are a Meta employee, you can view the originating Diff in D102030169.

[github-actions]

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.

[digantdesai]

Review automatically exported from Phabricator review in Meta.