Gradient scale under sequence parallelism differs between ZeRO reduction paths

Author: alanhuangyooCreated Sep 28, 2026Updated Sep 28, 2026

With sequence_parallel_size > 1, ZeRO averages gradients over two different divisors, depending on the path:

world_size / sequence_parallel_size: ZeRO-1/2, and ZeRO-3's contiguous path (reduce_scatter: false, contiguous_gradients: true, bucket within reduce_bucket_size).
The whole sequence-data-parallel group: ZeRO-3's __avg_scatter_grads. That path is taken with reduce_scatter: true (the default), with contiguous_gradients: false, or for a bucket over reduce_bucket_size.

So the same model and loss get gradients that differ by a factor of sp depending on those settings, and in one case between parameters of the same model (below).

I measured this on master (df00e0bae), on H20 GPUs. The model is one linear layer; each rank feeds its own shard, the optimizer is SGD with lr=0, and the gradient is read with safe_get_full_grad. The table shows ZeRO's gradient divided by the sum of the per-shard gradients:

loss	ZeRO-1	ZeRO-2	ZeRO-3, reduce_scatter on	ZeRO-3, reduce_scatter off, contiguous_gradients on
each rank's shard loss, sp=2	1.0	1.0	0.5	1.0
all-gathered loss (ALST recipe), sp=2	2.0	2.0	1.0	2.0
each rank's shard loss, sp=4	1.0	1.0	0.25	1.0
all-gathered loss (ALST recipe), sp=4	4.0	4.0	1.0	4.0

ZeRO-1 and ZeRO-2 give the same numbers with reduce_scatter on or off. ZeRO-3 with both reduce_scatter and contiguous_gradients off matches the "reduce_scatter on" column (0.5 / 1.0 at sp=2).

Within one model. With ZeRO-3, reduce_scatter: false and contiguous_gradients: true, a parameter larger than reduce_bucket_size goes through __avg_scatter_grads while the rest take the contiguous path. I measured Linear(4, 4) with reduce_bucket_size: 10, sp=2 and each rank's shard loss: the 16-element weight gets 0.5x the summed shard gradient and the 4-element bias gets 1.0x, in the same step. With the default reduce_bucket_size of 5e8 this hits any single parameter above 5e8 elements, e.g. an embedding with a 152k vocab and hidden size 4096 (6.2e8). No loss convention makes that consistent, so at least __avg_scatter_grads and the contiguous path should agree.

Where it comes from. #4530 changed the averaging divisor to world_size / sequence_parallel_size in two places:

ZeRO-1/2: average_tensor, gradient_reduction_w_predivide and allreduce_bucket.
ZeRO-3: __avg_scatter_contiguous_grads and allreduce_bucket.

ZeRO-3's __avg_scatter_grads wasn't changed. It still averages inside reduce_scatter_coalesced, which divides by the full group size.

Why ALST is only right on that one path. The ALST recipe is used by the tutorial, tests/unit/ulysses_alst, and the HF Trainer (deepspeed_sp_compute_loss in transformers, which is the loss the Trainer backprops). It all-gathers the shard losses with the differentiable all_gather, whose backward sums over the SP group, so each rank's local gradient is sp times its share. That gives the right gradient on the __avg_scatter_grads path (ZeRO-3's default) and sp times too much on the others. At sp=2, I stepped SGD with lr=1 and checked the step it actually took:

Without clipping, ZeRO-2 steps 2.0x the true gradient and reports 2.0x the grad norm; ZeRO-3 gives 1.0x for both.
With gradient_clipping set to 1.5x the true norm, ZeRO-2 still clips (its step comes out 1.5x), while ZeRO-3 doesn't.

Why the test doesn't catch it. TestUlyssesSPHF does run ZeRO-2, but the grads it compares are at most 5.6e-4, which is under its atol=1e-3. On master its ZeRO-2 case passes with grad_b ≈ 2.0005 * grad_a; the ZeRO-3 case gives 1.0003.

Which convention should hold? If the divisor should be world_size / sp everywhere, as ZeRO-1/2 and the ZeRO-3 contiguous path already do, two changes are needed:

ZeRO-3's __avg_scatter_grads needs the same factor.
The ALST recipe needs to stop multiplying by sp, e.g. each rank backprops loss * good_tokens / total_good_tokens and all-gathers only for the reported value.

These have to land together, or ALST + ZeRO-3 changes scale. I'm happy to send a PR once it's clear which way you want it, including a tighter tolerance for the ZeRO-2 case in that test.

Repro (torchrun --nproc_per_node 2 or 4)

python

