Struct FsdpState 

pub struct FsdpState { /* private fields */ }

FSDP state tracking for one training step.

Implementations

impl FsdpState

pub fn new(config: FsdpConfig) -> Self

Create a new FSDP state with the given configuration.

pub fn wrap_unit( &mut self, param_names: Vec<String>, param_values: HashMap<String, Vec<f64>>, ) -> Result<usize, FsdpError>

Register a group of parameters as one FSDP unit.

Returns the unit_id of the newly created unit.

pub fn allgather_unit(&mut self, unit_id: usize) -> Result<Vec<f64>, FsdpError>

AllGather parameters for a unit (before forward pass).

Simulates reconstruction by tiling the shard world_size times, then truncating to total_params length.

pub fn discard_unit_params(&mut self, unit_id: usize) -> Result<(), FsdpError>

Discard gathered parameters for a unit (after forward pass, to save memory).

pub fn reduce_scatter_grads( &mut self, unit_id: usize, grads: Vec<f64>, ) -> Result<Vec<f64>, FsdpError>

ReduceScatter gradients for a unit (after backward pass).

Simulates by averaging grads / world_size and returning the local shard.

pub fn memory_saving_ratio(&self) -> f64

Total memory saved vs DDP (replicated params) as a ratio.

Returns world_size since each rank holds 1/world_size of total params.

pub fn per_rank_memory_bytes(&self) -> usize

Per-rank memory footprint in bytes across all units.

pub fn unit_count(&self) -> usize

Number of FSDP units registered.

pub fn total_params(&self) -> usize

Total parameter count across all units.