Repository navigation
Third public-beta release of NVIDIA ALCHEMI Toolkit. The headline is distillation: teacher scoring, offline and on-policy training loops that run data-parallel, and an acceptance evaluation suite. This release also adds enhanced sampling, batched minimum-energy paths, a generative model API, L-BFGS optimizers, and LoRA fine-tuning.
Distillation
Train a compact student potential against a larger teacher, from a fixed dataset or from the student's own trajectories.
- Teacher scoring and offline labeling — a
TeacherScorerprotocol and publicTeacherSignalspecs; built-in signals (energy,forces,stress,atomic_energies,embeddings,hessian) are requested by name.label_datasetpersists the teacher fields to a resumable Zarr store. - Offline strategy —
DistillationStrategytrains a student against a"teacher"; teacher signals arrive asteacher_*batch fields, so any built-in loss term distills by pointing itstarget_keyat one. NewAtomicEnergyMatchingLossmatches per-atom energy decompositions. - On-policy segment loop — the student's own propagator generates frames,
TeacherLabelHookattachesteacher_*fields, and aReplayBufferfeeds exact reference/replay mixtures. Relaxation propagators are supported viafmaxand a convergence hook. Opt-in settings cap trajectory length (max_trajectory_steps, which recycles the seed pool) and freeze diverged frames (quarantine_diverged). - Multi-GPU on-policy — data-parallel under a
DDPHook: each rank labels its own shard with its own teacher replica. - Embedding, Hessian, and Boltzmann objectives — three loss terms distill what no reference dataset carries: per-atom representations, Hessian-vector products, and Boltzmann relative entropy.
- Evaluation and acceptance suite —
evaluate_accuracy(MAE/RMSE, force cosine similarity), stability and extensivity measurements, andbuild_acceptance_report, which judges each student against an extensible bar table, including a from-scratch baseline.non_conservative_residualreports the lower bound that a non-conservative teacher puts on a conservative student's force error.
Enhanced Sampling
EnhancedSampling— aDynamicsStrategythat adds bias, replica-exchange and epoch-commit hooks to the engine it builds. A bias is an ordinary additive potential, summed once against the unmodified model output. Built-ins:HarmonicUmbrellaBias,UpperWall/LowerWall,FlatBottomRestraint, well-tempered and RMSD metadynamics,AdaptiveBiasingForce.ReplicaExchange— synchronous temperature or umbrella exchange over a ladder ofThermodynamicStates, with exact checkpoint/restore.
Minimum-Energy Paths
- Batched NEB — the new
nvalchemi.dynamics.mepsubpackage provides endpoint interpolation with optional alignment, IDPP initialization, and a GPU-firstNEBstrategy for regular and climbing-image nudged elastic band.
Geometry Optimizers
LBFGS/LBFGSVariableCell— batched L-BFGS fromnvalchemi-toolkit-ops, withcell_force_scaleon both variable-cell optimizers. Fixed-cellLBFGSalso takes grouped updates that share curvature history across a group of graphs, one history per NEB path.
Generative Models
AtomisticGenerator—nvalchemi.genruns a user-supplied generating function and returns aBatchon the contract path (hooks, device checks, composition with pipelines and dynamics). Generators compose viagen_a | gen_band round-trip through pydantic.
Crystal Structure Prediction
nvalchemi.csp— molecular crystal-structure prediction: random rigid-molecule packing, atom-bounded radial matching, and greedy de-duplication.
Training
- LoRA fine-tuning — PEFT support within
FineTuningStrategy:LoRAConfig,load_peft_checkpoint_into_model, and base-model fingerprint checks. - Checkpoint loading and losses —
load_checkpoint/from_spec_dictforward**runtime_overrides; newgraph_balanced_mean,per_graph_sum, andmasked_meanlosses.
Data
GroupLayout— treat contiguous graphs as logical groups, with graph/node mappings, reductions, broadcasts, and grouped-batch append support.- Batch level API —
drop_level,pop_level,set_level,level_keys;index_select(idx, drop=),clone(drop=),without_keys(...).AtomicDataZarrReadergainsschema(),level_sizes(),check_integrity(), andstore. - Precision-preserving fields —
AtomicData.precision_preserving_keyskeeps chosen floating-point fields at their own dtype instead of the positions dtype, so a high-precision label is not quantized by float32. - Transactional checkpoints —
save_checkpoint/load_checkpoint: checksummed, pickle-free Zarr checkpoints for named stateful objects.
Dynamics
- Fused-stage lifecycle fixes — fused-level hooks now fire at all four pre/post-update boundaries, sub-stage
BEFORE_COMPUTEhooks fire, and fusedAFTER_COMPUTEhooks run after the sub-stage loop. Force priming preserves adaptive optimizer state, andreprime_on_entryrefreshes forces for graphs entering a new stage. PairSwapHook— pairwise swaps of per-system state with a caller-supplied acceptance rule;BaseDynamicsgainsapply_per_system_params()and state save/load.- New stages and structure sources —
ON_ADMISSIONinitializes per-system state on admission;ON_GRADUATEfires withctx.graduated_maskon the step a system exits;OrderedStructureSamplerserves dataset rows in order with sharding and restartable state.
Models
- Position Hessian APIs — local Hessian-vector products, reusable Hessian operators, and in-place dense Hessians for AIMNet2, MACE, Ewald/PME, and flat pipelines.
- UMA periodic OMol stress — periodic stress outputs for
task_name="omol".
Developer & Agent Experience
- Agent skills — 14 Claude Code skills (
.claude/skills/), with newnvalchemi-distillationandnvalchemi-mepguides. - User guides and examples — new guides for distillation, enhanced sampling, generative models, minimum-energy paths and crystal structure prediction, with matching gallery examples.
Breaking Changes
MACEWrappertakes atrainargument (defaultFalse). By default the weights are now frozen and the wrapper is in eval mode. Full-model training must passtrain=True; otherwiseTrainingStrategyandFineTuningStrategyraise "no trainable parameters" unless parameters are selected explicitly or added through PEFT. Thespec runCLI and theload_checkpoint(adapter="mace")adapter passtrain=True.from_checkpoint(compile_model=True)no longer freezes the weights:trainandcompile_modelare independent.MACEWrapperno longer declares an ordinary"hessian"model output. Custom MACE modules relying on the raw upstream value must usehessian_vector_product(),prepare_hessian(), orcompute_hessian().
Deprecations
BiasedPotentialHook— superseded bynvalchemi.enhanced_sampling. Itsbias_fncontract cannot express stress, so a bias applied through it is invisible to the NPT/NPH barostat. Still works, emits aDeprecationWarning.FusedStage.register_fused_hook()— use the inheritedregister_hook().FIREandFIREVariableCell— useFIRE2/FIRE2VariableCellorLBFGS/LBFGSVariableCell. Still works; emits aDeprecationWarning.
Requirements
- Python 3.11–3.14
- PyTorch >= 2.8
nvalchemi-toolkit-ops>= 0.5.0- numpy >= 2, < 3
- Optional:
[mace],[aimnet],[uma],[uma-cu12],[uma-cu13],[ase],[pymatgen],[rdkit],[tensorboard],[cu12],[cu13]extras
Multi-GPU on-policy distillation runs under a DDPHook; every other new feature runs on a single GPU.