NVIDIA / NVIDIA/Megatron-LM

[ENHANCEMENT] Accept ProcessGroupCollection in inference communication_utils broadcast helpers

Open
#7,119 1 comment 0 reactions 0 assignees View on GitHub
community-request waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

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

Megatron Core is migrating production code off direct `parallel_state.get_*_group()` reads toward caller-supplied `ProcessGroupCollection` / `ProcessGroup` objects. Several helpers in `megatron/core/inference/communication_utils.py` still unconditionally read MPU globals:

- `broadcast_tensor`
- `broadcast_list` / `broadcast_int_list` / `broadcast_float_list`

The P2P helpers already accept an optional `pp_group`. The remaining broadcast helpers used by the text-generation server do not.

Tag @NVIDIA/mcore-oncall

**Describe the solution you'd like**

- Accept an explicit `ProcessGroup` (or `ProcessGroupCollection`) on the broadcast helpers.
- Keep a `parallel_state` fallback only as a compatibility path, with an explicit comment that it is a migration fallback.
- Thread the new argument through in-repo callers that already have a process group available.
- Extend `tests/unit_tests/inference/test_communication_utils.py`.

**Describe alternatives you've considered**

Leaving the helpers on MPU globals forever. That works today but blocks callers that construct custom process groups (multimodal / HyperCommGrid inference).

**Additional context**

Scoped to `communication_utils.py` and its existing callers. Not a repo-wide `parallel_state` sweep.

Contributor guide

Open the contributing guide

Research direction

Start in megatron/core/inference/communication_utils.py and inspect the broadcast helpers and their existing callers, then read tests/unit_tests/inference/test_communication_utils.py. Trace how the P2P helpers receive process groups and identify callers that already have one available. Done means the broadcast helpers accept explicit ProcessGroup or ProcessGroupCollection values, retain a documented compatibility fallback, and have passing tests covering the new behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.