mstar.model.components.linear

mstar.model.components.linear#

Non-distributed fused linear projections.

FusedColumnLinear concatenates several output shards into a single weight (and optional bias) so a transformer’s QKV or gate/up projections run as one GEMM instead of several. Each shard is loaded independently via a weight_loader so HF checkpoints that keep the shards as separate tensors (q/k/v, gate/up) load straight into the fused parameter.

This mirrors the role of the TP-aware QKVParallelLinear / MergedColumnParallelLinear in model.components.distributed but without any tensor-parallel sharding — use it for models that fuse projections but don’t need TP.

Classes

FusedColumnLinear(input_size, shard_sizes[, ...])

Linear whose output is the concatenation of several shards along dim 0, fused into a single weight (and optional bias).
class mstar.model.components.linear.FusedColumnLinear(input_size, shard_sizes, bias=False, dtype=None)[source]#

Bases: Module

Linear whose output is the concatenation of several shards along dim 0, fused into a single weight (and optional bias).

Parameters:

  • input_size (int) – input feature dim.

  • shard_sizes (dict[str | int, int]) – maps a shard id to its output size. Shards are laid out along dim 0 in iteration order, so the order of this dict defines the layout (and the split used at forward time).

  • bias (bool) – whether to include a (fused) bias.

  • dtype (torch.dtype | None)

The fused weight / bias carry a weight_loader(param, tensor, shard_id) method that copies one checkpoint shard into its slice of the fused parameter, dispatched by shard_id (a key of shard_sizes). Wire the per-shard checkpoint keys to it with the loader’s stacked-param rules.

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

weight_loader(param, loaded_weight, loaded_shard_id=None)[source]#
Parameters:
Return type:

None