Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

[RFC] Skip collectives on absent or single-rank groups with deepspeed.comm.pg_utils.is_parallel

Aperta
  • Qualcosa non va?
#8,814 2 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
35/100
Tipo di issue
Funzionalità
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
backend, distributed-systems

Direzione di ricerca

Start with deepspeed/comm/comm.py and its existing collective helpers, then read the tests for deepspeed.comm and the guarded call sites in deepspeed/sequence/cross_entropy.py. Implement and test the proposed synchronous helper so absent and single-rank groups skip the reduction while multi-rank groups preserve existing behavior. Done means the helper is exported through deepspeed.comm and the cross-entropy callers use it; the naming and scope still need maintainer feedback.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

Update (2026-10-16): the proposal changed during review of #8815. Instead of an exported all_reduce_if_parallel wrapper, DeepSpeed gets a single predicate, deepspeed.comm.pg_utils.is_parallel(group), which guards every collective. The sections below reflect the current design; see the status table for progress.

Motivation

Following @delock's suggestion in #8791: many modules guard collectives with the same hand-written check, so that code written for an optional parallel group also runs without it:

if tp_group is not None and dist.get_world_size(tp_group) > 1:
    dist.all_reduce(tensor, op=..., group=tp_group)

This pattern (or a close variant) appeared in sequence parallel / vocab-parallel ops, MoE, inference v1 kernels, ZeRO model-parallel reductions, AutoTP layers and runtime norm helpers. #8791 added yet another local copy (_tp_active / _all_reduce in deepspeed/sequence/vocab_parallel_ops.py).

The copies differ in subtle ways (some only check group is None, some only check the size), and they all rely on a convention that conflicts with deepspeed.comm: in these helpers group=None means "this parallel dimension is absent", while dist.all_reduce(group=None) reduces over the world group. Dropping the guard therefore silently turns "no reduction" into a world reduction, which is wrong whenever DP > 1.

Proposal (current)

Add deepspeed/comm/pg_utils.py with one predicate:

def is_parallel(group):
    """Return whether ``group`` spans more than one rank. ``None`` marks an absent parallel dimension."""
    return group is not None and dist.get_world_size(group) > 1

Call sites keep their own collective:

from deepspeed.comm import pg_utils

if pg_utils.is_parallel(group):
    dist.all_reduce(tensor, group=group)

Design points:

  • One predicate for every primitive. Sync and async all-reduce, all-gather, inference all-reduce, etc. all use the same guard, so there is no per-collective wrapper and no question of what a skipped async_op=True call should return.
  • None means "no parallelism", documented and tested. It never touches the process group, so it works before torch.distributed is initialized.
  • Not star-exported from deepspeed.comm. Callers write pg_utils.is_parallel, so the two meanings of None (absent dimension vs. world group) stay visibly apart.
  • Non-members skip. For a group that excludes the current rank, get_world_size returns -1, and > 1 skips it.
  • Hot paths can cache it. AutoTP layers compute is_tp_parallel = is_parallel(mp_group) once at construction, so forward/backward pay no extra cost.
  • The backend collectives (cdb.*, timed_op, comms logger) are unchanged.

Why not skip inside dist.all_reduce itself? That would change behavior for every caller (group=None = world, async handles, comms logging), so an opt-in predicate is safer.

Status

Step Scope Status
PR-1 pg_utils.is_parallel + tests; migrate sequence (vocab_parallel_ops.py, cross_entropy.py, linear_cross_entropy.py), MoE, inference v1, ZeRO MP reductions, AutoTP layers, runtime norm helpers #8815 (originally planned as PR-1 + PR-2)
Out of scope Guards on a separately supplied size (len(partition_sizes), a tp_world_size argument), mpu-based world sizes, sites where None means the world group (kv_cache warm-up, moe_inference._alltoall), AutoEP folding correction kept as is, listed in #8815
Performance (separate, optional) At DP=1 the engine / ZeRO-1/2 gradient paths still flatten, all-reduce over a single rank and copy back, and ZeRO-2 overlap keeps an IPG double buffer. Caller-level short-circuit, not this predicate: skip the stage-0 reduction and disable the IPG double buffer when DP=1. Measured on 1× Arc Pro B70: stage 0 saves ~13 ms/step (4.3% at 8K tokens, 12.8% at 1K), ZeRO-1/2 average_tensor skip ~4 ms/step (1–3%), double buffer costs +0.94 GiB; losses bit-identical planned, small PR

Related fix found during the migration (independent of this RFC): #8837 — dist.all_reduce(x.contiguous()) in RowParallel / ColumnParallel / AsyncColumnParallel discarded the reduced result for non-contiguous tensors.

Questions

  • Naming: resolved as pg_utils.is_parallel(group) in #8815 review. Objections welcome.
  • Feedback on the out-of-scope list is welcome.
Lingua principale
Python
Stelle
43.3k
Fork
5.1k
Merge medio
4g 7h
PR unite (30g)
66

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di deepspeedai/DeepSpeed

Tutte le issue di deepspeedai/DeepSpeed

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.