Struct ZeRoConfig 

pub struct ZeRoConfig {
    pub stage: ZeRoStage,
    pub world_size: usize,
    pub rank: usize,
    pub cpu_offload: bool,
    pub nvme_offload: bool,
    pub overlap_comm: bool,
    pub bucket_size: usize,
}

ZeRO configuration

Fields

stage: ZeRoStage

ZeRO stage

world_size: usize

Number of processes/GPUs

rank: usize

Current process rank

cpu_offload: bool

Enable CPU offloading

nvme_offload: bool

Enable NVMe offloading

overlap_comm: bool

Overlap communication with computation

bucket_size: usize

Bucket size for gradient accumulation

Implementations

impl ZeRoConfig

pub fn stage1(world_size: usize, rank: usize) -> Self

Stage 1 configuration (optimizer state partitioning)

pub fn stage2(world_size: usize, rank: usize) -> Self

Stage 2 configuration (gradient partitioning)

pub fn stage3(world_size: usize, rank: usize) -> Self

Stage 3 configuration (parameter partitioning)

pub fn with_offload(self, cpu: bool, nvme: bool) -> Self

ZeRO-Offload configuration

Trait Implementations

impl Clone for ZeRoConfig

fn clone(&self) -> ZeRoConfig

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for ZeRoConfig

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

Formats the value using the given formatter.

impl Default for ZeRoConfig

fn default() -> Self

Returns the “default value” for a type.