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. Args: input_size: input dimension of the linea
| 220 | |
| 221 | |
| 222 | class MergedColumnParallelLinear(ColumnParallelLinear): |
| 223 | """Packed linear layers with column parallelism. |
| 224 | |
| 225 | Similar to ColumnParallelLinear, but the weight matrix is concatenated |
| 226 | along the output dimension. When the weight matrix is loaded, the |
| 227 | different partitions are sharded separately. |
| 228 | |
| 229 | Args: |
| 230 | input_size: input dimension of the linear layer. |
| 231 | output_sizes: list of output dimensions of the linear layer. |
| 232 | bias: If true, add bias. |
| 233 | gather_output: If true, call all-gather on output and make the output |
| 234 | available to all GPUs, otherwise, every GPU will have |
| 235 | its own output. |
| 236 | skip_bias_add: This was added to enable performance optimizations where |
| 237 | bias can be fused with other element-wise operations. we |
| 238 | skip adding bias but instead return it. |
| 239 | params_dtype: Data type for the parameters. |
| 240 | linear_method: (Maybe quantized) linear method. |
| 241 | """ |
| 242 | |
| 243 | def __init__( |
| 244 | self, |
| 245 | input_size: int, |
| 246 | output_sizes: List[int], |
| 247 | bias: bool = True, |
| 248 | gather_output: bool = False, |
| 249 | skip_bias_add: bool = False, |
| 250 | params_dtype: Optional[torch.dtype] = None, |
| 251 | linear_method: Optional[LinearMethodBase] = None, |
| 252 | ): |
| 253 | self.output_sizes = output_sizes |
| 254 | tp_size = get_tensor_model_parallel_world_size() |
| 255 | assert all(output_size % tp_size == 0 for output_size in output_sizes) |
| 256 | super().__init__(input_size, sum(output_sizes), bias, gather_output, |
| 257 | skip_bias_add, params_dtype, linear_method) |
| 258 | |
| 259 | def weight_loader(self, |
| 260 | param: Parameter, |
| 261 | loaded_weight: torch.Tensor, |
| 262 | loaded_shard_id: Optional[int] = None): |
| 263 | param_data = param.data |
| 264 | output_dim = getattr(param, "output_dim", None) |
| 265 | if loaded_shard_id is None: |
| 266 | # Loaded weight is already packed. |
| 267 | if output_dim is None: |
| 268 | assert param_data.shape == loaded_weight.shape |
| 269 | param_data.copy_(loaded_weight) |
| 270 | return |
| 271 | current_shard_offset = 0 |
| 272 | shard_offsets = [] |
| 273 | for i, output_size in enumerate(self.output_sizes): |
| 274 | shard_offsets.append((i, current_shard_offset, output_size)) |
| 275 | current_shard_offset += output_size |
| 276 | packed_dim = getattr(param, "packed_dim", None) |
| 277 | for shard_id, shard_offset, shard_size in shard_offsets: |
| 278 | # If quantized, we need to adjust the offset and size to account |
| 279 | # for the packing. |