Fix triton import guards

[michal2409]

What does this PR do ?

As far as pkg_resources is concerned, there's no such thing as triton. Instead, it recognizes pytorch-triton. Here's the fix. Also, if the pytorch-triton lib is missing, it raises pkg_resources.DistributionNotFound which currently is not being handled properly.
Collection: nlp

Changelog

• Add specific line by line info of high level changes in this PR.

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Jenkins CI

To run Jenkins, a NeMo User with write access must comment jenkins on the PR.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[michal2409]

jenkins

[michal2409]

jenkins

[hsiehjackson]

@michal2409 as far as I know we install triton in here instead of pytorch-triton. Looks like you install pytorch-triton in your container?

[mmarcinkiewicz]

@michal2409 as far as I know we install triton in here instead of pytorch-triton. Looks like you install pytorch-triton in your container?

You're right. I checked the DLFW containers and they do have pytorch-triton and it's different from triton that gets installed by NeMo. We'll revert this change.

The other outstanding issue is - what is the expected behavior if triton is not installed? Right now it raises an error that is not being handled, but I would imagine that we should catch it and set flash_attn_func_triton = None and continue the execution. Is that expectation correct?

[hsiehjackson]

If triton is not installed, we'll get a ModuleNotFoundError when we try to do from flash_attn.flash_attn_triton import flash_attn_func as flash_attn_func_triton in here. Then, we'll have flash_attn_func_triton = None

[mmarcinkiewicz]

I don't think this is true. I'm getting:

>>> a= pkg_resources.get_distribution("triton")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/pkg_resources/__init__.py", line 528, in get_distribution
    dist = get_provider(dist)
  File "/usr/local/lib/python3.10/dist-packages/pkg_resources/__init__.py", line 400, in get_provider
    return working_set.find(moduleOrReq) or require(str(moduleOrReq))[0]
  File "/usr/local/lib/python3.10/dist-packages/pkg_resources/__init__.py", line 968, in require
    needed = self.resolve(parse_requirements(requirements))
  File "/usr/local/lib/python3.10/dist-packages/pkg_resources/__init__.py", line 829, in resolve
    dist = self._resolve_dist(
  File "/usr/local/lib/python3.10/dist-packages/pkg_resources/__init__.py", line 870, in _resolve_dist
    raise DistributionNotFound(req, requirers)
pkg_resources.DistributionNotFound: The 'triton' distribution was not found and is required by the application

which is not being handled. Could you please double check @hsiehjackson ? Maybe I'm doing something wrong

EDIT:
Essentially what we need is to swap L76 from:
except (ImportError, ModuleNotFoundError, AssertionError):
to
except (ImportError, ModuleNotFoundError, AssertionError, pkg_resources.DistributionNotFound):
(the second line that is changed in this PR)

[hsiehjackson]

When my container install pytorch-triton instead of triton, I'll get the following error. How do you pass from flash_attn.flash_attn_triton import flash_attn_func as flash_attn_func_triton when you just install pytorch-triton?

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/flash_attn/flash_attn_triton.py", line 45, in <module>
    import triton
  File "/usr/local/lib/python3.10/dist-packages/triton/__init__.py", line 1, in <module>
    raise RuntimeError("Should never be installed")
RuntimeError: Should never be installed

[mmarcinkiewicz]

I don't know, it works for me. But sure, if you need triton instead of pytorch-triton then we'll revert the change, no problem - this is actually not a blocker for us. The blocker is that you are not handling the pkg_resources.DistributionNotFound exception and my code crashes. Can we handle that exception as well?

[hsiehjackson]

Sure! @michal2409 Can you help me revert the first line assert pkg_resources.get_distribution("pytorch-triton").version == '2.0.0.dev20221202'? I think we only need the second line as @mmarcinkiewicz suggested.

[michal2409]

@hsiehjackson done

[michal2409]

jenkins

[michal2409]

jenkins

[michal2409]

jenkins

[hsiehjackson]

LGTM!

[michal2409]

jenkins

[michal2409]

jenkins