Repository: microsoft/DeepSpeed
File: deepspeed/linear/optimized_linear.py (Python)

Code before:
        self.lora_config = lora_config
        self.quantization_config = quantization_config
        self.device = get_accelerator().current_device_name() if device is None else device
        assert self.lora_config is not None, "DSOptimizedLinear requires a LoRA config"

        self.zero_shards = self.lora_config.base_weight_sharding
        self.sharded_weight_size = int(float(self.input_dim) // self.zero_shards)
        if self.zero_shards > 1:
            w = torch.nn.Parameter(torch.empty(self.output_dim * self.sharded_weight_size, dtype=dtype), requires_grad=False)
        else:
            w = torch.nn.Parameter(torch.empty((self.output_dim, self.input_dim), dtype=dtype), requires_grad=False)
        torch.nn.init.xavier_uniform_(w)

        if self.quantization_config is not None:
            assert dtype == torch.bfloat16, "only bfloat16 is supported when using quantization"
            self.weight = QuantizedParameter(w, quantization_config=quantization_config)
        else:
            self.weight = w

        self.weight.requires_grad = False

        # Mark base weight to prevent broadcast and ensure proper offload behavior
        self.weight.ds_optim_param = True

        # Use RS lora for now.
        self.lora_scaling_factor = self.lora_config.lora_alpha / self.lora_config.lora_r
        # Keeping lora weights in bf16 precision for ease of training.
        self.lora_weight_1 = nn.Linear(self.input_dim,
                                       self.lora_config.lora_r,
                                       bias=self.bias,
                                       device=self.device,
                                       dtype=dtype)
        self.lora_weight_2 = nn.Linear(self.lora_config.lora_r,
                                       self.output_dim,
                                       bias=self.bias,
                                       device=self.device,
                                       dtype=dtype)
        
        # initialize "A" with kaiming uniform and "B" with zeros following this
        # https://github.com/huggingface/peft/blob/62122b5add8d6892f70c82eaef2147a6ba33b90b/src/peft/tuners/lora/layer.py#L155
        nn.init.kaiming_uniform_(self.lora_weight_1.weight, a=math.sqrt(5))
        nn.init.zeros_(self.lora_weight_2.weight)
        self.lora_weight_1.weight.requires_grad = True
        self.lora_weight_2.weight.requires_grad = True

    def full_weight(self):
        # This assumes weights are evenly sharded across gpus. which might not be correct.
        # in that case, we should flatten before all_gather.
        base_weight = self.weight
        if getattr(base_weight, 'ds_offload', False):
            # move to gpu so we can dequant and all-gather

Review comment (suggestion):
You may want to remove `# Use RS lora for now.` comment as the `math.sqrt(...)` that was removed is the rs-lora technique.

Code after:
                 linear_cls=nn.Linear):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.bias = bias
        self.lora_config = lora_config
        self.quantization_config = quantization_config
        self.device = get_accelerator().current_device_name() if device is None else device
        self.linear_cls = linear_cls
        self.dtype = dtype
        assert self.lora_config is not None, "DSOptimizedLinear requires a LoRA config"
        assert not self.bias, "bias=True is not supported by LoRAOptimizedLinear"
        self.zero_shards = self.lora_config.base_weight_sharding
        self.sharded_weight_size = int(float(self.input_dim) // self.zero_shards)
        if self.zero_shards > 1:
            assert self.zero_shards == dist.get_world_size(
            ), "base weight sharding is only supported across world size"
            w = torch.nn.Parameter(torch.empty(self.output_dim * self.sharded_weight_size, dtype=dtype),
                                   requires_grad=False)
        else:
            w = torch.nn.Parameter(torch.empty((self.output_dim, self.input_dim), dtype=dtype), requires_grad=False)
        torch.nn.init.xavier_uniform_(w.reshape(self.sharded_weight_size, self.output_dim))

        if self.quantization_config is not None:
            assert dtype == torch.bfloat16, "only bfloat16 is supported when using quantization"
            self.weight = QuantizedParameter(w, quantization_config=quantization_config)
        else:
            self.weight = w

        self.disabled = False
        self._initialized = False
        if not self.lora_config.delay_lora_init:
            self.init_lora()

    def disable(self):
        self.disabled = True
        self.weight = torch.nn.Parameter(torch.empty((self.output_dim, self.input_dim), dtype=self.dtype),
                                         requires_grad=False)

    def init_lora(self):
        if self.disabled:
            return

        if self.quantization_config is not None:
            # ensure quant-param wasn't stripped, in some cases transformers will do this during model init
            if not isinstance(self.weight, QuantizedParameter):
                self.weight = QuantizedParameter(self.weight, quantization_config=self.quantization_config)

        self._initialized = True
        self.weight.requires_grad = False