mstar.model.components.distributed

mstar.model.components.distributed#

Tensor-parallel building blocks.

Parallel linears (ColumnParallelLinear, RowParallelLinear, MergedColumnParallelLinear, QKVParallelLinear), vocab-parallel embedding (VocabParallelEmbedding), and the composed parallel Attention / GatedMLP blocks. Each parallel parameter carries a weight_loader attribute used by the model-level weight loader to slice checkpoint tensors per-rank on load.

For vocab parallelism, pair VocabParallelEmbedding with ColumnParallelLinear(gather_output=True) on the LM head: the embedding all-reduces shard contributions before the first transformer layer, and the LM head all-gathers logits along the vocab dim before returning, so the sampler stays vocab-oblivious.

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

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_)[source]#

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)[source]#

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.MergedColumnParallelLinear(comm_group, input_size, output_sizes, bias=True, gather_output=False, skip_bias_add=False, dtype=None)[source]#

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)[source]#

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:
class mstar.model.components.distributed.ParallelAttention(*, comm_group=None, hidden_size, num_heads, num_kv_heads, head_dim, qkv_bias=False, o_bias=False, qk_norm=False, rms_norm_eps=1e-6, rope_theta=10_000.0, rope_scale=1.0, rope_low_freq_factor=1.0, rope_high_freq_factor=1.0, rope_old_context_len=8192, input_hidden_size=None, attn_key='attn', kv_key='kv', pos_key='rope')[source]#

Bases: Module

Parameters:

  • comm_group (CommGroup | None)

  • hidden_size (int)

  • num_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • qkv_bias (bool)

  • o_bias (bool)

  • qk_norm (bool)

  • rms_norm_eps (float)

  • rope_theta (float)

  • rope_scale (float)

  • rope_low_freq_factor (float)

  • rope_high_freq_factor (float)

  • rope_old_context_len (int)

  • input_hidden_size (int | None)

  • attn_key (str)

  • kv_key (str)

  • pos_key (str | None)

bind_resources(resources)[source]#

Resolve the resources this layer calls. See NodeSubmodule.bind_node_resources.

Parameters:

resources (dict)

Return type:

None

forward(hidden_states)[source]#

Label and layer index come off the resources’ cursors; see Attention.forward.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

class mstar.model.components.distributed.ParallelGatedMLP(hidden_size, intermediate_size, comm_group=None, activation='silu', bias=False)[source]#

Bases: Module

SwiGLU-style gated MLP partitioned across TP ranks.

Parameters:

  • comm_group (CommGroup | None) – TP comm group for this MLP’s parallel linears.

  • hidden_size (int) – model hidden dim (full, not per-partition).

  • intermediate_size (int) – SwiGLU intermediate dim (full).

  • activation (str | Callable) – HF activation name (silu, gelu, gelu_tanh).

  • bias (bool) – whether the linears have a bias term.

forward(x)[source]#

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.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.components.distributed.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)[source]#

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)[source]#

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:
class mstar.model.components.distributed.RowParallelLinear(comm_group, input_size, output_size, bias=True, input_is_parallel=True, skip_bias_add=False, dtype=None, reduce_results=True)[source]#

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_)[source]#

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)[source]#

Copy this rank’s slice of loaded_weight into param.

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

Parameters:
class mstar.model.components.distributed.VocabParallelEmbedding(num_embeddings, embedding_dim, comm_group=None, padding_idx=None, dtype=None)[source]#

Bases: Module

Row-parallel token embedding.

Each rank holds rows [tp_rank * V/tp : (tp_rank + 1) * V/tp] of the full [V, H] embedding matrix. The forward zeroes contributions for tokens outside this rank’s slice and all_reduce sums shards into the replicated full embedding.

Parameters:

  • num_embeddings (int)

  • embedding_dim (int)

  • comm_group (CommGroup | None)

  • padding_idx (int | None)

  • dtype (torch.dtype | None)

forward(input_)[source]#

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.

Parameters:

input_ (Tensor)

Return type:

Tensor

weight_loader(param, loaded_weight, loaded_shard_id=None)[source]#

Slice [V, H] HF embedding into this rank’s [V/tp, H].

Parameters:
Return type:

None

Modules

attention

TP-aware multi-head attention.

embedding

Vocabulary-parallel embedding for TP'd token LM stacks.

linear

mlp

TP-aware SwiGLU MLPs (parallel counterparts of mstar.model.components.GatedMLP).

sequence_parallel

Ulysses sequence parallelism — model-agnostic primitives.