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 / dense MLP 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.KVColumnParallelLinear(comm_group, input_size, head_size, total_num_kv_heads, bias=False, dtype=None)[source]#

Bases: ColumnParallelLinear

One K or V projection, sharded by head, replicated when heads run out.

For models whose checkpoint keeps k_proj / v_proj apart rather than fused into a QKV block. QKVParallelLinear already handles this case, but only for the fused layout.

The replication is not an optimization — it is what keeps the model and the KV cache agreeing. KVConfig.shard gives every rank one whole head once tp_size >= total_num_kv_heads, so a plain column shard, which would hand each rank a fraction of a head, would silently disagree with the cache it writes into.

Parameters:

  • comm_group (CommGroup)

  • input_size (int)

  • head_size (int)

  • total_num_kv_heads (int)

  • bias (bool)

  • dtype (dtype | 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.ParallelCrossAttention(*, comm_group=None, hidden_size, num_heads, num_kv_heads=None, head_dim, q_bias=True, k_bias=False, v_bias=True, o_bias=True, source='default', cross_key=None, context_kv_key=None)[source]#

Bases: Module

Cross-attention split across tensor-parallel ranks.

Each rank handles a subset of attention heads, then the output projection combines their contributions with an all-reduce.

Parameters:

  • comm_group (CommGroup | None)

  • hidden_size (int)

  • num_heads (int)

  • num_kv_heads (int | None)

  • head_dim (int)

  • q_bias (bool)

  • k_bias (bool)

  • v_bias (bool)

  • o_bias (bool)

  • source (str)

  • cross_key (str | None)

  • context_kv_key (str | None)

bind_resources(resources)[source]#

Resolve this source’s cross-attention resource and the cache holding its context. See NodeSubmodule.bind_node_resources.

Parameters:

resources (dict)

Return type:

None

bind_step(label)[source]#
Parameters:

label (str)

Return type:

None

compute_kv(encoder_states)[source]#

Project the encoder context to K/V for the cross-attention pool.

(enc_len, hidden) -> (k, v), each (enc_len, num_kv_heads, head_dim). Override to reshape for a model-specific pool layout.

Parameters:

encoder_states (Tensor)

Return type:

tuple[Tensor, Tensor]

forward(hidden_states)[source]#

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

Parameters:

hidden_states (Tensor)

Return type:

Tensor

set_layer_idx(layer_idx)[source]#
Parameters:

layer_idx (int)

Return type:

None

class mstar.model.components.distributed.ParallelGatedDeltaNet(*, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size, layout=GDNProjLayout.SPLIT, proj_bias=False, conv_bias=False, rms_norm_eps=1e-6, linear_attn_key='linear_attn', state_key='gdn_state', comm_group=None)[source]#

Bases: GatedDeltaNet

GatedDeltaNet over one rank’s heads.

Takes the checkpoint’s (total) head counts and divides them here.

Parameters:

  • hidden_size (int)

  • num_k_heads (int)

  • num_v_heads (int)

  • head_k_dim (int)

  • head_v_dim (int)

  • conv_kernel_size (int)

  • layout (GDNProjLayout)

  • proj_bias (bool)

  • conv_bias (bool)

  • rms_norm_eps (float)

  • linear_attn_key (str)

  • state_key (str)

  • comm_group (CommGroup | None)

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.ParallelGatedMLPUnfused(hidden_size, intermediate_size, comm_group=None, activation='silu', bias=False)[source]#

Bases: Module

SwiGLU MLP with separate (unfused) gate/up column shards.

Same math and sharding as ParallelGatedMLP, with gate_proj and up_proj kept as separate ColumnParallelLinear projections so the parameter names match a checkpoint’s gate_proj.weight / up_proj.weight one-to-one (plain name-matched loading, no stacked-parameter loader rules). A trivial comm group (world size 1) makes the projections plain linears.

Parameters:

  • hidden_size (int)

  • intermediate_size (int)

  • comm_group (CommGroup | None)

  • activation (str | Callable)

  • bias (bool)

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.ParallelMLP(hidden_size, intermediate_size, comm_group=None, activation='relu2', bias=False)[source]#

Bases: Module

Dense two-projection MLP down_proj(act(up_proj(x))) across TP ranks.

The parallel counterpart of mstar.model.components.MLP, named after the dense-LLM checkpoint convention (up_proj / down_proj) so a Nemotron-style FFN loads by name. activation accepts the HF names _resolve_activation knows (relu2 is the Nemotron squared ReLU) or a callable. A trivial comm group (world size 1) makes both projections plain linears.

Parameters:

  • hidden_size (int)

  • intermediate_size (int)

  • comm_group (CommGroup | None)

  • activation (str | Callable)

  • bias (bool)

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 self- and cross-attention.

embedding

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

linear

linear_attn

The gated delta-net mixer, sharded across tensor-parallel ranks.

mlp

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

sequence_parallel

Ulysses sequence parallelism — model-agnostic primitives.