Struct FsdpUnit 

pub struct FsdpUnit {
    pub unit_id: usize,
    pub param_names: Vec<String>,
    pub total_params: usize,
    pub shard_size: usize,
    pub is_offloaded: bool,
}

One FSDP unit — wraps a group of parameters that are collectively sharded.

Fields

unit_id: usize

Unique identifier for this unit

param_names: Vec<String>

Names of parameters in this unit

total_params: usize

Total number of parameters across all ranks

shard_size: usize

Number of parameters this rank holds (shard_size = ceil(total/world_size))

is_offloaded: bool

Whether the shard is currently offloaded to CPU

Implementations

impl FsdpUnit

pub fn new( unit_id: usize, param_names: Vec<String>, total_params: usize, world_size: usize, ) -> Self

Create a new FSDP unit.

pub fn memory_bytes_per_rank(&self) -> usize

Memory in bytes used by this rank’s shard.

Uses 8 bytes per parameter (f64).

pub fn local_params(&self) -> usize

Number of parameters held locally by this rank (= shard_size).

pub fn all_gather_buffer_size(&self) -> usize

Size of the AllGather output buffer (= total_params — full reconstruction needed for forward pass).

pub fn reduce_scatter_buffer_size(&self) -> usize

Size of the ReduceScatter input/output buffer (= total_params).

Trait Implementations

impl Clone for FsdpUnit

fn clone(&self) -> FsdpUnit

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for FsdpUnit

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

Formats the value using the given formatter.