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

[tandede]

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.