Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for LoRA adapters trained with Rank-Stabilized scaling #299

Merged
merged 4 commits into from
Mar 5, 2024

Conversation

arnavgarg1
Copy link
Contributor

@arnavgarg1 arnavgarg1 commented Mar 4, 2024

In PEFT 0.9, support was added for adapters trained using a new flag called use_rslora.

When set to True, Rank-Stabilized LoRA sets the adapter scaling factor to lora_alpha/math.sqrt(r), since it was proven to work better. Otherwise, it will use the original default value of lora_alpha/r.

In equation form:

  • Normal LoRA layers: W0X + (lora_alpha/r)(BAX) where W0 is the base model, BA are the lora weight matrices and X is the input from the embedding layer/previous transformer layer.
  • LoRA layers with RSLoRA enabled: W0X + (lora_alpha/sqrt(r))(BAX)

In particular, this is useful when using larger ranks since it prevents the gradient from collapsing as rank increases, which may result in higher ranks actually leading to better performance (not true by default today and in the original LoRA paper). Paper: https://arxiv.org/pdf/2312.03732.pdf.

image

@arnavgarg1 arnavgarg1 changed the title Add support for Rank-Stabilized LoRA adapters Add support for LoRA adapters trained with Rank-Stabilized scaling Mar 4, 2024
lora_a, lora_a_name = module_map[weight_name]["lora_A"]
lora_a = lora_a.to(base_device, self.dtype)

lora_b, lora_b_name = module_map[weight_name]["lora_B"]
lora_b = lora_b.to(base_device, self.dtype)

scale = adapter_config.lora_alpha / adapter_config.r
scale: float = get_scaling_factor(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we remove : float here?

@arnavgarg1 arnavgarg1 merged commit 21631fa into main Mar 5, 2024
1 check passed
@arnavgarg1 arnavgarg1 deleted the rslora branch March 5, 2024 00:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants