[DTensor] Avoid forced redistribution in layer_norm strategy
- 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
Assessment
This issue has not been assessed yet.