Struct OptimizerShard 

pub struct OptimizerShard {
    pub rank: usize,
    pub world_size: usize,
    pub param_start: usize,
    pub param_end: usize,
    pub total_params: usize,
}

Optimizer state shard assignment for a single worker.

Each worker owns a contiguous range of parameter indices and holds only the Adam m/v states for those parameters.

Fields

rank: usize

This worker’s rank

world_size: usize

Total workers in the ZeRO group

param_start: usize

Start index (inclusive) in the flattened parameter vector

param_end: usize

End index (exclusive) in the flattened parameter vector

total_params: usize

Total parameter count across all shards

Implementations

impl OptimizerShard

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

Compute shard assignment for a given rank.

Divides total_params into world_size contiguous shards. Last shard absorbs any remainder.

Contract (C-ZERO-001)

• Union of all shards == [0, total_params)
• Intersection of any two shards == empty
• Each shard size within 1 of floor(total_params / world_size)

pub fn shard_size(&self) -> usize

Number of parameters in this shard.

pub fn owns_param(&self, param_idx: usize) -> bool

Check if a parameter index belongs to this shard.

pub fn memory_savings(&self) -> f64

Estimated memory savings ratio compared to full replication.

Returns fraction of memory saved (e.g., 0.5 for 2 workers).

pub fn shard_memory_bytes(&self) -> usize

Estimated optimizer memory for this shard in bytes.

AdamW stores m + v (2 × shard_size × sizeof(f32)).

pub fn full_memory_bytes(&self) -> usize

Full optimizer memory without sharding in bytes.

Trait Implementations

impl Clone for OptimizerShard

fn clone(&self) -> OptimizerShard

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for OptimizerShard

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

Formats the value using the given formatter.