Skip to content

throw MegatronOptimizerModule warning only with mcore models #12085

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

Merged
akoumpa merged 2 commits into from
Feb 6, 2025
Merged

throw MegatronOptimizerModule warning only with mcore models #12085

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5f29d8b
throw MegatronOptimizerModule warning only with mcore models
akoumpa Feb 6, 2025
1fabae2
Apply isort and black reformatting
akoumpa Feb 6, 2025
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
13 changes: 8 additions & 5 deletions nemo/lightning/pytorch/callbacks/peft.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
from nemo.lightning.megatron_parallel import MegatronParallel
from nemo.lightning.pytorch.callbacks.model_transform import ModelTransform
from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule
from nemo.lightning.pytorch.utils import is_trainer_attached
from nemo.lightning.pytorch.utils import get_automodel_from_trainer, is_trainer_attached
from nemo.utils import logging
from nemo.utils.callbacks.dist_ckpt_io import AsyncCompatibleCheckpointIO

Expand Down Expand Up @@ -221,10 +221,12 @@ def apply_transform(self, trainer):
cb.on_fit_start(trainer, trainer.lightning_module)
break
else:
logging.warning(
"MegatronOptimizerModule not found in trainer callbacks. finalize_model_grads is not "
"properly set up for PEFT."
)
# i.e., this is an mcore model; elif not supported here.
if get_automodel_from_trainer(trainer) is None:
logging.warning(
"MegatronOptimizerModule not found in trainer callbacks. finalize_model_grads is not "
"properly set up for PEFT."
)

def adapter_key_filter(self, key: str) -> bool:
"""
Expand Down Expand Up @@ -431,6 +433,7 @@ def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_optio

def _create_lora_hf_config(self):
from peft import LoraConfig

from nemo.collections.llm.peft import DoRA

lora_config = LoraConfig(
Expand Down
Loading