Struct TransformerParallelMapping 

pub struct TransformerParallelMapping {
    pub heads_per_rank: usize,
    pub kv_heads_per_rank: usize,
    pub head_dim: usize,
    pub hidden_per_rank: usize,
    pub intermediate_per_rank: usize,
}

Tensor parallel layer mapping for transformer models

Fields

heads_per_rank: usize

Number of attention heads per rank

kv_heads_per_rank: usize

Number of KV heads per rank

head_dim: usize

Head dimension

hidden_per_rank: usize

Hidden dimension per rank

intermediate_per_rank: usize

Intermediate dimension per rank (for MLP)

Implementations

impl TransformerParallelMapping

pub fn new( num_heads: usize, num_kv_heads: usize, head_dim: usize, hidden_dim: usize, intermediate_dim: usize, tp_size: usize, ) -> Result<Self>

Create mapping for a transformer model

pub fn q_proj_size(&self) -> usize

Get Q projection output dimension per rank

pub fn k_proj_size(&self) -> usize

Get K projection output dimension per rank

pub fn v_proj_size(&self) -> usize

Get V projection output dimension per rank

pub fn o_proj_in_size(&self) -> usize

Get O projection input dimension per rank

pub fn attention_weight_shards( &self, layer_idx: usize, config: &TensorParallelConfig, ) -> Vec<WeightShard>

Get weight shards for attention layer

pub fn mlp_weight_shards( &self, layer_idx: usize, config: &TensorParallelConfig, ) -> Vec<WeightShard>

Get weight shards for MLP layer

Trait Implementations

impl Clone for TransformerParallelMapping

fn clone(&self) -> TransformerParallelMapping

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for TransformerParallelMapping

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

Formats the value using the given formatter.