Merge branch 'ctranslate-multigpu' into ibaldoall

ivanbaldo · Apr 11, 2024 · cae4a29 · cae4a29
2 parents 9223180 + 9dd2213
commit cae4a29
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/ctranslate/README.md b/ctranslate/README.md
@@ -1 +1,6 @@
-[CTranslate](https://opennmt.net/CTranslate2/guides/transformers.html#llama-2)
+[CTranslate](https://opennmt.net/CTranslate2/guides/transformers.html#llama-2)
+
+In case of multiple GPUs, if for example we have 4, we run
+```sh
+mpirun -np 4 python3 bench.py
+```
diff --git a/ctranslate/bench.py b/ctranslate/bench.py
@@ -6,7 +6,7 @@
 from questions import questions
 import pandas as pd
 
-generator = ctranslate2.Generator("llama-2-7b-ct2", device="cuda")
+generator = ctranslate2.Generator("llama-2-7b-ct2", device="cuda", tensor_parallel=True, flash_attention=True)
 tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
 
 def predict(prompt:str):