Monkeypatch flash attention in for llama

[dakinggg]

This PR adds a monkeypatch of triton flash attention for llama2 models. If we start doing this for more models we can try to generalize into an algorithm, but until that time I think this monkeypatch is good enough.

TODO:

• Copy profiling results for 7b
• Copy profiling results for 70b
• Paste in evidence that the test passes

Result for 7b scale to compare implementations:

Test results:

tests/test_llama_patch.py::test_patch_equivalence[meta-llama/Llama-2-7b-hf-True-torch] 
------------------------------------------------------------------------------------------------------------------- live log call -------------------------------------------------------------------------------------------------------------------
2023-08-11 22:18:17 [    INFO] Setting seed to 42 (reproducibility.py:159)
2023-08-11 22:18:23 [    INFO] Setting seed to 42 (reproducibility.py:159)
PASSED                                                                                                                                                                                                                                        [ 12%]
tests/test_llama_patch.py::test_patch_equivalence[meta-llama/Llama-2-7b-hf-True-triton] 
------------------------------------------------------------------------------------------------------------------- live log call -------------------------------------------------------------------------------------------------------------------
2023-08-11 22:18:23 [    INFO] Setting seed to 42 (reproducibility.py:159)
2023-08-11 22:18:24 [    INFO] Setting seed to 42 (reproducibility.py:159)
PASSED                                                                                                                                                                                                                                        [ 25%]
tests/test_llama_patch.py::test_patch_equivalence[meta-llama/Llama-2-7b-hf-False-torch] 
------------------------------------------------------------------------------------------------------------------- live log call -------------------------------------------------------------------------------------------------------------------
2023-08-11 22:18:25 [    INFO] Setting seed to 42 (reproducibility.py:159)
2023-08-11 22:18:26 [    INFO] Setting seed to 42 (reproducibility.py:159)
PASSED                                                                                                                                                                                                                                        [ 37%]
tests/test_llama_patch.py::test_patch_equivalence[meta-llama/Llama-2-7b-hf-False-triton] 
------------------------------------------------------------------------------------------------------------------- live log call -------------------------------------------------------------------------------------------------------------------
2023-08-11 22:18:27 [    INFO] Setting seed to 42 (reproducibility.py:159)
2023-08-11 22:18:27 [    INFO] Setting seed to 42 (reproducibility.py:159)
PASSED                                                                                                                                                                                                                                        [ 50%]
tests/test_llama_patch.py::test_patch_equivalence[meta-llama/Llama-2-70b-hf-True-torch] 
------------------------------------------------------------------------------------------------------------------- live log call -------------------------------------------------------------------------------------------------------------------
2023-08-11 22:18:28 [    INFO] Setting seed to 42 (reproducibility.py:159)
2023-08-11 22:18:29 [    INFO] Setting seed to 42 (reproducibility.py:159)
PASSED                                                                                                                                                                                                                                        [ 62%]
tests/test_llama_patch.py::test_patch_equivalence[meta-llama/Llama-2-70b-hf-True-triton] 
------------------------------------------------------------------------------------------------------------------- live log call -------------------------------------------------------------------------------------------------------------------
2023-08-11 22:18:30 [    INFO] Setting seed to 42 (reproducibility.py:159)
2023-08-11 22:18:32 [    INFO] Setting seed to 42 (reproducibility.py:159)
PASSED                                                                                                                                                                                                                                        [ 75%]
tests/test_llama_patch.py::test_patch_equivalence[meta-llama/Llama-2-70b-hf-False-torch] 
------------------------------------------------------------------------------------------------------------------- live log call -------------------------------------------------------------------------------------------------------------------
2023-08-11 22:18:33 [    INFO] Setting seed to 42 (reproducibility.py:159)
2023-08-11 22:18:34 [    INFO] Setting seed to 42 (reproducibility.py:159)
PASSED                                                                                                                                                                                                                                        [ 87%]
tests/test_llama_patch.py::test_patch_equivalence[meta-llama/Llama-2-70b-hf-False-triton] 
------------------------------------------------------------------------------------------------------------------- live log call -------------------------------------------------------------------------------------------------------------------
2023-08-11 22:18:36 [    INFO] Setting seed to 42 (reproducibility.py:159)
2023-08-11 22:18:37 [    INFO] Setting seed to 42 (reproducibility.py:159)
PASSED                                                                                                                                                                                                                                        [100%]

================================================================================================================= warnings summary ==================================================================================================================
../../miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/accelerate/utils/dataclasses.py:29
  /mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/accelerate/utils/dataclasses.py:29: DeprecationWarning: The distutils package is deprecated and slated for removal in Python 3.12. Use setuptools or check PEP 632 for potential alternatives
    from distutils.util import strtobool

../../miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py:15
  /mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py:15: DeprecationWarning: The distutils.sysconfig module is deprecated, use sysconfig instead
    import distutils.sysconfig

../../miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/jupyter_client/connect.py:20
  /mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/jupyter_client/connect.py:20: DeprecationWarning: Jupyter is migrating its paths to use standard platformdirs
  given by the platformdirs library.  To remove this warning and
  see the appropriate new directories, set the environment variable
  `JUPYTER_PLATFORM_DIRS=1` and then run `jupyter --paths`.
  The use of platformdirs will be the default in `jupyter_core` v6
    from jupyter_core.paths import jupyter_data_dir, jupyter_runtime_dir, secure_write

../../miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/comet_ml/monkey_patching.py:19
  /mnt/workdisk/danielking/miniconda3/envs/foundry-3.10/lib/python3.10/site-packages/comet_ml/monkey_patching.py:19: DeprecationWarning: the imp module is deprecated in favour of importlib and slated for removal in Python 3.12; see the module's documentation for alternative uses
    import imp

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
========================================================================================================== 8 passed, 4 warnings in 22.00s ===========================================================================================================

70b running on 16 and 32 gpus:

[vchiley]

left a few comments

[germanjke]

Hi guys! @vchiley @dakinggg when do you planning to release this request? Thanks

[vchiley]

lgtm