Map norm_zero_centered_gamma onto the existing rms_norm_add_unit_offset - #22573
Map norm_zero_centered_gamma onto the existing rms_norm_add_unit_offset#22573telgamal-1 wants to merge 1 commit into
Conversation
Summary:
ALTERNATIVE to D117233461, implementing minguo's review suggestion. Same bug, one
fewer concept.
rlformers stores gamma offset by -1 when `norm_zero_centered_gamma` is set, so the
effective scale is `weight + 1`. ET dropped the flag on the params -> ModelArgs
conversion, so a checkpoint carrying it loaded without error and produced
systematically wrong activations.
D117233461 fixes that by adding a **new** `ModelArgs.norm_zero_centered_gamma` and
wiring it into `RMSNormWithInputScale`. But ET already has a field meaning exactly
this: `ModelArgs.rms_norm_add_unit_offset` (`model_args.py:56`), and `RMSNorm.forward`
already implements it as `output * (1 + weight)` (`norm.py:37-41`). It is already wired
into `attention_norm` and `ffn_norm` (`llama_transformer.py:226,232`), `q_norm`/`k_norm`
(`attention.py:468,474`) and the final norm (`llama_transformer.py:345`) — nothing ever
populated it.
On the rlformers side the single `norm_zero_centered_gamma` drives all of those too
(`transformer.py:3618,3634,3710,3715,5468`). It is one checkpoint property, not two.
So the two-field version leaves a checkpoint with `norm_zero_centered_gamma: true` and
affine norms computing `rms_norm((weight + 1) * x)` in `post_ffn_norm` and
`rms_norm(weight * x)` in the other five — the same silent wrong-activation bug,
relocated rather than fixed. `f1129404401` reaches `inf` under either version only
because its other norms happen to be unaffected; that is a property of that checkpoint,
not of the fix.
This version instead populates the existing field from the existing checkpoint key and
passes it into `RMSNormWithInputScale`. One name, one source of truth, six norms covered
instead of one, and no new `ModelArgs` surface: 10 files instead of 12.
It also removes the reason `patch_rms_unit_offset` exists in `utils/omni_patches.py:76`,
whose own docstring says ET "silently crushes activations" because "the omni
params->ModelArgs conversion never sets it". That monkeypatch is left in place here —
deleting it belongs in its own change.
Two things a reviewer should decide, both raised by minguo and NOT resolved here:
- rlformers gates the **final** norm on `output_norm_gain_center_type` ("one" -> +1,
"zero" -> no shift), independent of the per-layer flag. ET has one flag covering the
final norm too, so mapping them 1:1 is a deliberate simplification that needs sign-off.
- Widening the mapping changes numerics for any already-validated checkpoint whose
params.json carries `norm_zero_centered_gamma: true`. The conversion scripts bake the
+1 in and then set the flag false, so most are unaffected — but this wants an A/B on
the existing `backbone_cuda_test` checkpoints rather than an assumption.
Differential Revision: D118485339
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22573
Note: Links to docs will display an error until the docs builds have been completed. ❌ 3 Awaiting Approval, 3 New FailuresAs of commit 1301fdb with merge base 4fc2acc ( AWAITING APPROVAL - The following workflows need approval before CI can run:
NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
|
@telgamal-1 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D118485339. |
This PR needs a
|
Summary:
ALTERNATIVE to D117233461, implementing minguo's review suggestion. Same bug, one
fewer concept.
rlformers stores gamma offset by -1 when
norm_zero_centered_gammais set, so theeffective scale is
weight + 1. ET dropped the flag on the params -> ModelArgsconversion, so a checkpoint carrying it loaded without error and produced
systematically wrong activations.
D117233461 fixes that by adding a new
ModelArgs.norm_zero_centered_gammaandwiring it into
RMSNormWithInputScale. But ET already has a field meaning exactlythis:
ModelArgs.rms_norm_add_unit_offset(model_args.py:56), andRMSNorm.forwardalready implements it as
output * (1 + weight)(norm.py:37-41). It is already wiredinto
attention_normandffn_norm(llama_transformer.py:226,232),q_norm/k_norm(
attention.py:468,474) and the final norm (llama_transformer.py:345) — nothing everpopulated it.
On the rlformers side the single
norm_zero_centered_gammadrives all of those too(
transformer.py:3618,3634,3710,3715,5468). It is one checkpoint property, not two.So the two-field version leaves a checkpoint with
norm_zero_centered_gamma: trueandaffine norms computing
rms_norm((weight + 1) * x)inpost_ffn_normandrms_norm(weight * x)in the other five — the same silent wrong-activation bug,relocated rather than fixed.
f1129404401reachesinfunder either version onlybecause its other norms happen to be unaffected; that is a property of that checkpoint,
not of the fix.
This version instead populates the existing field from the existing checkpoint key and
passes it into
RMSNormWithInputScale. One name, one source of truth, six norms coveredinstead of one, and no new
ModelArgssurface: 10 files instead of 12.It also removes the reason
patch_rms_unit_offsetexists inutils/omni_patches.py:76,whose own docstring says ET "silently crushes activations" because "the omni
params->ModelArgs conversion never sets it". That monkeypatch is left in place here —
deleting it belongs in its own change.
Two things a reviewer should decide, both raised by minguo and NOT resolved here:
output_norm_gain_center_type("one" -> +1,"zero" -> no shift), independent of the per-layer flag. ET has one flag covering the
final norm too, so mapping them 1:1 is a deliberate simplification that needs sign-off.
params.json carries
norm_zero_centered_gamma: true. The conversion scripts bake the+1 in and then set the flag false, so most are unaffected — but this wants an A/B on
the existing
backbone_cuda_testcheckpoints rather than an assumption.Differential Revision: D118485339