Applying LayerNorm After TEColumnParallelLinear in Tensor Parallel Setup

[ftgreat]

Hello,

I am working on a two-layer MLP setup within a tensor parallelism context (e.g., tp=2), which consists of a TEColumnParallelLinear followed by a TERowParallelLinear.
I’m looking to perform a LayerNorm operation right after the output from TEColumnParallelLinear, and importantly, I want each tensor parallel split (tp slice) to have its own independent LayerNorm.

Could you provide any suggestions on whether there is an existing module that supports this, or how one might go about adding this functionality?
Thank you for your assistance!

[timmoon10]

You could accomplish this with a Linear module followed by a LayerNormLinear:

layer1 = te.Linear(
    in_features,
    layer1_features,
    tp_group=tensor_parallel_group,
    tp_size=tensor_parallel_size,
    parallel_mode="column",
)
layer2 = te.LayerNormLinear(
    layer1_features,
    layer2_features,
    tp_group=tensor_parallel_group,
    tp_size=tensor_parallel_size,
    parallel_mode="row",
)
y = layer2(layer1(x))

If DDP/FSDP is configured as usual (i.e. doing grad reductions over the data-parallel group), then the LayerNorms will train independently for each tensor parallel group.

If you need more flexibility, you could try using the experimental operation-based API:

mlp = te.ops.Sequential(
    te.ops.Linear(
        in_features,
        layer1_features,
        tensor_parallel_mode="column",
        tensor_parallel_group=tensor_parallel_group,
    ),
    te.ops.LayerNorm(layer1_features),
    te.ops.Linear(
        layer1_features,
        layer2_features,
        tensor_parallel_mode="row",
        tensor_parallel_group=tensor_parallel_group,
    ),
)
y = mlp(x)

This requires #1033.

[ftgreat]

Thanks a lot.