Struct ModelParallel 

pub struct ModelParallel { /* private fields */ }

Model parallel training

Splits the model across multiple GPUs. Different layers run on different devices.

Implementations

impl ModelParallel

pub fn new(config: DistributedConfig) -> Self

pub fn place_layer(&mut self, layer_name: &str, device_id: usize)

Assign a layer to a specific device

pub fn get_device(&self, layer_name: &str) -> Option<usize>

Get device for a layer

pub fn auto_place_layers(&mut self, layer_names: &[String], num_devices: usize)

Automatic layer placement using a simple strategy

pub fn transfer( &self, tensor: &Tensor, _from_device: usize, _to_device: usize, ) -> Tensor

Transfer tensor between devices