huggingface / huggingface/candle
metal -> CPU Tensor::to_device is slow for model output
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
When using `sentence-transformers/all-MiniLM-L6-v2` on m1 PRO mac (32GB ram) with --features metal on,
attempting to generate embeddings for 1000 sentences and attempting to convert the final tensor to a Vec> the `to_device` function of Tensor seems very slow.
I believe to_device is important and must be performant in order to switch from Tensor to Vec and potentially Serialize the output from a model in an application
## How to reproduce:
place the following code snipet (below) in candle-examples/examples/bert-metal/main.rs
cargo run --example bert-metal --features metal
```
use candle_transformers::models::bert::{BertModel, Config, DTYPE};
use anyhow::{Error as E, Result};
use candle::{Device, Tensor};
use candle_nn::VarBuilder;
use hf_hub::{api::sync::Api, Repo, RepoType};
use tokenizers::{PaddingParams, Tokenizer};
fn build_model_and_tokenizer() -> Result<(BertModel, Tokenizer)> {
let device = candle_examples::device(false)?;
let repo = Repo::with_revision(
"sentence-transformers/all-MiniLM-L6-v2".to_string(),
RepoType::Model,
"refs/pr/21".to_string(),
);
let (config_filename, tokenizer_filename, weights_filename) = {
let api = Api::new()?.repo(repo);
let config = api.get("config.json")?;
let tokenizer = api.get("tokenizer.json")?;
let weights = api.get("model.safetensors")?;
(config, tokenizer, weights)
};
let config = std::fs::read_to_string(config_filename)?;
let config: Config = serde_json::from_str(&config)?;
let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights_filename], DTYPE, &device)? };
let model = BertModel::load(vb, &config)?;
Ok((model, tokenizer))
}
pub fn normalize_l2(v: &Tensor) -> Result {
Ok(v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)?)
}
fn get_mock_2d_tensor(vec_count: usize, vec_size: usize, d: &Device) -> Tensor {
let mut data = Vec::with_capacity(vec_count);
for i in 0..vec_count {
data.push(vec![i as f32; vec_size]);
}
Tensor::new(data, d).expect("failed to generate tensor")
}
fn main() -> Result<()> {
let (model, mut tokenizer) = build_model_and_tokenizer()?;
let device = &model.device;
let sentence_count = 1_000;
let sentences: Vec<&str> =
vec!["This is a sample sentence that should be short enough"; sentence_count];
let pp = PaddingParams {
strategy: tokenizers::PaddingStrategy::Fixed(128),
..Default::default()
};
tokenizer.with_padding(Some(pp));
let tokens = tokenizer.encode_batch(sentences, true).map_err(E::msg)?;
let mut token_ids = Vec::with_capacity(tokens.len());
for encoding in tokens {
token_ids.push(encoding.get_ids().to_vec());
}
let token_ids = Tensor::new(token_ids, device)?;
let token_type_ids = token_ids.zeros_like()?;
println!("running inference on batch {:?}", token_ids.shape());
let start = std::time::Instant::now();
let embeddings = model.forward(&token_ids, &token_type_ids)?;
println!("model forward timing {:?}", start.elapsed());
let start = std::time::Instant::now();
let (_n_sentence, n_tokens, _hidden_size) = embeddings.dims3()?;
let embeddings = (embeddings.sum(1)? / (n_tokens as f64))?;
let embeddings = normalize_l2(&embeddings)?;
println!(
"pooled embeddings {:?} in {:?}",
embeddings.shape(),
start.elapsed()
);
let start = std::time::Instant::now();
embeddings.to_device(&Device::Cpu)?;
println!(
"moved embedding tensor back to cpu in {:?}",
start.elapsed()
);
let mock_tensor = get_mock_2d_tensor(1000, 384, device);
let start = std::time::Instant::now();
let old_device = mock_tensor.device();
mock_tensor.to_device(&Device::Cpu)?;
println!(
"moved mock tensor of size {:?} from {:?} back to cpu in {:?}",
mock_tensor.shape(),
old_device,
start.elapsed()
);
Ok(())
}
```
output
```
running inference on batch [1000, 128]
model forward timing 185.251708ms
pooled embeddings [1000, 384] in 233.083µs
moved embedding tensor back to cpu in 18.590362041s
moved mock tensor of size [1000, 384] from Metal(MetalDevice(4294969642)) back to cpu in 517.167µs
```
I added a sample move of a manually generated tensor at the end to show that the issue might not be related to the size of the tensor.
## Investigation
- The issue might be related to the`MetalStorage.to_cpu`, when debugging I added a `self.device.wait_until_completed` statement at the top of the function and found that most of the time was spent there, could there be commands waiting to be flushed?
- Is there a way to generate a new Tensor.storage or Tensor.storage.device that would only contain the [1000, 384] floats for an existing tensor? with a fresh `command_queue`, `command_buffer` and `buffers`from an existing Tensor to maybe bypass this issue?
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with candle-examples/examples/bert-metal/main.rs and reproduce the timing difference between the model output and the manually generated tensor using cargo run --example bert-metal --features metal. Then inspect MetalStorage.to_cpu and the wait_until_completed call mentioned in the investigation. Done means explaining the excessive wait and validating improved transfer timing without changing the reproduction's expected output.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100