Trait LayerParallelism

Source

pub trait LayerParallelism: Module {
    type InputLayer: DistributedLayer<Input: Module, Output: Module>;
    type HiddenLayer: DistributedLayer<Input = HiddenLayerSignal<Self>, Output = HiddenLayerSignal<Self>>;
    type OutputLayer: DistributedLayer<Input = HiddenLayerSignal<Self>>;

    // Required methods
    fn layer_input(&self) -> &Self::InputLayer;
    fn layers_hidden(&self) -> impl Iterator<Item = &Self::HiddenLayer>;
    fn layer_output(&self) -> &Self::OutputLayer;
    fn into_layers(self) -> LayerSequence<Self>;
}

Expand description

A model that runs as an input layer, then hidden layers in order, then an output layer, so its layers can be split across devices by a DistributedLayeredModel.

It is usually not the struct the model trained with, since a layer split across devices may be represented differently. Weights trained on another struct load into this one by remapping their keys. Each layer is built on the device a LayerPlacement gives it, so the model never exists whole on one device.

impl Model {
    fn new(config: &ModelConfig, placement: &LayerPlacement) -> Self {
        Self {
            embedding: Embedding::new(config, &placement.input),
            blocks: placement
                .hidden
                .iter()
                .map(|device| Block::new(config, device))
                .collect(),
            head: Head::new(config, &placement.output),
        }
    }
}

impl LayerParallelism for Model {
    type InputLayer = Embedding;
    type HiddenLayer = Block;
    type OutputLayer = Head;

    fn layer_input(&self) -> &Embedding {
        &self.embedding
    }

    fn layers_hidden(&self) -> impl Iterator<Item = &Block> {
        self.blocks.iter()
    }

    fn layer_output(&self) -> &Head {
        &self.head
    }

    fn into_layers(self) -> LayerSequence<Self> {
        LayerSequence {
            input: self.embedding,
            hidden: self.blocks,
            output: self.head,
        }
    }
}

let placement = LayerPlacement::even(&devices, config.num_blocks);
let model = Model::new(&config, &placement);
let model = DistributedLayeredModel::new(model, &placement);
let predictions = model.forward(features);