Skip to content

Adding training functionalities to Toolkit#108

Open
laserkelvin wants to merge 378 commits into
NVIDIA:mainfrom
laserkelvin:training-epic
Open

Adding training functionalities to Toolkit#108
laserkelvin wants to merge 378 commits into
NVIDIA:mainfrom
laserkelvin:training-epic

Conversation

@laserkelvin

@laserkelvin laserkelvin commented Jun 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

ALCHEMI Toolkit Pull Request

Description

This PR introduces the core functionalities required to support training and fine-tuning of models in nvalchemi-toolkit.

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

Changes Made

  • create_model_spec methods and dynamic pydantic model creation for pickle-less serialization of configuration
  • Adds a few base loss functions, the general loss abstraction including individual losses and a composed loss function. The latter can be adjusted with weight scheduling, allowing the relative weighting of different losses to be adjusted over the course of training
  • Adds a TrainingStrategy pydantic model as a recipe validation and loop executor. The execution is highly modular and extendible, allowing for (hopefully) arbitrarily complex training workflows to be built, and not limited to MLIPs
  • Adds a FineTuningStrategy that specializes TrainingStrategy for...fine-tuning workflows by making pre-existing checkpoints and layer addition/modification integral to the workflow
  • Adds data loading optimizations; the main changes is addition of "batched" pre-fetching, which amortizes I/O for non-contiguous data samples. This is crucial for Zarr performance when shuffling data
  • Adds multidataset support, with a "meta" sampler that allows users to implement different cross-dataset sampling strategies (e.g. to account for dataset size imbalances)
  • Adds several training-related hooks, such as model averaging, mixed precision, checkpointing
  • Adds a CLI for training and fine-tuning: the intended use of this CLI is to provide a relatively straightforward on-ramp for users looking to get fine-tune (or train a model from scratch) quickly without needing to know the full training API

Testing

  • 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

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.

laserkelvin added 30 commits May 26, 2026 08:16
@laserkelvin
docs(training): document mixed precision hooks
627b16a 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
fix(training): narrow AMP autocast scope
ee9c3e0 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
refactor(training): dispatch mixed precision hook stages
75ea950 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
test(training): align mixed precision tests with train batch helper
c63e3c1 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
fix(training): prevent duplicate mixed precision hooks
8b75aad 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
test(training): align update hook API expectations
2396c0c 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
Merge pull request #9 from laserkelvin/feat-training-update-orchestrator
14580ac 
Add TrainingUpdateHook framework and orchestrator
@laserkelvin
Merge remote-tracking branch 'fork/training-epic' into feat-mixed-pre…
b9aa80b 
…cision-hook
@laserkelvin
test: consolidating and using existing device fixture
ad7ba4c 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
Merge pull request #7 from laserkelvin/feat-mixed-precision-hook
92eaa19 
Add `MixedPrecisionHook`
@laserkelvin
Merge branch 'feat-training-update-orchestrator' into feat-ema-hook
c85c44f 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>

# Conflicts:
#	nvalchemi/training/hooks/update.py
#	test/training/test_strategy.py
@laserkelvin
Merge remote-tracking branch 'fork/training-epic' into feat-ema-hook
03b307b 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>

# Conflicts:
#	docs/modules/training/hooks.rst
#	test/training/test_strategy.py
@laserkelvin
Merge remote-tracking branch 'fork/training-epic' into feat-ema-hook
cdbe62e 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>

