Drop the unused NCCL process group from the distributed workers#224
Open
clemsgrs wants to merge 2 commits into
Open
Drop the unused NCCL process group from the distributed workers#224clemsgrs wants to merge 2 commits into
clemsgrs wants to merge 2 commits into
Conversation
The distributed workers built an NCCL process group that ran zero live
collectives: every cross-rank exchange happens through the filesystem
(shard files / final artifacts + the JSONL progress bridge). The group
existed only so get_global_rank()/get_global_size() could call
dist.get_rank()/dist.get_world_size() — values torchrun already exports
as RANK / WORLD_SIZE. It cost an NCCL init, a per-rank comms buffer, and
a 14-day-timeout hang surface for no benefit.
distributed.enable() now records the torchrun-exported topology
(RANK / WORLD_SIZE / LOCAL_RANK / LOCAL_WORLD_SIZE) — or single-process
defaults when launched bare — into module state and keeps only the
torch.cuda.set_device() binding; it no longer calls init_process_group.
is_enabled()/get_global_* read that module state instead of a live
process group. The now-dead destroy_process_group() teardown and the
torch.distributed import are removed from pipeline_worker and
direct_embed_worker (their encode logic is unchanged).
Also removes the dead dist.barrier() from hf_login() (and the now-unused
is_enabled_and_multiple_gpus helper). The barrier only ever ran in the
single-process parent, but was a latent hang: had it run multi-rank with
HF_TOKEN unset, non-main ranks would return early while rank 0 blocked on
the barrier alone.
Validated with a live 2-rank torchrun launch: enable() reads the env,
binds cuda:{local_rank}, and initializes no process group.
Review follow-ups on the process-group removal (#219): - slide2vec.distributed no longer reads or exports MASTER_ADDR / MASTER_PORT and loses _get_available_port() plus the socket import: those only located the env:// rendezvous for the process group that no longer exists. enable() still records and exports the four torchrun topology variables, and the bare-launch error now names the real failure (no torchrun env, no CUDA device) instead of "PyTorch distributed". - Reword the get_global_rank/get_local_rank/get_local_size docstrings that still spoke of a "process group". - Remove the five dead 'import torch.distributed as dist' + is_available/is_initialized stubs in the worker tests; they only neutralized the destroy_process_group teardown the workers no longer have. The AST contract test now also locks out the socket import. Validated: full non-heavy suite green (500 passed); live 2-rank torchrun against the working tree reports the correct topology, binds cuda:{0,1}, and initializes no process group; a bare single-process enable() yields the rank-0/world-1 defaults.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The distributed workers built an NCCL process group with zero live collectives — every cross-rank exchange goes through the filesystem (shard files / final artifacts + the JSONL progress bridge). The group existed only so
get_global_rank()/get_global_size()could calldist.get_rank()/dist.get_world_size(), values torchrun already exports asRANK/WORLD_SIZE. Cost: an NCCL init, a per-rank comms buffer, and a 14-day-timeout hang surface — for no benefit.Changes
slide2vec/distributed/__init__.py—enable()records the torchrun-exported topology (RANK/WORLD_SIZE/LOCAL_RANK/LOCAL_WORLD_SIZE), or single-process defaults when launched bare, into module state and keeps only thetorch.cuda.set_device()binding. It no longer callsinit_process_group(theallow_nccl_timeoutknob is gone with it).is_enabled()/get_global_*read that module state instead of a live process group. Drops thetorch.distributed,datetime, and unusedrandomimports.pipeline_worker.py/direct_embed_worker.py— remove the now-deaddist.destroy_process_group()teardown and thetorch.distributedimport. Their encode logic is byte-identical modulo indentation (verified programmatically).hf_login()— removes the deaddist.barrier()(and the now-unusedis_enabled_and_multiple_gpushelper). The barrier only ever ran in the single-process parent, but was a latent hang: had it run multi-rank withHF_TOKENunset, non-main ranks wouldreturnearly while rank 0 blocked on the barrier alone.Validation
init_process_group/destroy_process_group/barrier/ collectives in the module (AST),set_devicesurvives, notorch.distributed/datetimeimport, andget_global_rank/get_global_size/is_enabledread module state rather than a process group.torchrunlaunch (--nproc_per_node=2, 2× GPU):enable()reads the env, bindscuda:{local_rank}, and initializes no process group (dist.is_initialized() == False).Closes #219