[ET-VK] Synthesize the layer_norm affine parameters that are absent - #22506
[ET-VK] Synthesize the layer_norm affine parameters that are absent#22506msluszniak wants to merge 2 commits into
Conversation
The layer_norm shader reads a weight and a bias binding unconditionally, and add_native_layer_norm_node already synthesized a zero bias for nn.LayerNorm(bias=False). A missing weight was still a hard error, so any model that calls F.layer_norm with no affine parameters at all aborts at prepack with "native_layer_norm requires weight to be non-None". That call is not unusual: kokoro's AdaLayerNorm normalizes with no affine parameters and applies its own style-conditioned scale and shift afterwards. Synthesize a unit weight the same way, sized from normalized_shape when there is no weight tensor to take a shape from. A unit weight and a zero bias reproduce out = (x - mean) * rstd exactly. The two paths now share one helper.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22506
Note: Links to docs will display an error until the docs builds have been completed.
|
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
SS-JIA
left a comment
SS-JIA
left a comment
There was a problem hiding this comment.
AI-assisted review. Inline findings below.
| is_buffer ? layer_norm_buffer_gwg : default_pick_gwg, | ||
| is_buffer ? pick_required_lwg : default_pick_lwg, | ||
| is_buffer ? layer_norm_buffer_global_wg_size | ||
| : default_pick_global_wg_size, |
There was a problem hiding this comment.
[High] This does not compile against the current dispatch API. default_pick_global_wg_size and default_pick_local_wg_size do not exist; the callbacks are default_pick_gwg and default_pick_lwg, and custom pickers must use the current GlobalWorkGrid/LocalWorkGroup signatures. Please restore those names and types while retaining the affine synthesis.
AI reviewer note: this comment was generated by an AI reviewer.
There was a problem hiding this comment.
@msluszniak for context, the global work grid / local work group determination API was updated recently, so this comment is valid. Please update the functions to return GlobalWorkGrid / LocalWorkGroup, and also rename functions to match the *_pick_gwg and *_pick_lwg pattern. Thanks in advance!
There was a problem hiding this comment.
Done. layer_norm_buffer_gwg is back, returning a GlobalWorkGrid with its required LocalWorkGroup, the extra local-size picker is gone since pick_required_lwg covers it, and the dispatch uses default_pick_gwg / default_pick_lwg again. The affine synthesis is unchanged.
One thing while in there: it was including runtime/core/portable_type/half.h just to write 1.0 in fp16. Since ones and zeros are the only values it ever needs, that is now a single bit pattern and the file no longer reaches outside the backend for a float16 type, matching what StagingBuffer.cpp already does.
The affine synthesis was written against the dispatch API as it stood before the GlobalWorkGrid / LocalWorkGroup rework, so it renamed layer_norm_buffer_gwg back to layer_norm_buffer_global_wg_size, added a local size picker that the required-lwg path already covers, and called default_pick_global_wg_size / default_pick_local_wg_size, none of which exist any more. Restore the picker that is there, returning a GlobalWorkGrid with its required local work group, and go back to default_pick_gwg / pick_required_lwg / default_pick_lwg. The synthesis itself is unchanged apart from dropping the include of portable_type/half.h: the ones and zeros it writes are the only two values it ever needs, so the fp16 case is one bit pattern and the file no longer reaches outside the backend for a float16 type. The backend avoids that dependency elsewhere too, in StagingBuffer.cpp.
3 participants
Fixes #22504.
add_native_layer_norm_nodealready synthesizes a zero bias fornn.LayerNorm(bias=False), but still throws on a missing weight.F.layer_normcalled with no affine parameters at all is an ordinary call, used whenever a module applies its own scale and shift afterwards, and it currently makes the model abort at prepack.This synthesizes a unit weight the same way, sizing it from
normalized_shapewhen there is no weight tensor to take a shape from. A unit weight and a zero bias reproduceout = (x - mean) * rstdexactly. The two paths now share one helper.Verified on device: kokoro's duration predictor, whose
AdaLayerNormnormalizes with no affine parameters, aborts onmainand runs to completion with this change on an Adreno 840.cases.pygains threenative_layer_normcases covering weight-only, bias-only and neither.