Struct ZeroShardMap 

pub struct ZeroShardMap {
    pub block_owners: Vec<usize>,
    pub lm_head_owner: usize,
    pub final_norm_owner: usize,
    pub embedding_owner: usize,
    pub world_size: usize,
}

Block-level optimizer shard map.

Maps transformer blocks to worker ranks. In ZeRO-1, each block’s optimizer state is owned by exactly one worker. The owner runs the optimizer step for that block after gradient AllReduce, then broadcasts updated weights.

Fields

block_owners: Vec<usize>

Which worker owns each block’s optimizer state. block_owners[i] = rank that owns block i.

lm_head_owner: usize

Which worker owns the LM head optimizer state.

final_norm_owner: usize

Which worker owns the final norm optimizer state.

embedding_owner: usize

Which worker owns the embedding optimizer state.

world_size: usize

World size

Implementations

impl ZeroShardMap

pub fn round_robin(num_blocks: usize, world_size: usize) -> Self

Create a shard map distributing blocks round-robin across workers.

Non-block components (LM head, final norm, embedding) are assigned to rank 0 by default since they’re relatively small.

pub fn contiguous(num_blocks: usize, world_size: usize) -> Self

Create a shard map distributing blocks in contiguous chunks.

Preferred for pipeline parallelism compatibility: worker 0 gets blocks 0..N/W, worker 1 gets N/W..2N/W, etc.

pub fn block_owner(&self, block_idx: usize) -> usize

Get the owning rank for a given block index.

pub fn rank_owns_block(&self, rank: usize, block_idx: usize) -> bool

Check if this rank owns a given block’s optimizer state.

pub fn blocks_for_rank(&self, rank: usize) -> Vec<usize>

Get all block indices owned by a given rank.

pub fn num_blocks_for_rank(&self, rank: usize) -> usize

Number of blocks owned by a given rank.

pub fn memory_fraction_for_rank(&self, rank: usize) -> f64

Memory savings for a given rank (ratio of blocks owned / total blocks).

Returns the fraction of optimizer memory this rank holds.

Trait Implementations

impl Clone for ZeroShardMap

fn clone(&self) -> ZeroShardMap

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for ZeroShardMap

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

Formats the value using the given formatter.