Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

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

Open
  • Is something wrong?
#8,814 2 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
35/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
backend, distributed-systems

Research direction

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.

Written by the indexing model from the issue text.

Description

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.
Dominant language
Python
Stars
43.3k
Forks
5.1k
Avg merge
4d 17h
Merged PRs (30d)
63

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepspeedai/DeepSpeed

All issues in deepspeedai/DeepSpeed

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.