mstar.model.components.distributed.embedding

mstar.model.components.distributed.embedding#

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

Embedding weight [V, H] is row-sharded across TP ranks (each rank holds V/tp rows). Forward masks tokens outside the local rank’s slice to zero, runs F.embedding on the local shard, then sums across ranks with a single all_reduce. The output is [..., H] replicated on every rank — i.e. the LM stack’s subsequent layers see the same input they would in non-TP mode, no further coordination required.

Pairs with ColumnParallelLinear(gather_output=True) as the LM head: the head produces [B, V/tp] per rank and gather_output collects to [B, V] so the sampler stays vocab-oblivious.

Requires vocab_size divisible by tp_size. Padding-to-divisible is a follow-up.

Classes

VocabParallelEmbedding(num_embeddings, ...)

Row-parallel token embedding.
class mstar.model.components.distributed.embedding.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