More precise Errors and logger warnings for failing to import triton

[b-burton]

🚀 Feature

Better handling of errors supplied to users when attempting (and failing) to import triton - namely when their OS does not support it.

Motivation

Aside from the current error being annoying to keep receiving (not a big deal as it can be suppressed), it is not very informative. When switch between OS's it can appear that an error you would ignore on Windows is the same as careless mistake (not installing triton) in a Linux environment.

Pitch

Changing the init.py file to have clearer and more useful error handling.

Current

xformers/xformers/__init__.py

Lines 56 to 65 in 6e10bd2

	try:
	import triton # noqa
	return True
	except (ImportError, AttributeError):
	logger.warning(
	"A matching Triton is not available, some optimizations will not be enabled",
	exc_info=True,
	)
	return False

Suggested Change

try:
    import triton  # noqa

    return True
except (ModuleNotFoundError):

    if os.name == "nt":  # If the operating system is windows
        logger.warning(
            "Triton is currently not available on Windows, some optimizations will not be enabled",
            exc_info=False,
        )
    else:

        if os.uname().sysname == 'Darwin':  # If the operating system is MacOS (extra check because MacOS is also Unix-like)
            logger.warning(
                "Triton is currently not available on MacOS, some optimizations will not be enabled",
                exc_info=False,
            )
        raise ModuleNotFoundError(
            "No Triton module found. For installation instructions, see: https://github.com/triton-lang/triton/?tab=readme-ov-file#quick-installation"
        )  # Rasie an actual error - the user can install triton but has not. Could be changed to a logger.warning if too overkill
    return False

except (ImportError, AttributeError):
    logger.warning(
        "A matching Triton is not available, some optimizations will not be enabled",
        exc_info=True,
    )
    return False

Alternatives

Alternatives could include:

1. The one suggested above
2. Leaving the code as it is - not that big of a deal
3. Leaving the code as is but changing the current logger message
4. An earlier catch telling the user that their experience will be degraded on windows

Additional context

There seems to be a lot of interest in triton for windows so making a clearer error may help:
triton-lang/triton#1640

[lw]

Could you submit a PR with your proposal?

I'm not sure if I'm a big fan of all these checks and branches in that snippet: they do deliver a very clear and targeted message, but they add complexity. I wonder if just changing the message (e.g., evoking that it's unavailable on Windows and MacOS) might be enough?

Also note that worrying about MacOS is a very niche concern, since MacOS doesn't even support CUDA. If one day we'll support MPS (highly unlikely in the short term) we might worry more about it.