# Conflicts:
#	nvalchemi/training/hooks/update.py
#	test/training/test_strategy.py
@laserkelvin
feat(training): add strategy checkpoint restart loading
6b810c1 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
fix(training): restore checkpoint restart consistency
e441123 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
docs(training): note checkpoint restart workflow
690b5d3 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
Merge pull request #8 from laserkelvin/feat-ema-hook
8acb8b2 
Add EMAHook for exponential moving average of model weights
@laserkelvin
fix(data): generate edge rows in io benchmark
9720c5b
@laserkelvin
refactor(data): profile io benchmark readback
af3095d
@laserkelvin
refactor(data): batch zarr dataloader reads
a424720
@laserkelvin
feat(data): compare zarr readback modes
01bc4f3
@laserkelvin
docs(data): document zarr readback modes
e1a23e8
@laserkelvin
docs(data): refresh zarr benchmark examples
849adf0
@laserkelvin
Merge remote-tracking branch 'origin/main' into training-epic
eb51e24 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
Merge remote-tracking branch 'fork/training-epic' into feat-checkpoin…
e774991 
…t-loading

Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
feat(training): add periodic checkpoint hook
35f76ee 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
fix(training): respect checkpoint hook lifecycle
def6893 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
fix(training): make checkpoint hook cadence explicit
2b64eac 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
refactor: simplifying mutual exclusion
03b5b8e 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
test(training): cover checkpoint hook restart cycles
5263c85 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
feat(training): add validation batch preparation callback
22ecded 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
Revert "feat(training): add validation batch preparation callback"
e542169 
This reverts commit 22ecded.

Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
docs: clarifying validation
8c01fea 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
feat(training): add CLI runtime hook stage overrides
f5dc0ef 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
docs(training): align fine-tuning CLI guidance
0c27326 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
docs: noting equivariant models need specialized readout
8073ecf 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin

laserkelvin commented Jun 23, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 8073ecf

@laserkelvin

laserkelvin commented Jun 23, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test bda67ad

laserkelvin added 10 commits June 23, 2026 11:41
@laserkelvin
fix(training): wire CLI validation data
04b0ebe 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
fix(training): require checkpoint cadence
09ae1a0 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
perf(data): streamline multidataset samplers
5caa269 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
refactor(distributed): share rank and collective helpers
1aabcac 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
feat(training): add loss dtype alignment policy
485b81f 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
docs(training): document loss dtype alignment
691d06a 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
feat(training): expose loss dtype policy in CLI
05ae3e0 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
docs(training): document CLI dtype alignment
44dc143 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
fix(training): restore filtered optimizer checkpoints
6c18cd6 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
feat(training): add MACE source E0 options
4d4e024 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin

laserkelvin commented Jun 23, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 8aa39b4

@laserkelvin
fix(training): preserve optimizer state on checkpoint resume
2b35272 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin
feat(training): add CLI resume and validation options
4d49093 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin

laserkelvin commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 4d49093

@laserkelvin
docs: adding agent skill for reporting abstraction
9d99fc3 
Signed-off-by: Kelvin Lee <kinlongkelvi@nvidia.com>
@laserkelvin

laserkelvin commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test 9d99fc3

@laserkelvin

laserkelvin commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test b66aeaf

dallasfoster
dallasfoster reviewed Jun 24, 2026
View reviewed changes
Comment thread docs/modules/dynamics/hooks.rst
Comment on lines +260 to +269
Use :class:`~nvalchemi.dynamics.hooks.StageTimingHook` for lightweight stage
timing and optional NVTX ranges.

.. code-block:: python

from nvalchemi.dynamics.hooks import ProfilerHook
from nvalchemi.dynamics.hooks import StageTimingHook

hook = ProfilerHook(enable_nvtx=True, enable_timer=True, frequency=10)
hook = StageTimingHook("step", frequency=10, log_path="stage_timing.csv")
dynamics = DemoDynamics(model=model, n_steps=1_000, dt=0.5, hooks=[hook])
dynamics.run(batch)

@dallasfoster dallasfoster Jun 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

This example doesn't really tell me what the heck this hook is doing, what "step" refers to, what "frequency" means. We don't need full api doc here but I would expect just another sentence with sufficient exposition explaining what is going on here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ys-teh ys-teh ys-teh left review comments
@dallasfoster dallasfoster dallasfoster requested changes
@zubatyuk zubatyuk Awaiting requested review from zubatyuk
@atulcthakur atulcthakur Awaiting requested review from atulcthakur

Requested changes must be addressed to merge this pull request.

Assignees

No one assigned

Labels

enhancement New feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@laserkelvin @ys-teh @dallasfoster @EricZQu