[AutoParallel]Revise Infermeta of LayerNorm for Sequence-Data Hybrid Parallelism

[JZ-LIANG]

PR types

Function optimization

PR changes

Others

Description

Pcard-76459

The Current LayerNorm Implement would flatten the broadcast axes (axis before "begin_norm_axis") of LayerNorm, which would hinder broadcast axes to be sharded by different Mesh dimensions.

In Sequence-Data Hybrid Parallelism, we need to shard both batch and sequence axes (broadcast axes) of LayerNorm input to get the best Performance.

Therefore we plan to remove the "flatten logic" in LayerNorm Op.

Before PR:
all axes before 『begin_norm_axis』will be flatten into one axis for var and std of LN.

After PR:
the shape of mean and var will keep the same as the shape of input 『begin_norm_axis』.

# pseudocode
in = paddle.randn(shape = [8, 1024, 768])
var, mean, out = layer_norm(in, begin_norm_axis=2)

# before PR
out.shape == [8, 1024, 768]
var.shape == [8192]
mean.shape == [8192]

# after PR
out.shape == [8, 1024, 768]
var.shape == [8, 1024]
mean.shape == [8, 1024]

The Logic modified in this PR:

• Layernorm InferMeta
• Layernorm InferSpmd
• Layernorm Prim composite rule
• Layernorm CPU kernel
• Layernorm XPU & MLKDNN Unitest

[pkuzyc]

Are there some corresponding cases in unit test? If so, unit test should also be revised.

[xiaoguoguo626807]

LGTM for composite rules

[zhiqiu]

LGTM

[chenwhql]

LGTM

[pkuzyc]

LGTM for spmd rules