import os
import torch
import torch.distributed.nn.functional as dist_fn
import deepspeed
from deepspeed import comm as dist
from deepspeed.runtime.sequence_parallel import parallel_state_sp as mpu
from deepspeed.utils import safe_get_full_grad

W0 = torch.arange(16, dtype=torch.float32).view(4, 4) / 16


def run(stage, reduce_scatter, loss_kind):
    model = torch.nn.Linear(4, 4, bias=False)
    with torch.no_grad():
        model.weight.copy_(W0)
    config = {
        "train_micro_batch_size_per_gpu": 1,
        "gradient_clipping": 0.0,
        "zero_allow_untested_optimizer": True,
        "sequence_parallel_size": dist.get_world_size(),
        "optimizer": {"type": "SGD", "params": {"lr": 0.0}},
        "zero_optimization": {"stage": stage, "reduce_scatter": reduce_scatter},
    }
    engine, _, _, _ = deepspeed.initialize(model=model, config=config, mpu=mpu)
    x = torch.full((1, 4), float(dist.get_rank() + 1), device=engine.device)
    local = engine(x).pow(2).sum()
    loss = local if loss_kind == "shard" else sum(dist_fn.all_gather(local, group=mpu.get_sequence_parallel_group()))
    engine.backward(loss)
    grad = safe_get_full_grad(engine.module.weight).float().cpu()
    engine.destroy()
    return grad


def reference():
    total = torch.zeros(4, 4)
    for r in range(int(os.environ["WORLD_SIZE"])):
        w = W0.clone().requires_grad_(True)
        (torch.full((1, 4), float(r + 1)) @ w.t()).pow(2).sum().backward()
        total += w.grad
    return total


deepspeed.init_distributed()
mpu.initialize_sequence_parallel(dist.get_world_size())
ref = reference().flatten()
for loss_kind in ("shard", "all_gather"):
    for stage in (1, 2, 3):
        for rs in (True, False):
            g = run(stage, rs, loss_kind).flatten()
            if dist.get_rank() == 0:
                print(f"{loss_kind:10s} stage{stage} reduce_scatter={rs!s:5s} {(g @ ref / (ref @ ref)).item():.4f}")

Repro for the within-one-model case (torchrun --nproc_per_node 2)

python

import os

import torch

import deepspeed
from deepspeed import comm as dist
from deepspeed.runtime.sequence_parallel import parallel_state_sp as mpu
from deepspeed.utils import safe_get_full_grad

W0 = torch.arange(16, dtype=torch.float32).view(4, 4) / 16
B0 = torch.arange(4, dtype=torch.float32) / 4


def reference():
    gw, gb = torch.zeros(4, 4), torch.zeros(4)
    for r in range(int(os.environ["WORLD_SIZE"])):
        lin = torch.nn.Linear(4, 4)
        with torch.no_grad():
            lin.weight.copy_(W0)
            lin.bias.copy_(B0)
        lin(torch.full((1, 4), float(r + 1))).pow(2).sum().backward()
        gw += lin.weight.grad
        gb += lin.bias.grad
    return gw.flatten(), gb


deepspeed.init_distributed()
mpu.initialize_sequence_parallel(dist.get_world_size())
model = torch.nn.Linear(4, 4)
with torch.no_grad():
    model.weight.copy_(W0)
    model.bias.copy_(B0)
config = {
    "train_micro_batch_size_per_gpu": 1,
    "gradient_clipping": 0.0,
    "zero_allow_untested_optimizer": True,
    "sequence_parallel_size": dist.get_world_size(),
    "optimizer": {"type": "SGD", "params": {"lr": 0.0}},
    "zero_optimization": {
        "stage": 3,
        "reduce_scatter": False,
        "contiguous_gradients": True,
        "reduce_bucket_size": 10,  # the 16-element weight does not fit, the 4-element bias does
    },
}
engine, _, _, _ = deepspeed.initialize(model=model, config=config, mpu=mpu)
x = torch.full((1, 4), float(dist.get_rank() + 1), device=engine.device)
engine.backward(engine(x).pow(2).sum())  # each rank's own shard loss
gw = safe_get_full_grad(engine.module.weight).float().cpu().flatten()
gb = safe_get_full_grad(engine.module.bias).float().cpu()
rw, rb = reference()
if dist.get_rank() == 0:
    print(f"weight (16 elements) grad / sum of shard grads = {(gw @ rw / (rw @ rw)).item():.4f}")
    print(f"bias   (4 elements)  grad / sum of shard grads = {(gb @ rb / (rb @ rb)).item():.4f}")

Source: deepspeedai/DeepSpeed

View original on GitHub View discussion on GitHub