API for TensorIDs in GradStore for O(num_relevant_vars) backward_step

[spaghetti-source]

Hello folks.

I implemented word2vec in candle, but it was very slow. The reason was that the candle_nn's SGD::step takes O(num_variables) regardless of the number of relevant variables. (rem: In word2vec training, we take words in sentences and update only their embeddings.) Below is the relevant code fragment in candle.

// candle-nn/src/optim.rs
impl SGD {
    fn step(&mut self, grads: &candle::backprop::GradStore) -> Result<()> {
        for var in self.vars.iter() {
            if let Some(grad) = grads.get(var) {
                var.set(&var.sub(&(grad * self.learning_rate)?)?)?;
            }
        }
        Ok(())
    }
    // ...
}

To mitigate this issue, I propose to add an API to retrieve a list of TensorId in the GradStore. The following is my suggested implementation.

// candle-core/src/backprop.rs
impl GradStore {
    pub fn get_ids(&self) -> impl Iterator<Item = &TensorId> {
        self.0.keys()
    }
    // ...
}

By having this API, we can implement our own SGD as follows to mitigate the issue.

// my own code
struct MySGD {
    vars: HashMap<TensorId, Var>,
    learning_rate: f64,
}
impl MySGD {
    fn step(&mut self, grads: &GradStore) -> Result<()> {
        for id in grads.get_ids() {
            if let Some(var) = self.vars.get(id) {
                if let Some(grad) = grads.get(var) {
                    var.set(&var.sub(&(grad * self.learning_rate)?)?)?;
                }
            }
        }
        Ok(())
    }
    // ...
}

This provides a significant speed-up in use cases that update only relevant variables like word2vec. I implemented a simple benchmark (https://gist.github.com/spaghetti-source/f0630f1d0ad1b98f736a1d8e9719ff6d) and observed the following speed-up on my local computer.

vocabulary_size	original	proposed
10_000	2.3 s	0.025 s
50_000	80 s	0.13 s

[LaurentMazare]

Sounds good, feel free to make a PR adding get_ids.

[spaghetti-source]

Thanks. I've created a PR #2379