[BUG]: FP8 intranode detection ignores OpenMPI and Slurm local size values

Open Beginner friendly
#6,437 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
75/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Quiet
Tech stack
python, pytorch
Domain
distributed-systems

Research direction

Start at process_group_is_intranode() in colossalai.quantization.fp8 and run the minimal reproduction with OpenMPI or Slurm variables set without LOCAL_WORLD_SIZE. The fix is done when the detected variable supplies the local size, the torchrun/OpenMPI/Slurm priority remains unchanged, and both cases determine intranode status without KeyError.

Written by the indexing model from the issue text.

Description

Bug

process_group_is_intranode() checks three environment variables for the local world size, but it always reads LOCAL_WORLD_SIZE after finding a match. When a job exposes only the OpenMPI or Slurm variable, FP8 collective setup fails with KeyError: 'LOCAL_WORLD_SIZE' before it can determine whether the process group is intranode.

Minimal reproduction
import os
from unittest.mock import patch

from colossalai.quantization import fp8

os.environ.pop("LOCAL_WORLD_SIZE", None)
os.environ["OMPI_COMM_WORLD_LOCAL_SIZE"] = "4"

with patch.object(fp8.dist, "get_process_group_ranks", return_value=[0, 1, 2, 3]):
    print(fp8.process_group_is_intranode(object()))

The same failure occurs when only SLURM_TASKS_PER_NODE=4 is set.

Expected behavior

The function should read the value from the environment variable that was detected. If more than one variable is present, the existing order should remain the priority: torchrun, then OpenMPI, then Slurm.

Environment
  • ColossalAI main at 4f9953b
  • Python 3.11
  • PyTorch 2.5.1
  • macOS arm64 (the failure occurs before any GPU operation)
Checks
  • I searched the existing issues.
  • I reproduced this on the latest main branch.
  • I included a minimal reproduction.
Dominant language
Python
Stars
41.4k
Forks
4.5k
No merged PRs in 30d

Contributor guide

Open the contributing guide

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 hpcaitech/ColossalAI

All issues in hpcaitech/ColossalAI

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.