Linear method without quantization. Args: separate_bias_add: If true, add bias separately after matrix multiplication.
| 38 | |
| 39 | |
| 40 | class UnquantizedLinearMethod(LinearMethodBase): |
| 41 | """Linear method without quantization. |
| 42 | |
| 43 | Args: |
| 44 | separate_bias_add: If true, add bias separately after matrix |
| 45 | multiplication. |
| 46 | """ |
| 47 | |
| 48 | def __init__(self, separate_bias_add: bool = False): |
| 49 | self.separate_bias_add = separate_bias_add |
| 50 | |
| 51 | def create_weights(self, input_size_per_partition: int, |
| 52 | output_size_per_partition: int, input_size: int, |
| 53 | output_size: int, |
| 54 | params_dtype: torch.dtype) -> Dict[str, Any]: |
| 55 | weight = Parameter(torch.empty(output_size_per_partition, |
| 56 | input_size_per_partition, |
| 57 | device=torch.cuda.current_device(), |
| 58 | dtype=params_dtype), |
| 59 | requires_grad=False) |
| 60 | set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) |
| 61 | return {"weight": weight} |
| 62 | |
| 63 | def apply_weights(self, |
| 64 | weights: Dict[str, torch.Tensor], |
| 65 | x: torch.Tensor, |
| 66 | bias: Optional[torch.Tensor] = None) -> torch.Tensor: |
| 67 | weight = weights["weight"] |
| 68 | if self.separate_bias_add: |
| 69 | if bias: |
| 70 | return F.linear(x, weight) + bias |
| 71 | return F.linear(x, weight) |
| 72 | return F.linear(x, weight, bias) |
| 73 | |
| 74 | |
| 75 | class ReplicatedLinear(torch.nn.Module): |