Fix triton import guards
State: closed | Pull request: yes
Labels: name: NLP

# 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

```python
# 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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] 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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


jenkins

jenkins

@michal2409 as far as I know we install `triton` in [here](https://github.com/NVIDIA/NeMo/blob/main/requirements/requirements.txt#L13) instead of `pytorch-triton`.  Looks like you install `pytorch-triton` in your container? 

> @michal2409 as far as I know we install `triton` in [here](https://github.com/NVIDIA/NeMo/blob/main/requirements/requirements.txt#L13) 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?

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](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py#L45). Then, we'll have `flash_attn_func_triton = None`

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)

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
```

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?

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.

@hsiehjackson done

jenkins

jenkins

jenkins

jenkins

jenkins