burn_train/learner/
sharder.rs

1use burn_core::{
2    Tensor,
3    module::{ModuleMapper, Param},
4};
5
6use crate::{Learner, LearnerModel};
7
8/// Describes how the module is distributed across multiple devices.
9pub struct ModuleSharder;
10
11impl ModuleMapper for ModuleSharder {
12    fn map_float<const D: usize>(&mut self, param: Param<Tensor<D>>) -> Param<Tensor<D>> {
13        let (id, tensor, mapper) = param.consume();
14        let tensor = tensor.set_distributed(id);
15        Param::from_mapped_value(id, tensor, mapper)
16    }
17}
18
19impl<M: LearnerModel> Learner<M> {
20    /// Mark the model as sharded across multiple devices.
21    pub fn grad_sharded(&mut self) {
22        self.model = self.model.clone().map(&mut ModuleSharder);
23    }
24}

burn_train/learner/sharder.rs

burn_train/learner/
sharder.rs