Optimized version of nn.Linear that adds features such as: * LoRA w. base weight sharding * FP [6,8,12] quantization Arguments: input_dim: Required: size of each input sample output_dim: Required: size of each output sample bias: Optional: If set to Fals
| 16 | |
| 17 | |
| 18 | class OptimizedLinear(nn.Module): |
| 19 | """ |
| 20 | Optimized version of nn.Linear that adds features such as: |
| 21 | * LoRA w. base weight sharding |
| 22 | * FP [6,8,12] quantization |
| 23 | |
| 24 | Arguments: |
| 25 | input_dim: Required: size of each input sample |
| 26 | output_dim: Required: size of each output sample |
| 27 | bias: Optional: If set to False, the layer will not learn an additive bias. Default: False |
| 28 | lora_config: Optional: LoRAConfig defining lora features and base-weight-sharding degree |
| 29 | quantization_config: Optional: QuantizationConfig defining quantization features |
| 30 | dtype: Optional: parameter dtype, only supports bfloat16 currently |
| 31 | |
| 32 | Returns: |
| 33 | Returns a new nn.Module depending on the input config. Either native |
| 34 | torch.nn.Linear, QuantizedLinear, or the full-featured DSOptimizedLinear. |
| 35 | """ |
| 36 | |
| 37 | def __new__(self, |
| 38 | input_dim: int, |
| 39 | output_dim: int, |
| 40 | bias: bool = False, |
| 41 | lora_config: LoRAConfig = None, |
| 42 | quantization_config: QuantizationConfig = None, |
| 43 | device=None, |
| 44 | dtype=torch.bfloat16, |
| 45 | linear_cls=nn.Linear): |
| 46 | |
| 47 | if quantization_config is not None and not is_dataclass(quantization_config): |
| 48 | raise ValueError(f"Expecting QuantizationConfig but received {type(quantization_config)}") |
| 49 | if lora_config is not None and not is_dataclass(lora_config): |
| 50 | raise ValueError(f"Expecting LoRAConfig but received {type(lora_config)}") |
| 51 | if lora_config is None and quantization_config is None: |
| 52 | # Everything disabled, fall back to normal nn.Linear |
| 53 | self = linear_cls(input_dim, output_dim, bias=bias, dtype=dtype, device=device) |
| 54 | |
| 55 | elif lora_config: |
| 56 | # lora enabled, quantization may or may not be |
| 57 | self = LoRAOptimizedLinear(input_dim=input_dim, |
| 58 | output_dim=output_dim, |
| 59 | bias=bias, |
| 60 | lora_config=lora_config, |
| 61 | quantization_config=quantization_config, |
| 62 | dtype=dtype, |
| 63 | device=device, |
| 64 | linear_cls=linear_cls) |
| 65 | |
| 66 | elif quantization_config: |
| 67 | # only quantization enabled, no lora |
| 68 | self = QuantizedLinear(input_dim=input_dim, |
| 69 | output_dim=output_dim, |
| 70 | bias=bias, |
| 71 | quantization_config=quantization_config, |
| 72 | dtype=dtype) |
| 73 | return self |
| 74 | |
| 75 |