Skip to content

v0.3.0

Latest

Choose a tag to compare

@laserkelvin laserkelvin released this 04 Nov 19:43
77d3c9e

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 TeacherScorer protocol and public TeacherSignal specs; built-in signals (energy, forces, stress, atomic_energies, embeddings, hessian) are requested by name. label_dataset persists the teacher fields to a resumable Zarr store.
  • Offline strategy — DistillationStrategy trains a student against a "teacher"; teacher signals arrive as teacher_* batch fields, so any built-in loss term distills by pointing its target_key at one. New AtomicEnergyMatchingLoss matches per-atom energy decompositions.
  • On-policy segment loop — the student's own propagator generates frames, TeacherLabelHook attaches teacher_* fields, and a ReplayBuffer feeds exact reference/replay mixtures. Relaxation propagators are supported via fmax and 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, and build_acceptance_report, which judges each student against an extensible bar table, including a from-scratch baseline. non_conservative_residual reports the lower bound that a non-conservative teacher puts on a conservative student's force error.

Enhanced Sampling

  • EnhancedSampling — a DynamicsStrategy that 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 of ThermodynamicStates, with exact checkpoint/restore.

Minimum-Energy Paths

  • Batched NEB — the new nvalchemi.dynamics.mep subpackage provides endpoint interpolation with optional alignment, IDPP initialization, and a GPU-first NEB strategy for regular and climbing-image nudged elastic band.

Geometry Optimizers

  • LBFGS / LBFGSVariableCell — batched L-BFGS from nvalchemi-toolkit-ops, with cell_force_scale on both variable-cell optimizers. Fixed-cell LBFGS also takes grouped updates that share curvature history across a group of graphs, one history per NEB path.

Generative Models

  • AtomisticGenerator — nvalchemi.gen runs a user-supplied generating function and returns a Batch on the contract path (hooks, device checks, composition with pipelines and dynamics). Generators compose via gen_a | gen_b and 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_dict forward **runtime_overrides; new graph_balanced_mean, per_graph_sum, and masked_mean losses.

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(...). AtomicDataZarrReader gains schema(), level_sizes(), check_integrity(), and store.
  • Precision-preserving fields — AtomicData.precision_preserving_keys keeps 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_COMPUTE hooks fire, and fused AFTER_COMPUTE hooks run after the sub-stage loop. Force priming preserves adaptive optimizer state, and reprime_on_entry refreshes forces for graphs entering a new stage.
  • PairSwapHook — pairwise swaps of per-system state with a caller-supplied acceptance rule; BaseDynamics gains apply_per_system_params() and state save/load.
  • New stages and structure sources — ON_ADMISSION initializes per-system state on admission; ON_GRADUATE fires with ctx.graduated_mask on the step a system exits; OrderedStructureSampler serves 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 new nvalchemi-distillation and nvalchemi-mep guides.
  • 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

  • MACEWrapper takes a train argument (default False). By default the weights are now frozen and the wrapper is in eval mode. Full-model training must pass train=True; otherwise TrainingStrategy and FineTuningStrategy raise "no trainable parameters" unless parameters are selected explicitly or added through PEFT. The spec run CLI and the load_checkpoint(adapter="mace") adapter pass train=True. from_checkpoint(compile_model=True) no longer freezes the weights: train and compile_model are independent.
  • MACEWrapper no longer declares an ordinary "hessian" model output. Custom MACE modules relying on the raw upstream value must use hessian_vector_product(), prepare_hessian(), or compute_hessian().

Deprecations

  • BiasedPotentialHook — superseded by nvalchemi.enhanced_sampling. Its bias_fn contract cannot express stress, so a bias applied through it is invisible to the NPT/NPH barostat. Still works, emits a DeprecationWarning.
  • FusedStage.register_fused_hook() — use the inherited register_hook().
  • FIRE and FIREVariableCell — use FIRE2 / FIRE2VariableCell or LBFGS / LBFGSVariableCell. Still works; emits a DeprecationWarning.

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.