pytorch / pytorch/pytorch

[DTensor] Avoid forced redistribution in layer_norm strategy

Open
#174,276 2 comments 1 reaction 0 assignees View on GitHub
bot-triaged enhancement module: dtensor module: performance oncall: distributed oncall: distributed parallelisms ptd-bot-triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🚀 The enhancement, motivation and pitch

Layer norm strategy forces input/weight/bias to replicate on normalized dimensions before calling the op. This causes unnecessary all-gather when input is sharded on those dims.

Current behavior https://github.com/pytorch/pytorch/blob/07671f333e3495a7b2307a6f43e583a8e4f88106/torch/distributed/tensor/_ops/_math_ops.py#L1003-L1051
- Input: replicated on dims >= axis
- Weight/bias: fully replicated

For a `[32, 2048, 4096]` tensor sharded on hidden dim across 8 GPUs, this moves 512MB. Decomposing into local stats + all-reduce would move ~0.5MB.

can we decompose `layer_norm` into primitives that work with partial reductions? i.e.,
1. compute local sum(x), sum(x²) → Partial placement
2. all-reduce statistics only
3. normalize locally with global mean/rstd
4. element-wise weight/bias (stays sharded)

cc @jerryzh168 @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @tianyu-l @XilunWu @SherlockNoMad @ppwwyyxx @H-Huang

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.