| 75 | |
| 76 | |
| 77 | class QuantizedColumnParallelLinear(ColumnParallelLinear): |
| 78 | def __init__( |
| 79 | self, |
| 80 | input_size: int, |
| 81 | output_size: int, |
| 82 | weight_bit_width: int, |
| 83 | weight: torch.Tensor = None, |
| 84 | bias: torch.Tensor = None, |
| 85 | *args, |
| 86 | **kwargs, |
| 87 | ): |
| 88 | super(QuantizedColumnParallelLinear, self).__init__(input_size, output_size, *args, **kwargs) |
| 89 | self.input_size = input_size |
| 90 | self.output_size = output_size |
| 91 | self.weight_bit_width = weight_bit_width |
| 92 | if "skip_bias_add" in kwargs: |
| 93 | self.skip_bias_add = kwargs["skip_bias_add"] |
| 94 | else: |
| 95 | self.skip_bias_add = False |
| 96 | del self.weight |
| 97 | |
| 98 | if weight is None: |
| 99 | self.weight = torch.empty( |
| 100 | self.output_size, self.input_size * weight_bit_width // 8, dtype=torch.int8, device=kwargs["device"] |
| 101 | ) |
| 102 | self.weight_scale = torch.empty(self.output_size, dtype=kwargs["params_dtype"], device=kwargs["device"]) |
| 103 | else: |
| 104 | self.weight_scale = (weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).half() |
| 105 | self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8) |
| 106 | if weight_bit_width == 4: |
| 107 | self.weight = compress_int4_weight(self.weight) |
| 108 | |
| 109 | if bias is None: |
| 110 | self.register_parameter('bias', None) |
| 111 | else: |
| 112 | del self.bias |
| 113 | self.bias = bias |
| 114 | |
| 115 | self.weight = Parameter(self.weight.to(kwargs["device"]), requires_grad=False) |
| 116 | self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False) |
| 117 | |
| 118 | def forward(self, input_): |
| 119 | # Set up backprop all-reduce. |
| 120 | input_parallel = copy_to_tensor_model_parallel_region(input_) |
| 121 | # Matrix multiply. |
| 122 | output_parallel = W8A16Linear.apply(input_parallel, self.weight, self.weight_scale, self.weight_bit_width) |
| 123 | if self.bias is not None and not self.skip_bias_add: |
| 124 | output_parallel = output_parallel + self.bias |
| 125 | if self.gather_output: |
| 126 | # All-gather across the partitions. |
| 127 | output = gather_from_tensor_model_parallel_region(output_parallel) |
| 128 | else: |
| 129 | output = output_parallel |
| 130 | |
| 131 | output_bias = self.bias if self.skip_bias_add else None |
| 132 | |
| 133 | return output, output_bias |
| 134 | |