Struct ParameterPartition 

pub struct ParameterPartition {
    pub name: String,
    pub local_shard: Tensor,
    pub partition_info: PartitionInfo,
    pub is_gathered: bool,
    pub full_parameter: Option<Tensor>,
}

Parameter partition for ZeRO Stage 3

Fields

name: String

Parameter name

local_shard: Tensor

Local shard of the parameter

partition_info: PartitionInfo

Partition metadata

is_gathered: bool

Whether this parameter is currently gathered

full_parameter: Option<Tensor>

Full parameter (only valid when is_gathered = true)

Implementations

impl ParameterPartition

pub fn new( name: String, local_shard: Tensor, partition_info: PartitionInfo, ) -> Self

pub fn memory_usage(&self) -> usize

Get memory usage of this partition

pub fn gather(&mut self, mp_context: &ModelParallelContext) -> Result<()>

Gather full parameter from all partitions

pub fn release(&mut self)

Release gathered parameter to save memory

Trait Implementations

impl Clone for ParameterPartition

fn clone(&self) -> ParameterPartition

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for ParameterPartition

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.