Update after review

Review comments: - foundation-model-stack#52 (review) Signed-off-by: Martin Hickey <[email protected]>
VassilisVassiliadis · Feb 26, 2024 · 647da3e · 647da3e
1 parent e898b1d
commit 647da3e
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -13,6 +13,8 @@ pip install -U datasets
 pip install -e .
 ```
 
+> Note: If you wish to use [FlashAttention](https://github.com/Dao-AILab/flash-attention), then you need to install these requirements: `pip install -r flashattn_requirements`. [FlashAttention](https://github.com/Dao-AILab/flash-attention) requires the [CUDA Toolit](https://developer.nvidia.com/cuda-toolkit) to be pre-installed.
+
 ## Data format
 The data format expectation is a single column text. The trainer is configured to expect a response template as a string. For example, if one wants to prepare the `alpaca` format data to feed into this trainer, it is quite easy and can be done with the following code.
 

diff --git a/torch_requirements.txt → flashattn_requirements.txt b/torch_requirements.txt → flashattn_requirements.txt
@@ -1,4 +1,2 @@
-packaging
-torch
 wheel
 flash-attn
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,8 @@
 numpy
 accelerate>=0.20.3
+packaging
 transformers>=4.34.1
+torch
 aim==3.17.5
 sentencepiece
 tokenizers>=0.13.3
@@ -10,3 +12,4 @@ ninja
 peft>=0.8.0
 datasets>=2.15.0
 fire
+