Skip to content

[ET-VK] native_layer_norm rejects a missing weight, so affine-free F.layer_norm cannot run#22504

Description

@msluszniak

馃悰 Describe the bug

add_native_layer_norm_node throws when weight is None:

if (graph.val_is_none(weight_data)) {
  VK_THROW("native_layer_norm requires weight to be non-None");
}

The neighbouring code already handles a missing bias, synthesizing a zero-bias TensorRef for nn.LayerNorm(bias=False). A missing weight is the symmetric case and is just as ordinary: F.layer_norm(x, (channels,), eps=...) with no affine parameters is a normal call, used whenever a module applies its own scale and shift afterwards.

Concrete case: kokoro's AdaLayerNorm normalizes with no affine parameters and then applies a style-conditioned (1 + gamma) * x + beta. Lowering that model to Vulkan produces a .pte that aborts at prepack:

Exception raised from add_native_layer_norm_node at
backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp:95:
native_layer_norm requires weight to be non-None

The partitioner claims the op, so there is no way to avoid this from the export side other than blocklisting native_layer_norm entirely, which splits the graph.

A unit weight and a zero bias reproduce out = (x - mean) * rstd exactly, so synthesizing the missing parameter is enough. When there is no weight tensor to take a shape from, normalized_shape supplies it.

Versions

Reproduces on main (a6b115b).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions