Summary
gpt_oss.torch.utils.init_distributed() uses global RANK as the CUDA device index:
rank = int(os.environ.get("RANK", 0))
...
torch.cuda.set_device(rank)
device = torch.device(f"cuda:{rank}")That happens to work for single-node launches where global and local ranks are identical. Under standard multi-node torchrun, RANK is global across all workers while LOCAL_RANK identifies the GPU/process on the current node.
Reproducer
On node 2 of a two-node, four-GPU-per-node launch, a worker can have:
RANK=4
LOCAL_RANK=0
WORLD_SIZE=8
The current code attempts to bind cuda:4 on a host whose local devices are indexed 0-3.
Impact
The reference Torch generation/chat path cannot use standard multi-node distributed launches correctly and can fail with an invalid CUDA device ordinal on every node after the first.
Proposed resolution
Keep global RANK for dist.init_process_group() and rank-aware output suppression, but use LOCAL_RANK for torch.cuda.set_device() and the returned CUDA device. Fall back to RANK when LOCAL_RANK is not present to preserve existing custom/single-node launch behavior.
Add a unit regression with different global and local ranks that verifies process-group initialization receives the global rank while CUDA device selection uses the local rank.
Summary
gpt_oss.torch.utils.init_distributed()uses globalRANKas the CUDA device index:That happens to work for single-node launches where global and local ranks are identical. Under standard multi-node
torchrun,RANKis global across all workers whileLOCAL_RANKidentifies the GPU/process on the current node.Reproducer
On node 2 of a two-node, four-GPU-per-node launch, a worker can have:
The current code attempts to bind
cuda:4on a host whose local devices are indexed 0-3.Impact
The reference Torch generation/chat path cannot use standard multi-node distributed launches correctly and can fail with an invalid CUDA device ordinal on every node after the first.
Proposed resolution
Keep global
RANKfordist.init_process_group()and rank-aware output suppression, but useLOCAL_RANKfortorch.cuda.set_device()and the returned CUDA device. Fall back toRANKwhenLOCAL_RANKis not present to preserve existing custom/single-node launch behavior.Add a unit regression with different global and local ranks that verifies process-group initialization receives the global rank while CUDA device selection uses the local rank.