mstar.model.components.distributed.linear

Classes

ColumnParallelLinear(comm_group, input_size, ...)	Linear layer with column parallelism.
MergedColumnParallelLinear(comm_group, ...)	Packed linear layers with column parallelism.
QKVParallelLinear(comm_group, hidden_size, ...)	Linear layers for the attention's QKV transformation.
RowParallelLinear(comm_group, input_size, ...)	Linear layer with row parallelism.

class mstar.model.components.distributed.linear.ColumnParallelLinear(comm_group, input_size, output_size, bias=False, gather_output=False, skip_bias_add=False, dtype=None)

Bases: Module

Linear layer with column parallelism.

The linear layer is defined as Y = XA + b. A is parallelized along its second dimension as A = [A_1, …, A_p].

Parameters:

• comm_group (CommGroup)

• input_size (int)

• output_size (int)

• bias (bool)

• gather_output (bool)

• skip_bias_add (bool)

• dtype (dtype | None)

forward(input_)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor | tuple[Tensor, Parameter | None]

weight_loader(param, loaded_weight, loaded_shard_id=None)

Copy this rank’s slice of loaded_weight into param.

For plain ColumnParallelLinear there’s a single shard so loaded_shard_id must be None. Subclasses (MergedColumnParallelLinear, QKVParallelLinear) override to dispatch by shard id.

Parameters:

• param (Parameter) – the destination parameter (self.weight or self.bias).

• loaded_weight (Tensor) – the full-checkpoint tensor — shape (output_size, input_size) for weight, (output_size,) for bias.

• loaded_shard_id (int | str | None)

class mstar.model.components.distributed.linear.MergedColumnParallelLinear(comm_group, input_size, output_sizes, bias=True, gather_output=False, skip_bias_add=False, dtype=None)

Bases: ColumnParallelLinear

Packed linear layers with column parallelism.

Similar to ColumnParallelLinear, but the weight matrix is concatenated along the output dimension. When the weight matrix is loaded, the different partitions are sharded separately.

Parameters:

• comm_group (CommGroup)

• input_size (int)

• output_sizes (list[int])

• bias (bool)

• gather_output (bool)

• skip_bias_add (bool)

• dtype (dtype | None)

weight_loader(param, loaded_weight, loaded_shard_id=None)

Copy this rank’s slice of one of the merged sub-shards into param at the correct offset.

loaded_weight is the full checkpoint tensor for the sub-shard identified by loaded_shard_id (an index into output_sizes).

Parameters:

• param (Parameter)

• loaded_weight (Tensor)

• loaded_shard_id (int | None)

class mstar.model.components.distributed.linear.QKVParallelLinear(comm_group, hidden_size, head_size, total_num_heads, total_num_kv_heads=None, bias=True, skip_bias_add=False, dtype=None, v_head_size=None)

Bases: ColumnParallelLinear

Linear layers for the attention’s QKV transformation.

Linear layers for the linear transformation of the query, key, and value vectors in the attention layer. The weight matrix is concatenated along the output dimension. The layer is parallelized along the head dimension. When the number of key/value heads is smaller than the number of query heads (e.g., multi-query/grouped-query attention), the key/value head may be replicated while the query heads are partitioned.

Parameters:

• comm_group (CommGroup)

• hidden_size (int)

• head_size (int)

• total_num_heads (int)

• total_num_kv_heads (int | None)

• bias (bool)

• skip_bias_add (bool)

• dtype (dtype | None)

• v_head_size (int | None)

weight_loader(param, loaded_weight, loaded_shard_id=None)

Copy this rank’s q / k / v slice into param at the right offset within the merged qkv parameter.

loaded_shard_id must be one of "q", "k", or "v". For GQA where tp_size > total_num_kv_heads the K / V heads are replicated across ranks (each rank in the same KV-replica group loads the same KV head).

Parameters:

• param (Parameter)

• loaded_weight (Tensor)

• loaded_shard_id (str | None)

class mstar.model.components.distributed.linear.RowParallelLinear(comm_group, input_size, output_size, bias=True, input_is_parallel=True, skip_bias_add=False, dtype=None, reduce_results=True)

Bases: Module

Linear layer with row parallelism.

The linear layer is defined as Y = XA + b. A is parallelized along its first dimension and X along its second dimension as:

A_1 |

. |

A = | . | X = [X_1, …, X_p]

. |

A_p | - -

Parameters:

• comm_group (CommGroup)

• input_size (int)

• output_size (int)

• bias (bool)

• input_is_parallel (bool)

• skip_bias_add (bool)

• dtype (dtype | None)

• reduce_results (bool)

forward(input_)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor | tuple[Tensor, Parameter | None]

weight_loader(param, loaded_weight, loaded_shard_id=None)

Copy this rank’s slice of loaded_weight into param.

Partition is along dim 1 (input dim) for the weight.

Parameters:

• param (Parameter)

• loaded_weight (Tensor)

• loaded_shard_id (int | str | None)