Class SubParamLinearLayer

deepspeed/module_inject/layers.py:1745–1880

Column-parallel linear layer with sub-parameter support. Handles cases where weights contain multiple logical sub-parameters that need to be partitioned separately (e.g., fused QKV, chunked MLP, GQA). The `shape` parameter controls how the weight is viewed and partitioned: - (

class SubParamLinearLayer(TensorParallel_Layer):
    """
    Column-parallel linear layer with sub-parameter support.

    Handles cases where weights contain multiple logical sub-parameters
    that need to be partitioned separately (e.g., fused QKV, chunked MLP, GQA).

    The `shape` parameter controls how the weight is viewed and partitioned:
    - (3, -1) with partition_dim=0: 3 equal sub-params, partition each at dim 0
    - ((q, k, v), -1) with partition_dim=0: 3 unequal sub-params (1-level nesting)
    """

    def __init__(self, module, mp_group, shape, partition_dim=0, **kwargs):
        super(SubParamLinearLayer, self).__init__(mp_group, **kwargs)
        self.weight = module.weight
        self.bias = module.bias
        self.shape = shape
        self.partition_dim = partition_dim

        self._orig_weight_shape = tuple(module.weight.shape)
        self._orig_bias_shape = tuple(module.bias.shape) if self.bias is not None else None
        (self._logical_shape, self._output_shape, self._subparam_sizes,
         self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
                                                                    self.partition_dim, self.name)
        # Resolve the per-rank widths once, for the same reason _freeze_partition_sizes does:
        # the split depends on this model's tp_meta.
        self._subparam_shard_widths = _subparam_shard_widths(
            self._subparam_sizes or (self._logical_shape[self.partition_dim], ), self.tp_world_size, self.tp_meta,
            self.name)
        self._bias_shape_spec = _bias_subparam_shape_spec(self._output_shape, self._bias_partition_dim,
                                                          self._subparam_sizes)
        if self.bias is not None and self.bias.numel() != _shape_prod(self._output_shape):
            raise ValueError(f"AutoTP layer '{self.name}' bias size {self.bias.numel()} does not match output shape "
                             f"{self._output_shape}.")

        if self._should_materialize_tp_partition():
            self._tp_partition([self.weight, self.bias])
        self.support_training = True
        self.config_tp_params(self.weight)
        if self.bias is not None:
            self.config_tp_params(self.bias)
        self._mark_uc_metadata()

    def forward(self, input):
        self._assert_compiled_if_deferred()
        if getattr(self, 'mp_group', None) is not None and not self.defer_collectives_to_compiler:
            input = ColumnParallel.apply(self.mp_group, input)
        output = torch.matmul(input, self.weight.transpose(-1, -2))
        if self.bias is not None:
            output = add_bias(output, self.bias)
        return output

    @torch.no_grad()
    def gather_params(self, params_list):
        """Gather partitioned parameters back to full size."""
        for idx, param in enumerate(params_list):
            if param is None:
                continue