From 9dd2213e45a354666edd0be777208122966577ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Iv=C3=A1n=20Baldo?= Date: Thu, 11 Apr 2024 16:42:40 -0300 Subject: [PATCH] CTranslate2: support multiple GPUs (if run with mpirun) and Flash Attention 2. --- ctranslate/README.md | 7 ++++++- ctranslate/bench.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/ctranslate/README.md b/ctranslate/README.md index 4dbd2bd..d94ffd2 100644 --- a/ctranslate/README.md +++ b/ctranslate/README.md @@ -1 +1,6 @@ -[CTranslate](https://opennmt.net/CTranslate2/guides/transformers.html#llama-2) \ No newline at end of file +[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 +``` \ No newline at end of file diff --git a/ctranslate/bench.py b/ctranslate/bench.py index 94b59d5..0a161d3 100644 --- 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):