#6437·ColossalAI

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

Author: tandedeCreated Aug 20, 2026Updated Aug 20, 2026

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

python
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.