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

lora merge fix for O2 names #7325

Merged
merged 7 commits into from
Sep 1, 2023
Merged

lora merge fix for O2 names #7325

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
cd7e683
wip
arendu Aug 25, 2023
1769113
adjust key names based on O2
arendu Aug 25, 2023
e5dc831
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 25, 2023
3a336ae
update
arendu Aug 25, 2023
9db9fe9
fixed conflicts
arendu Aug 25, 2023
94dd420
minor
arendu Aug 25, 2023
3a419f8
Merge branch 'main' into adithyare/merge_fix_for_o2
arendu Aug 25, 2023
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
12 changes: 10 additions & 2 deletions scripts/nlp_language_modeling/merge_lora_weights/merge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,13 @@ def load_lora(lora_nemo, tp):
return lora_state_dict


def fix_for_O2(state_dict):
new_state_dict = {}
for k, v in state_dict.items():
new_state_dict[k.replace('model.language_model', 'model.module.language_model')] = v
return new_state_dict


def merge(
base_model_state_dict: Dict[str, Any], lora_state_dict: Dict[int, Any], tp: int, num_layers: int, curr_rank: int
):
Expand Down Expand Up @@ -103,6 +110,7 @@ def merge(
wt_self_attn = base_model_state_dict[key_self_attn_kqv]
wt_lora = wt_lora_out @ wt_lora_in
base_model_state_dict[key_self_attn_kqv] = wt_self_attn + wt_lora.type_as(wt_self_attn)
print("mergeing for weight", key_self_attn_kqv)
return base_model_state_dict


Expand Down Expand Up @@ -148,8 +156,6 @@ def main(cfg) -> None:
pretrained_cfg.activations_checkpoint_granularity = None
pretrained_cfg.activations_checkpoint_method = None
pretrained_cfg.precision = trainer.precision
if trainer.precision == "16":
pretrained_cfg.megatron_amp_O2 = False
model = MegatronGPTModel.restore_from(
restore_path=cfg.gpt_model_file,
trainer=trainer,
Expand Down Expand Up @@ -199,6 +205,8 @@ def main(cfg) -> None:
)

# load the merged_weights back into the base model, for this current rank.
if model.cfg.megatron_amp_O2:
merged_weights = fix_for_O2(merged_weights)
model.load_state_dict(merged_weights)

# Going to go through the motions of inference to force PTL to run subprocess for loading all base model's ranks.
Expand Down