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.

Sign up for GitHub

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

Make QLoRA more model-agnostic #9488

Merged
merged 2 commits into from
Jul 2, 2024
Merged

Make QLoRA more model-agnostic #9488

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
54ff538
qlora support more models
cuichenx Jun 17, 2024
af6f598
Merge branch 'main' into chcui/qlora2
cuichenx Jun 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 9 additions & 8 deletions nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
from megatron.core.models.common.embeddings.rotary_pos_embedding import apply_rotary_pos_emb
from megatron.core.tensor_parallel import ColumnParallelLinear
from megatron.core.transformer.attention import SelfAttention
from megatron.core.transformer.custom_layers.transformer_engine import SplitAlongDim
from megatron.core.transformer.mlp import MLP
Expand Down Expand Up @@ -305,14 +304,16 @@ def mcore_register_adapters(self):

def forward(self, hidden_states, expert_idx=None):
# [s, b, 4 * h/p]
if isinstance(self.linear_fc1, ColumnParallelLinear):
layernorm_output = hidden_states
intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states)
elif self.linear_fc1.te_return_bias:
intermediate_parallel, bias_parallel, layernorm_output = self.linear_fc1(hidden_states)
output = self.linear_fc1(hidden_states)
if isinstance(output, tuple) and len(output) == 2:
intermediate_parallel, bias_parallel = output
if isinstance(intermediate_parallel, tuple) and len(intermediate_parallel) == 2:
intermediate_parallel, layernorm_output = intermediate_parallel
else:
layernorm_output = hidden_states
else:
# bias_parallel is None
(intermediate_parallel, layernorm_output), bias_parallel = self.linear_fc1(hidden_states)
# self.linear_fc1.te_return_bias == True
intermediate_parallel, bias_parallel, layernorm_output = output

# LoRA logic
if self.is_adapter_available():
Expand Down
8 changes: 4 additions & 4 deletions nemo/collections/nlp/modules/common/megatron/adapters/qlora.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,12 +224,12 @@ def qlora_load_model(model: 'MCoreGPTModel', model_cfg: 'DictConfig', checkpoint
def replace_linear(module: nn.Module, prefix=""):
for name, child in module.named_children():
if name in qlora_targets:
bf16_weight = checkpoint[f"{prefix}.{name}.weight"]
bf16_weight = checkpoint[f"{prefix}.{name}.weight"].to(torch.bfloat16)
logging.info(f'QLoRA: Quantizing linear layer: {prefix}.{name}')
if name in ['linear_proj', 'linear_fc2']:
layer_norm_weight = checkpoint.get(f"{prefix}.{name}.layer_norm_weight", None)
if layer_norm_weight is None:
setattr(module, name, NF4LinearWrapper(bf16_weight))
else: # name in ['linear_qkv', 'linear_fc1']
layer_norm_weight = checkpoint[f"{prefix}.{name}.layer_norm_weight"]
else:
layer_norm_bias = checkpoint.get(f"{prefix}.{name}.layer_norm_bias", None)
normalization = module.config.normalization
zero_centered_gamma = module.config.layernorm_zero_centered_gamma
Expand Down
Loading