Skip to content

Support custom process groups in distributed streamer - #166

Draft
Broduker wants to merge 1 commit into
dsx-ai-factory:masterfrom
Broduker:process-group-support
Draft

Broduker wants to merge 1 commit into
dsx-ai-factory:masterfrom
Broduker:process-group-support

Conversation

@Broduker

@Broduker Broduker commented Aug 6, 2026

Copy link
Copy Markdown

Summary

add an optional custom process group to SafetensorsStreamer and DistributedStreamer;
scope distributed rank discovery, backend checks, partitioning, barriers, and broadcasts to the provided group;
treat a provided process group as borrowed and leave its lifecycle to the caller;
add document explaining how to use it.

When no process group is provided, the existing global/host-local group selection behavior is preserved.

Fixes #165

Tests

  • added unit tests for process-group-scoped rank discovery, direct group reuse, and caller ownership;
  • added a distributed test that streams through a borrowed all-rank process group;
  • added a distributed test where only a subset process group participates.

@Broduker
Broduker force-pushed the process-group-support branch from b197ca5 to cecd028 Compare August 6, 2026 09:15
@Broduker Broduker changed the title Support caller-provided process groups in distributed streamer Support custom process groups in distributed streamer Aug 6, 2026
@noa-neria

Copy link
Copy Markdown
Collaborator

@Broduker thank you! Happy to review
See my suggestions regarding the design in #165

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Distributed streaming cannot be scoped to model-parallel process groups

2 participants