feat(training): add LoRA fine-tuning strategy by ys-teh · Pull Request #149 · NVIDIA/nvalchemi-toolkit

ys-teh · 2026-07-24T21:22:16Z

ALCHEMI Toolkit Pull Request

Description

Adds LoRA-based parameter-efficient fine-tuning support within FineTuningStrategy for pretrained MLIP models. The implementation supports peft_config=LoRAConfig(...), equivariant LoRA wrappers for e3nn/cuEquivariance layers, trainable-state PEFT checkpoint loading, base-model fingerprint validation, and a MACE LoRA fine-tuning example.

FSDP support is out of scope for this PR. CLI support will be added in a separate PR.

Type of Change

Bug fix (non-breaking change that fixes an issue)
New feature (non-breaking change that adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Performance improvement
Documentation update
Refactoring (no functional changes)
CI/CD or infrastructure change

Related Issues

N/A

Changes Made

Updated existing training/checkpoint classes to support PEFT fine-tuning:
- Added peft_config support to FineTuningStrategy, with LoRA as the currently supported PEFT method.
- Added trainable and managed parameter registration to FineTuningStrategy.
- Updated module patch behavior so patched modules are trainable by default and do not need to be repeated in trainable_patterns.
- Added save_trainable_state_only support to TrainingStrategy.save_checkpoint and CheckpointHook for parameter-efficient checkpoint saving.
- Updated checkpoint restore/loading to support trainable-state PEFT checkpoints and partial model state loading.
- Updated ModulePatchHook and TrainableParameterHook to use registered trainable/managed parameter names.
- Added FineTuningSummaryHook to provide a concise report on trainable parameters.
- Added base fingerprint metadata and compatibility checks for loading PEFT checkpoints into base models.
- Fixed DemoModelWrapper.from_checkpoint / export_model full-model round-tripping.
Added LoRA PEFT support for workflows with LoRAConfig, fine-tuning serialization, checkpoint validation, PEFT checkpoint loading, and merge-into-base support.
Added LoRA PEFT support:
- Added LoRAConfig support for configuring LoRA adapter injection through FineTuningStrategy.
- Added merge-into-base support for LoRA workflows.
- Added reusable LoRA registration hooks and built-in wrapper registration.
- Added LoRA wrappers for e3nn.o3.Linear, e3nn fully connected layers, and cuEquivariance linear layers.
Added fine-tuning docs and examples/intermediate/08_lora_finetuning.py for MACE medium-mpa-0 LoRA fine-tuning on LPSC data.

Testing

Targeted tests run:

uv run --extra cu12 pytest test/training/test_peft_lora.py test/training/test_finetune.py test/training/test_reference_energies.py test/training/test_strategy.py test/training/test_checkpoint.py (passed: 252 tests)
Unit tests pass locally (make pytest)
Linting passes (make lint)
New tests added for new functionality meets coverage expectations?

Checklist

I have read and understand the Contributing Guidelines
I have updated the CHANGELOG.md
I have performed a self-review of my code
I have added docstrings to new functions/classes
I have updated the documentation (if applicable)

Additional Notes

Simple benchmarking sanity check with TRAINING_EPOCHS=100 using the provided LoRA example, compared against a similar setup trained from scratch:

Training method	Trainable parameters	Energy / atom MAE (eV/atom)	Forces MAE (eV/A)	Stress MAE (eV/A^3)
From scratch	9,063,204 / 9,063,204	0.00965	0.416	0.00347
LoRA fine-tuning	251,844 / 9,151,188	0.000363	0.0136	0.000173

This currently depends on a PhysicsNeMo version with the PEFT LoRA support used here. The version pin can be revisited and updated as the required PEFT APIs become available in a newer stable PhysicsNeMo release.

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

copy-pr-bot · 2026-07-24T21:22:19Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-07-24T21:29:51Z

Greptile Summary

This PR adds LoRA-based parameter-efficient fine-tuning (PEFT) to FineTuningStrategy, including equivariant LoRA wrappers for e3nn/cuEquivariance layers, trainable-state checkpoint saving/loading, base-model fingerprint validation, and a MACE fine-tuning example.

Core PEFT infrastructure (peft/ subpackage): LoRAConfig, LoRAHook, BaseFingerprintHook, and load_peft_checkpoint_into_model compose a clean pipeline — fingerprint before LoRA injection, inject adapters, register trainable/managed names, apply patches, then filter optimizer parameters.
Checkpoint changes (_checkpoint.py, hooks/checkpoint.py, hooks/ema.py): save_trainable_state_only saves only adapter and buffer weights; the EMA hook correctly propagates the partial-load flag through lazy initialization; _load_partial_model_state rejects unexpected keys while accepting missing frozen-parameter keys.
FineTuningStrategy extensions (finetune.py, hooks/finetune.py): Hook generation order (fingerprint → LoRA → patch → trainable filter → summary), registry-based trainable/managed name tracking, and the _replace_hooks_with_registry_validation identity check to avoid replaying on_register side effects are all well-structured.

Important Files Changed

Filename	Overview
nvalchemi/training/peft/init.py	New public PEFT subpackage; lazily exposes helpers via `__getattr__`. Contains a dead `is_lora_layer` entry in the fallback set that routes to `lora_wrappers` (which doesn't define it), shadowed by an earlier dedicated branch.
nvalchemi/training/peft/lora.py	New LoRA config, metadata serialization, and checkpoint helpers. Previously flagged dead-intersection checks have been addressed with comments and `load_partial_model_state` refactor.
nvalchemi/training/peft/loading.py	New standalone PEFT checkpoint loader; validates base fingerprint, applies LoRA and module patches, then loads full/partial state. Import path allowlisting and `trust_remote_code` guard look correct.
nvalchemi/training/peft/lora_hook.py	New `LoRAHook`; maps model-prefixed targets to model-local targets and registers trainable/managed names on the workflow.
nvalchemi/training/finetune.py	Extended `FineTuningStrategy` with `peft_config`, trainable/managed parameter registries, and `compute_base_fingerprints`. Hook ordering is correct.
nvalchemi/training/_checkpoint.py	Added `_filter_snapshot_to_trainable_state`, `_load_partial_model_state`, and `save_trainable_state_only` support. EMA hook filtering correctly uses `module.` prefix for the `AveragedModel` wrapper.
nvalchemi/training/hooks/ema.py	Added `_pending_averaged_state_load` to carry strict vs. partial load intent across lazy EMA initialization. State dict serialization and restoration paths are consistent.
nvalchemi/training/strategy.py	Updated `_replace_hooks_with_registry_validation` to skip re-registration of already-registered hooks by object identity.
pyproject.toml	Pins PhysicsNeMo to a specific pre-release git commit for PEFT APIs; acknowledged as temporary in the PR description.

_{Reviews (5): Last reviewed commit: "reduce epochs in example" | Re-trigger Greptile}

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

zubatyuk · 2026-07-28T18:43:17Z

+        # Register LoRA wrappers and apply LoRA adapters according to the saved strategy.
+        register_builtin_lora_wrappers()
+        for layer_cls, wrapper_cls in cls._wrapper_registrations_from_spec(strategy):
+            _peft.register_lora_wrapper(layer_cls, wrapper_cls)


A trust policy shoud be enforced for the custom classes in strategy.json. Coincidently, there is a related PR in another project: isayevlab/aimnetcentral#108

https://github.com/zubatyuk/aimnetcentral/blob/feat/artifact-trust-boundary/docs/model_format.md#model-yaml-import-policy

Good suggestion! In the refactored version, I have now introduced allowed_import_paths and trust_remote_code controls in nvalchemi/training/peft/loading.py, specifically for load_peft_checkpoint_into_model.

zubatyuk · 2026-07-28T18:46:18Z

+    if not isinstance(normalized, Mapping):
+        return
+    current = {
+        name: _peft.compute_base_fingerprint(model)


Note, this computes hash of the architecture, not the weights. LoRA weights should be tied to the weights of the model.

Yes, I am aware of that. I have not added weight hashes yet as I am still thinking through the best way to do it without hurting scalability for large models.

zubatyuk · 2026-07-28T18:59:13Z

+        for name in model_names:
+            # Collect this model's patch targets, metadata, and trainable states.
+            patch_prefix = f"{name}."
+            named_parameters = dict(checkpoint_models[name].named_parameters())


Not commonly used in MLIPs, but buffers like running_mean and running_var for BatchNorm should also be exported.

Thanks for pointing it out. I have now added buffers in nvalchemi.training._checkpoint._filter_snapshot_to_trainable_state.

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

Feature/lora finetuning refactor

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

ys-teh added 15 commits July 17, 2026 00:47

fix bug in DemoModelWrapper

e4f5598

allow loading partial state dict

8498174

avoid replaying hook registration

1e618a0

lora implementation

0ca34b5

generalize _matched_names

4590898

refactor lora implementation

467569d

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

add docstring to mention need for rank-zero guarding

292b86f

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

add reference energies calculations

dd3804f

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

temporarily update physicsnemo version for peft support

90ec7f8

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

add cueq wrappers

6c01a96

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

register trainable parameter names

efda44b

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

modify save_checkpoint to allow saving trainable states only

bfa0264

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

refactor to prevent introducing customized hooks

fa5cab1

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

add lora example

adc29d5

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

fix ruff formatting

293ae67

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

greptile-apps Bot reviewed Jul 24, 2026

View reviewed changes

Comment thread nvalchemi/training/_checkpoint.py Outdated

Comment thread nvalchemi/training/peft/lora.py Outdated

Comment thread nvalchemi/training/peft/lora.py Outdated

Comment thread nvalchemi/training/peft/wrappers.py Outdated

Comment thread examples/intermediate/08_lora_finetuning.py

dallasfoster requested review from dallasfoster and laserkelvin July 27, 2026 21:23

dallasfoster assigned ys-teh Jul 27, 2026

ys-teh added 3 commits July 28, 2026 02:39

partial state loading and naming

dc2d1d3

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

make is_lora_layer public

7e61113

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

handle import

0b487f2

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

zubatyuk reviewed Jul 28, 2026

View reviewed changes

ys-teh added 3 commits July 30, 2026 19:37

refactor finetuning strategy to include peft_config

c7894f5

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

refactor to merge lora specific strategy into finetuning strategy

678ac5d

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

update physicsnemo version

b5f9885

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

ys-teh and others added 4 commits July 30, 2026 19:41

fixup! refactor to merge lora specific strategy into finetuning strategy

ecd7de7

update doc

adb9739

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

Merge pull request #2 from ys-teh/feature/lora-finetuning-refactor

b75bbc7

Feature/lora finetuning refactor

reduce epochs in example

088f061

Signed-off-by: Ying Shi Teh <yteh@nvidia.com>

Navigation Menu

Uh oh!

feat(training): add LoRA fine-tuning strategy - #149

feat(training): add LoRA fine-tuning strategy#149
ys-teh wants to merge 25 commits into
NVIDIA:mainfrom
ys-teh:feature/lora-finetuning

ys-teh commented Jul 24, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 24, 2026

Uh oh!

greptile-apps Bot commented Jul 24, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zubatyuk Jul 28, 2026 •

edited

Loading

Uh oh!

zubatyuk Jul 28, 2026

Uh oh!

ys-teh Jul 30, 2026

Uh oh!

zubatyuk Jul 28, 2026

Uh oh!

ys-teh Jul 30, 2026

Uh oh!

zubatyuk Jul 28, 2026

Uh oh!

ys-teh Jul 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ys-teh commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

ALCHEMI Toolkit Pull Request

Description

Type of Change

Related Issues

Changes Made

Testing

Checklist

Additional Notes

Uh oh!

copy-pr-bot Bot commented Jul 24, 2026

Uh oh!

greptile-apps Bot commented Jul 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Important Files Changed

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zubatyuk Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zubatyuk Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

ys-teh Jul 30, 2026

Choose a reason for hiding this comment

Uh oh!

zubatyuk Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

ys-teh Jul 30, 2026

Choose a reason for hiding this comment

Uh oh!

zubatyuk Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

ys-teh Jul 30, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ys-teh commented Jul 24, 2026 •

edited

Loading

greptile-apps Bot commented Jul 24, 2026 •

edited

Loading

zubatyuk Jul 28, 2026 •

edited

Loading