burn_core/module/parallel/
base.rs

1use alloc::vec::Vec;
2
3use super::DistributedLayer;
4use crate::module::Module;
5
6/// A model that runs as an input layer, then hidden layers in order, then an output layer, so its
7/// layers can be split across devices by a [`DistributedLayeredModel`](super::DistributedLayeredModel).
8///
9/// It is usually not the struct the model trained with, since a layer split across devices may be
10/// represented differently. Weights trained on another struct load into this one by remapping
11/// their keys. Each layer is built on the device a [`LayerPlacement`](super::LayerPlacement) gives
12/// it, so the model never exists whole on one device.
13///
14/// ```rust,ignore
15/// impl Model {
16///     fn new(config: &ModelConfig, placement: &LayerPlacement) -> Self {
17///         Self {
18///             embedding: Embedding::new(config, &placement.input),
19///             blocks: placement
20///                 .hidden
21///                 .iter()
22///                 .map(|device| Block::new(config, device))
23///                 .collect(),
24///             head: Head::new(config, &placement.output),
25///         }
26///     }
27/// }
28///
29/// impl LayerParallelism for Model {
30///     type InputLayer = Embedding;
31///     type HiddenLayer = Block;
32///     type OutputLayer = Head;
33///
34///     fn layer_input(&self) -> &Embedding {
35///         &self.embedding
36///     }
37///
38///     fn layers_hidden(&self) -> impl Iterator<Item = &Block> {
39///         self.blocks.iter()
40///     }
41///
42///     fn layer_output(&self) -> &Head {
43///         &self.head
44///     }
45///
46///     fn into_layers(self) -> LayerSequence<Self> {
47///         LayerSequence {
48///             input: self.embedding,
49///             hidden: self.blocks,
50///             output: self.head,
51///         }
52///     }
53/// }
54///
55/// let placement = LayerPlacement::even(&devices, config.num_blocks);
56/// let model = Model::new(&config, &placement);
57/// let model = DistributedLayeredModel::new(model, &placement);
58/// let predictions = model.forward(features);
59/// ```
60pub trait LayerParallelism: Module {
61    /// Runs first, on what the model takes; what it returns is what the hidden layers pass along.
62    type InputLayer: DistributedLayer<Input: Module, Output: Module>;
63    /// Each hidden layer.
64    type HiddenLayer: DistributedLayer<Input = HiddenLayerSignal<Self>, Output = HiddenLayerSignal<Self>>;
65    /// Runs last; what it returns is what the model returns.
66    type OutputLayer: DistributedLayer<Input = HiddenLayerSignal<Self>>;
67
68    /// The input layer.
69    fn layer_input(&self) -> &Self::InputLayer;
70
71    /// The hidden layers in forward order, the same order on every call.
72    fn layers_hidden(&self) -> impl Iterator<Item = &Self::HiddenLayer>;
73
74    /// The output layer.
75    fn layer_output(&self) -> &Self::OutputLayer;
76
77    /// The model taken apart into its layers, the hidden ones in the order
78    /// [`layers_hidden`](Self::layers_hidden) gives them.
79    fn into_layers(self) -> LayerSequence<Self>;
80}
81
82/// What the hidden layers of a [`LayerParallelism`] model take and return: whatever its input
83/// layer produces.
84pub type HiddenLayerSignal<M> = <<M as LayerParallelism>::InputLayer as DistributedLayer>::Output;
85
86/// The layers of a [`LayerParallelism`] model, owned.
87#[derive(Debug)]
88pub struct LayerSequence<M: LayerParallelism> {
89    /// The input layer.
90    pub input: M::InputLayer,
91    /// The hidden layers, in forward order.
92    pub hidden: Vec<M::HiddenLayer>,
93    /// The output layer.
94    pub output: M::OutputLayer,
95}
burn_core/module/parallel/base.rs

burn_core/module/parallel/
base.rs
