(self,
input_dim: int,
output_dim: int,
bias: bool = False,
lora_config: LoRAConfig = None,
quantization_config: QuantizationConfig = None,
device=None,
dtype=torch.bfloat16,
linear_cls=nn.Linear)
| 76 | class LoRAOptimizedLinear(nn.Module): |
| 77 | |
| 78 | def __init__(self, |
| 79 | input_dim: int, |
| 80 | output_dim: int, |
| 81 | bias: bool = False, |
| 82 | lora_config: LoRAConfig = None, |
| 83 | quantization_config: QuantizationConfig = None, |
| 84 | device=None, |
| 85 | dtype=torch.bfloat16, |
| 86 | linear_cls=nn.Linear): |
| 87 | super().__init__() |
| 88 | self.input_dim = input_dim |
| 89 | self.output_dim = output_dim |
| 90 | self.bias = bias |
| 91 | self.lora_config = lora_config |
| 92 | self.quantization_config = quantization_config |
| 93 | self.device = get_accelerator().current_device_name() if device is None else device |
| 94 | self.linear_cls = linear_cls |
| 95 | self.dtype = dtype |
| 96 | assert self.lora_config is not None, "DSOptimizedLinear requires a LoRA config" |
| 97 | assert not self.bias, "bias=True is not supported by LoRAOptimizedLinear" |
| 98 | self.zero_shards = self.lora_config.base_weight_sharding |
| 99 | self.sharded_weight_size = int(float(self.input_dim) // self.zero_shards) |
| 100 | if self.zero_shards > 1: |
| 101 | assert self.zero_shards == dist.get_world_size( |
| 102 | ), "base weight sharding is only supported across world size" |
| 103 | w = torch.nn.Parameter(torch.empty(self.output_dim * self.sharded_weight_size, dtype=dtype), |
| 104 | requires_grad=False) |
| 105 | else: |
| 106 | w = torch.nn.Parameter(torch.empty((self.output_dim, self.input_dim), dtype=dtype), requires_grad=False) |
| 107 | torch.nn.init.xavier_uniform_(w.reshape(self.sharded_weight_size, self.output_dim)) |
| 108 | |
| 109 | if self.quantization_config is not None: |
| 110 | assert dtype == torch.bfloat16, "only bfloat16 is supported when using quantization" |
| 111 | self.weight = QuantizedParameter(w, quantization_config=quantization_config) |
| 112 | else: |
| 113 | self.weight = w |
| 114 | |
| 115 | self.disabled = False |
| 116 | self._initialized = False |
| 117 | if not self.lora_config.delay_lora_init: |
| 118 | self.init_lora() |
| 119 | |
| 120 | def disable(self): |
| 121 | self.disabled = True |
nothing calls this directly
no test coverage detected