Skip to content

Commit 4b2ab61

Browse files
OlivierDehaeneOlivierDehaene
authored
feat(candle): support Qwen2 on Cuda (#316)
1 parent 6c6cd93 commit 4b2ab61

18 files changed

+6805
-26
lines changed

‎README.md‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -64,20 +64,22 @@ Ember, GTE and E5. TEI implements many features such as:
6464
#### Text Embeddings
6565

6666
Text Embeddings Inference currently supports Nomic, BERT, CamemBERT, XLM-RoBERTa models with absolute positions, JinaBERT
67-
model with Alibi positions and Mistral, Alibabe GTE models with Rope positions.
67+
model with Alibi positions and Mistral, Alibaba GTE and Qwen2 models with Rope positions.
6868

6969
Below are some examples of the currently supported models:
7070

71-
| MTEB Rank | Model Size | Model Type | Model ID |
72-
|-----------|----------------|-------------|--------------------------------------------------------------------------------------------------|
73-
| 1 | 7B (Very Slow) | Mistral | [Salesforce/SFR-Embedding-2_R](https://hf.co/Salesforce/SFR-Embedding-2_R) |
74-
| 15 | 0.4B | Alibaba GTE | [Alibaba-NLP/gte-large-en-v1.5](Alibaba-NLP/gte-large-en-v1.5) |
75-
| 20 | 0.3B | Bert | [WhereIsAI/UAE-Large-V1](https://hf.co/WhereIsAI/UAE-Large-V1) |
76-
| 24 | 0.5B | XLM-RoBERTa | [intfloat/multilingual-e5-large-instruct](https://hf.co/intfloat/multilingual-e5-large-instruct) |
77-
| N/A | 0.1B | NomicBert | [nomic-ai/nomic-embed-text-v1](https://hf.co/nomic-ai/nomic-embed-text-v1) |
78-
| N/A | 0.1B | NomicBert | [nomic-ai/nomic-embed-text-v1.5](https://hf.co/nomic-ai/nomic-embed-text-v1.5) |
79-
| N/A | 0.1B | JinaBERT | [jinaai/jina-embeddings-v2-base-en](https://hf.co/jinaai/jina-embeddings-v2-base-en) |
80-
| N/A | 0.1B | JinaBERT | [jinaai/jina-embeddings-v2-base-code](https://hf.co/jinaai/jina-embeddings-v2-base-code) |
71+
| MTEB Rank | Model Size | Model Type | Model ID |
72+
|-----------|---------------------|-------------|--------------------------------------------------------------------------------------------------|
73+
| 1 | 7B (Very Expensive) | Mistral | [Salesforce/SFR-Embedding-2_R](https://hf.co/Salesforce/SFR-Embedding-2_R) |
74+
| 2 | 7B (Very Expensive) | Qwen2 | [Alibaba-NLP/gte-Qwen2-7B-instruct](https://hf.co/Alibaba-NLP/gte-Qwen2-7B-instruct) |
75+
| 9 | 1.5B (Expensive) | Qwen2 | [Alibaba-NLP/gte-Qwen2-1.5B-instruct](https://hf.co/Alibaba-NLP/gte-Qwen2-1.5B-instruct) |
76+
| 15 | 0.4B | Alibaba GTE | [Alibaba-NLP/gte-large-en-v1.5](Alibaba-NLP/gte-large-en-v1.5) |
77+
| 20 | 0.3B | Bert | [WhereIsAI/UAE-Large-V1](https://hf.co/WhereIsAI/UAE-Large-V1) |
78+
| 24 | 0.5B | XLM-RoBERTa | [intfloat/multilingual-e5-large-instruct](https://hf.co/intfloat/multilingual-e5-large-instruct) |
79+
| N/A | 0.1B | NomicBert | [nomic-ai/nomic-embed-text-v1](https://hf.co/nomic-ai/nomic-embed-text-v1) |
80+
| N/A | 0.1B | NomicBert | [nomic-ai/nomic-embed-text-v1.5](https://hf.co/nomic-ai/nomic-embed-text-v1.5) |
81+
| N/A | 0.1B | JinaBERT | [jinaai/jina-embeddings-v2-base-en](https://hf.co/jinaai/jina-embeddings-v2-base-en) |
82+
| N/A | 0.1B | JinaBERT | [jinaai/jina-embeddings-v2-base-code](https://hf.co/jinaai/jina-embeddings-v2-base-code) |
8183

8284

8385
To explore the list of best performing text embeddings models, visit the

‎backends/candle/src/lib.rs‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@ use crate::compute_cap::{
1212
};
1313
use crate::models::{
1414
BertConfig, BertModel, DistilBertConfig, DistilBertModel, GTEConfig, JinaBertModel,
15-
JinaCodeBertModel, MistralConfig, Model, NomicBertModel, NomicConfig,
15+
JinaCodeBertModel, MistralConfig, Model, NomicBertModel, NomicConfig, Qwen2Config,
1616
};
1717
#[cfg(feature = "cuda")]
1818
use crate::models::{
1919
FlashBertModel, FlashDistilBertModel, FlashGTEModel, FlashJinaBertModel,
20-
FlashJinaCodeBertModel, FlashMistralModel, FlashNomicBertModel,
20+
FlashJinaCodeBertModel, FlashMistralModel, FlashNomicBertModel, FlashQwen2Model,
2121
};
2222
use anyhow::Context;
2323
use candle::{DType, Device};
@@ -59,6 +59,7 @@ enum Config {
5959
Mistral(MistralConfig),
6060
#[serde(rename = "new")]
6161
Gte(GTEConfig),
62+
Qwen2(Qwen2Config),
6263
}
6364

6465
pub struct CandleBackend {
@@ -221,6 +222,10 @@ impl CandleBackend {
221222
"GTE is only supported on Cuda devices in fp16 with flash attention enabled"
222223
.to_string(),
223224
)),
225+
(Config::Qwen2(_), Device::Cpu | Device::Metal(_)) => Err(BackendError::Start(
226+
"Qwen2 is only supported on Cuda devices in fp16 with flash attention enabled"
227+
.to_string(),
228+
)),
224229
#[cfg(feature = "cuda")]
225230
(Config::Bert(config), Device::Cuda(_)) => {
226231
if cfg!(any(feature = "flash-attn", feature = "flash-attn-v1"))
@@ -342,14 +347,25 @@ impl CandleBackend {
342347
#[cfg(feature = "cuda")]
343348
(Config::Gte(config), Device::Cuda(_)) => {
344349
if dtype != DType::F16
345-
|| !cfg!(feature = "flash-attn")
346-
|| get_runtime_compute_cap().unwrap() < 80
350+
|| !cfg!(any(feature = "flash-attn", feature = "flash-attn-v1"))
347351
{
348-
return Err(BackendError::Start("GTE is only supported on Cuda devices in fp16 with flash attention v2 enabled".to_string()));
352+
return Err(BackendError::Start("GTE is only supported on Cuda devices in fp16 with flash attention enabled".to_string()));
349353
}
350354
tracing::info!("Starting FlashGTE model on {:?}", device);
351355
Ok(Box::new(FlashGTEModel::load(vb, &config, model_type).s()?))
352356
}
357+
#[cfg(feature = "cuda")]
358+
(Config::Qwen2(config), Device::Cuda(_)) => {
359+
if dtype != DType::F16
360+
|| !cfg!(any(feature = "flash-attn", feature = "flash-attn-v1"))
361+
{
362+
return Err(BackendError::Start("Qwen2 is only supported on Cuda devices in fp16 with flash attention v2 enabled".to_string()));
363+
}
364+
tracing::info!("Starting FlashQwen2 model on {:?}", device);
365+
Ok(Box::new(
366+
FlashQwen2Model::load(vb, &config, model_type).s()?,
367+
))
368+
}
353369
};
354370

355371
Ok(Self {

‎backends/candle/src/models/flash_bert.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -400,7 +400,10 @@ impl FlashBertModel {
400400
// Get token indices form cu_seqlens
401401
let mut indices = match self.pool {
402402
Pool::Cls => cu_seqlens.narrow(0, 0, batch_size)?,
403-
Pool::LastToken => cu_seqlens.narrow(0, 1, batch_size)?,
403+
Pool::LastToken => {
404+
let end = cu_seqlens.narrow(0, 1, batch_size)?;
405+
(&end - &end.ones_like()?)?
406+
}
404407
_ => unreachable!(),
405408
};
406409

‎backends/candle/src/models/flash_distilbert.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,10 @@ impl FlashDistilBertModel {
266266
// Get token indices form cu_seqlens
267267
let mut indices = match self.pool {
268268
Pool::Cls => cu_seqlens.narrow(0, 0, batch_size)?,
269-
Pool::LastToken => cu_seqlens.narrow(0, 1, batch_size)?,
269+
Pool::LastToken => {
270+
let end = cu_seqlens.narrow(0, 1, batch_size)?;
271+
(&end - &end.ones_like()?)?
272+
}
270273
_ => unreachable!(),
271274
};
272275

‎backends/candle/src/models/flash_gte.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,10 @@ impl FlashGTEModel {
352352
// Get token indices form cu_seqlens
353353
let mut indices = match self.pool {
354354
Pool::Cls => cu_seqlens.narrow(0, 0, batch_size)?,
355-
Pool::LastToken => cu_seqlens.narrow(0, 1, batch_size)?,
355+
Pool::LastToken => {
356+
let end = cu_seqlens.narrow(0, 1, batch_size)?;
357+
(&end - &end.ones_like()?)?
358+
}
356359
_ => unreachable!(),
357360
};
358361

‎backends/candle/src/models/flash_jina.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -327,7 +327,10 @@ impl FlashJinaBertModel {
327327
// Get token indices form cu_seqlens
328328
let mut indices = match self.pool {
329329
Pool::Cls => cu_seqlens.narrow(0, 0, batch_size)?,
330-
Pool::LastToken => cu_seqlens.narrow(0, 1, batch_size)?,
330+
Pool::LastToken => {
331+
let end = cu_seqlens.narrow(0, 1, batch_size)?;
332+
(&end - &end.ones_like()?)?
333+
}
331334
_ => unreachable!(),
332335
};
333336

‎backends/candle/src/models/flash_jina_code.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,10 @@ impl FlashJinaCodeBertModel {
380380
// Get token indices form cu_seqlens
381381
let mut indices = match self.pool {
382382
Pool::Cls => cu_seqlens.narrow(0, 0, batch_size)?,
383-
Pool::LastToken => cu_seqlens.narrow(0, 1, batch_size)?,
383+
Pool::LastToken => {
384+
let end = cu_seqlens.narrow(0, 1, batch_size)?;
385+
(&end - &end.ones_like()?)?
386+
}
384387
_ => unreachable!(),
385388
};
386389

‎backends/candle/src/models/flash_mistral.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -334,7 +334,10 @@ impl FlashMistralModel {
334334
// Get token indices form cu_seqlens
335335
let mut indices = match self.pool {
336336
Pool::Cls => cu_seqlens.narrow(0, 0, batch_size)?,
337-
Pool::LastToken => cu_seqlens.narrow(0, 1, batch_size)?,
337+
Pool::LastToken => {
338+
let end = cu_seqlens.narrow(0, 1, batch_size)?;
339+
(&end - &end.ones_like()?)?
340+
}
338341
_ => unreachable!(),
339342
};
340343

‎backends/candle/src/models/flash_nomic.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,10 @@ impl FlashNomicBertModel {
311311
// Get token indices form cu_seqlens
312312
let mut indices = match self.pool {
313313
Pool::Cls => cu_seqlens.narrow(0, 0, batch_size)?,
314-
Pool::LastToken => cu_seqlens.narrow(0, 1, batch_size)?,
314+
Pool::LastToken => {
315+
let end = cu_seqlens.narrow(0, 1, batch_size)?;
316+
(&end - &end.ones_like()?)?
317+
}
315318
_ => unreachable!(),
316319
};
317320

0 commit comments

Comments
 (0)