[REQUEST] DeepCompile ZeRO-1: reduce-to-owner instead of AllReduce

Is your feature request related to a problem? Please describe.

Under DeepCompile with ZeRO stage 1, every parameter's full gradient is AllReduced even though only one rank keeps the result.

ZeRO-1 partitions optimizer state, and DeepCompile's ZeRO-1 init (deepspeed/compile/init_z1_and_2.py) takes that literally: it installs a real gradient buffer only on the rank whose optimizer partition holds the parameter, and torch.empty([0]) everywhere else. Z1CustomOpExecutor::flushReduceBucket (csrc/compile/z1.cpp) nevertheless issues ncclAllReduce for every parameter in the bucket, and the copy-out loop right below it drops the result wherever the buffer is empty:

if (grad_buf.numel() == 0) { continue; }

So world_size - 1 of world_size ranks pay for a reduction result they discard immediately. Instrumenting the flush confirms it — discarded = (ws-1)/ws, i.e. 87.5% at ws=8:

[dc-probe][r0] flushReduceBucket #3 dtype=BFloat16 tasks=118 reduced_elems=52428800 kept_elems=6553600 discarded=87.5%
pre_div=1 owner_tasks=0 owner_elems=0 (0.0%)

This matters on interconnects where the gradient reduction is exposed rather than hidden behind backward compute, which is the common case once the DP degree grows or the fabric is weaker than NVLink.

Describe the solution you'd like

For every parameter that lies entirely inside one optimizer partition, reduce to the single rank that keeps it — ncclReduce(root=owner) instead of ncclAllReduce — inside the existing bucket flush. Same reduction, same operand order per element, one destination instead of ws.

Two opt-in CompileConfig flags, both default off:

{
  "zero_optimization": { "stage": 1 },
  "compile": {
    "deepcompile": true,
    "reduce_to_owner": true
  }
}

• reduce_to_owner — enable the ncclReduce path.
• reduce_probe — print one line per bucket flush (first 20) reporting reduced_elems, kept_elems, discarded, and how much of the flush took the owner path. Lets a user size the win on their own model before turning the optimization on, and is useful on its own for debugging the reduce path.

Design points:

• Owner map costs nothing to build. optimizer.param_to_partition_ids[group][param_id] lists every partition holding a piece of the parameter. Exactly one entry ⇒ exactly one owner, and that entry is the owner's rank. initialize_gradient_partitioning_data_structures populates this for every partition on every rank, so the map needs no communication and is identical everywhere — required, since the map decides which collective each bucket entry uses and NCCL needs every rank to issue the same sequence inside a group. roots is therefore precomputed for the whole bucket before any collective is issued.
• Partial parameters keep AllReduce. A parameter straddling a partition boundary (at most ws-1 per param group) gets 2+ entries and is simply not declared; a failed map lookup means "AllReduce", not an error.
• Copy-out is unchanged and still correct. The set of ranks whose bucket slice ncclReduce leaves untouched is exactly the set that registered torch.empty([0]) and is skipped in copy-out.
• No native flag. An empty owner map already means "AllReduce everything", so with the config flag off Python declares no owners and the code path is bit-for-bit the original one.
• Summation semantics untouched. applyPreDivision and getReductionOp are not modified, so pre_div_reduce still divides by the world size before ncclSum.
• Parameter groups reduced over a different communicator are skipped (MoE expert groups, for example): their partition ids do not name ranks of DeepCompile's communicator, so they are excluded by a world-size + group_rank → global_rank comparison against engine.data_parallel_group.

Stage 2 logs that the flag is ignored; stage 3 does not use this reduction path at all. No extra memory, no new buffer, no change to the reduce bucket size, nothing to tune, and independent of the gradient dtype.

Describe alternatives you've considered

• ncclReduceScatter over the bucket. The natural ZeRO choice and the same traffic class, but it requires the bucket payload to be laid out contiguously by destination rank. DeepCompile assembles buckets in in-graph gradient-arrival order, one full-parameter task each, with arbitrary owner distribution — so this needs a bucket relayout pass. Per-parameter ncclReduce reaches the same traffic class with no relayout, at the cost of more collective calls inside a single group launch. ReduceScatter remains the better end state if someone wants to restructure bucketing.
• Quantized gradient reduction (fp8). Measured on the same benchmark and net negative on NVLink: all-to-all + all-gather 34.75 ms (+53%) and all-to-all → owner 25.41 ms (+12%) against a 22.72 ms bf16 AllReduce baseline. Quantization overhead is not repaid when the link is fast; that half is deliberately not proposed here.
• Reusing ZeRO++ zero_quantized_gradients. all_to_all_quant_reduce() (deepspeed/runtime/comm/coalesced_collectives.py) is INT4, its only caller is stage3.py, and it is unreachable from DeepCompile's in-graph reduce_grad path. Different quantization target and different stage.
• Status quo. Leaves the 87.5% discard in place. Acceptable while the reduction is fully hidden behind backward compute, which is exactly what reduce_probe lets a user check.

Additional context

Measured on 8x A800 (NVLink), 402.8M parameters / 806MB of bf16 gradients, communication-bound microbenchmark, median bucket-flush wall time:

Run-to-run spread of the reduce-to-owner number is 18.4–20.5 ms. The gain is bounded by how much of the flush reduce_probe reports as owner_elems, and by how much of the reduction is exposed rather than overlapped with backward compute. At init, rank 0 reports the qualifying ratio so the ceiling is visible before any step runs:

DeepCompile reduce_to_owner: 291/302 gradients reduce to one rank

A model made of a few very large tensors spread across all ranks will report a ratio near 0, and the flag will do nothing for it — which is the honest failure mode and the reason for the probe.

Numerics. The reduction is mathematically equivalent to the AllReduce it replaces: no rank's contribution is dropped and nothing is scaled differently. It is not bit-identical under real data parallelism, because NCCL's summation order changes — with world_size >= 4 the observed maximum relative difference is 9.0e-4 in bf16, the same class of nondeterminism as changing NCCL_ALGO. This is why both flags default to off and why acceptance should be band-level (loss curves over a few hundred steps) rather than value-level. Feeding identical input to every rank keeps the summation order unchanged and does reproduce exactly, which is useful as a correctness check.

Scope of the change. deepspeed/compile/config.py (the two flags), deepspeed/compile/init_z1_and_2.py (owner map + declaration), csrc/compile/z1.cpp (owner map, ncclReduce branch, probe), csrc/compile/init.cpp (two bindings), plus GPU-free unit tests for the owner map and a deepspeed/compile/reduce_to_owner.md write-up.

Follow-up. Z2CustomOpExecutor::flushReduceBucket has the same shape, and ZeRO-2 registers real buffers with offsets at init, so the same optimization applies there. Left out of this proposal because the measurement above is ZeRO-1 only.

I have a working branch for this and will open the PR against this issue.