huggingface / huggingface/candle

Poor performance in back propagation

Open
#2,913 5 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

(All numbers are measured on an M1 Max, with samply -- a sampling CPU profiler)

The performance numbers I'll share are for the following model:

```rust
const N_EMBD: usize = 32;
const N_HIDDEN: usize = 128;

let model = Sequential::new()
// Embedding layer
.add(candle_nn::embedding(
num_hero_tokens.into(),
N_EMBD,
vb.push_prefix("embedding"),
)?)
// First hidden layer
.add(FlattenConsecutive::new(5))
.add(candle_nn::linear_no_bias(
N_EMBD * 5,
N_HIDDEN,
vb.push_prefix("linear[0]"),
)?)
.add(TransposedBatchNorm::new(candle_nn::batch_norm(
N_HIDDEN,
candle_nn::BatchNormConfig::default(),
vb.push_prefix(format!("batch_norm[0]")),
)?))
.add(|t: &Tensor| t.tanh())
// Second hidden layer
.add(FlattenConsecutive::new(2))
.add(|t: &Tensor| t.squeeze(1))
.add(candle_nn::linear_no_bias(
N_HIDDEN * 2,
N_HIDDEN,
vb.push_prefix("linear[1]"),
)?)
.add(candle_nn::batch_norm(
N_HIDDEN,
candle_nn::BatchNormConfig::default(),
vb.push_prefix(format!("batch_norm[1]")),
)?)
.add(|t: &Tensor| t.tanh())
// Output layer
.add(candle_nn::linear(
N_HIDDEN,
1,
vb.push_prefix("linear[output]"),
)?)
// Sigmoid ensures that probabilities are between 0 and 1
.add(candle_nn::Activation::Sigmoid)
.add(|t: &Tensor| t.squeeze(1));
```

And training is done with the following code:

```
fn train(
dev: &candle_core::Device,
vm: &candle_nn::VarMap,
model: &impl ModuleT,
x: &Tensor,
y: &Tensor,
) -> anyhow::Result<()> {
let mut opt = candle_nn::AdamW::new_lr(vm.all_vars(), 0.01)?;
let mut rng = rand::rng();
for i in 0..TRAINING_EPOCHS {
let idx = Tensor::from_iter(
(0..TRAINING_BATCH_SIZE).map(|_| rng.random_range(0..x.dims()[0]) as u32),
dev,
)?;
let x_batch = x.index_select(&idx, 0)?;
let y_batch = y.index_select(&idx, 0)?;

let logits = model.forward_t(&x_batch, true)?;
let loss = candle_nn::loss::binary_cross_entropy_with_logit(&logits, &y_batch)
.context("cross_entropy")?;

if i % 1000 == 0 {
println!("{i:7}/{TRAINING_EPOCHS}: {}", loss.to_scalar::()?);
}

opt.backward_step(&loss)?;
}

Ok(())
}
```

Running this on a `Cpu` device I see that:

- 95% of time is spent in `backward_step`
- 89% in `Tensor::backward`
- 6% in `AdamW::step`
- 4% of time is spent in `forward_t`

I'm not much of an ML-hand, but speaking to a friend with far more expertise, he tells me that:

> Very roughly, backprop is typically 2x the FLOPs of the forward pass, and between 1-2x the number of kernels/operators. So if it's much more than 2x the time, that's suspicious

Breaking down `Tensor::backward`:
- 57% in `Tensor::add`
- 15% in `Storage::reduce_op`
- 6% in `Arc::drop_slow`
- 5% in `GradStore::or_insert`
- 3.5% in `Tensor::matmul`
- etc.

I also ran with the Metal backend (and did not experience a significant speedup relative to CPU), and obtained the following rough breakdown:

- 60% of time in backprop
- 20% `GradStore::or_insert`
- 33% of time in `AdamW::step`
- 6% in `forward_t`

However, these numbers should be considered less reliable: because the numbers are captured with a sampling CPU profiler, time spent is really in synchronization points between the CPU/GPU or userspace/kernel, so these may not be a proper reflection of where computation time is spent.

With all that said, the conclusion that `backwards` is too much seems inescapable, and I want to offer a few observations:

- There is no way to reuse a `GradStore` (zeroing out the gradient `Tensor`s, instead of allocating them fresh on each `Tensor::backwards` call). This leads to additional allocations (I'm far from an expert, but my understanding is that many GPU allocators are quite slow, so avoiding tons of extra allocation traffic is desirable.)
- Because `Tensor` is immutable, all of the `Tensor::add` in `backwards` really add up (pun not intended): You're going to have `add` proportional to the number of ops in back propagation (because everyone has to add to the accmumulated gradient), and in the current design each one of those is an allocation (and then a free!) and not an in-place op.
- `AdamW::step` time feels excessive in my Metal measurements, but since it didn't reproduce under CPU, I'm ignoring it for now.

I'm not proposing any specific actions, as some of these cut right to the core of candle's design decisions and I don't want to get ahead of myself.

Happy to share more of the code if that'd be helpful, also happy to run any additional tests or experiments. Cheers!

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.