Allow ROCm arch override

[ailuntz]

Fixes #1444.\n\nSupport BNB_ROCM_GPU_ARCH override when rocminfo is unavailable.

[Abdennacer-Badaoui]

Thanks @ailuntz.
A few things need adjustment: First, the override should be a fallback after rocminfo fails, not a priority override that runs before it; for most users, auto-detection is more reliable and shouldn't be bypassed. Second, use only BNB_ROCM_GPU_ARCH and drop BNB_ROCM_ARCH, that name is already used as a CMake build-time variable and reusing it at runtime with different semantics is confusing. Third, add some basic validation on the value (e.g. check it matches gfx[a-zA-Z0-9]+) so garbage input doesn't silently propagate. And finally, the existing error message should mention BNB_ROCM_GPU_ARCH as a workaround so users know what to do when rocminfo fails (ask users to install rocminfo, if it's not possible then ask them to set the env variable)

[ailuntz]

Thanks for the feedback. I reviewed the code and agree the override should be a fallback only when rocminfo fails, and keeping both env vars is unnecessary. I've updated the PR to use only BNB_ROCM_GPU_ARCH as a fallback and pushed the changes.

[Abdennacer-Badaoui]

Thanks. Can you update the existing error message to mention BNB_ROCM_GPU_ARCH as a workaround so users know what to do when rocminfo fails (ask users to install rocminfo first, if it's not possible then ask them to set the env variable)

[ailuntz]

Updated this PR based on the feedback.

Changes in the latest push:

• keep BNB_ROCM_GPU_ARCH as a fallback only after rocminfo detection fails
• remove the extra runtime env name and use only BNB_ROCM_GPU_ARCH
• validate override values against the expected gfx... format before using them
• update the detection warning to tell users to install rocminfo first, and otherwise use BNB_ROCM_GPU_ARCH as the workaround
• add focused tests for detected value precedence, valid fallback override, and invalid override rejection
• align the ROCm installation docs with the same guidance

[Abdennacer-Badaoui]

Thanks @ailuntz. But I think we can keep everything in the get_rocm_gpu_arch() function, something like this, wdyt :

def get_rocm_gpu_arch() -> str:
    """Get ROCm GPU architecture."""
    logger = logging.getLogger(__name__)
    try:
        if not torch.version.hip:
            return "unknown"

        if platform.system() == "Windows":
            cmd = ["hipinfo.exe"]
            arch_pattern = r"gcnArchName:\s+gfx([a-zA-Z\d]+)"
        else:
            cmd = ["rocminfo"]
            arch_pattern = r"Name:\s+gfx([a-zA-Z\d]+)"

        result = subprocess.run(cmd, capture_output=True, text=True)
        match = re.search(arch_pattern, result.stdout)
        if match:
            return "gfx" + match.group(1)

    except Exception as e:
        logger.error(f"Could not detect ROCm GPU architecture: {e}")

    # Fallback: check BNB_ROCM_GPU_ARCH env variable
    env_arch = os.environ.get("BNB_ROCM_GPU_ARCH", "").strip()
    if env_arch:
        if re.fullmatch(r"gfx[a-zA-Z0-9]+", env_arch):
            logger.info(f"Using ROCm GPU architecture from BNB_ROCM_GPU_ARCH: {env_arch}")
            return env_arch
        else:
            logger.warning(
                f"BNB_ROCM_GPU_ARCH='{env_arch}' is not a valid architecture name. "
                "Expected format: gfx followed by alphanumeric characters (e.g. gfx942, gfx90a)."
            )

    if torch.cuda.is_available():
        logger.warning(
            "ROCm GPU architecture detection failed despite ROCm being available. "
            "Please ensure 'rocminfo' is installed and available in your PATH. "
            "If that is not possible, you can set the BNB_ROCM_GPU_ARCH environment variable "
            "manually (e.g. export BNB_ROCM_GPU_ARCH=gfx942)."
        )

    return "unknown"

And no need to add tests for this.