Conversation
…e issues, add example
Greptile SummaryLarge architectural refactor unifying Important Files Changed
Reviews (9): Last reviewed commit: "fix skin dropping" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| outputs = {"energy", "forces", "stress", "charges"} | ||
| if self._is_nse: | ||
| outputs.add("spin_charges") | ||
|
|
||
| self.model_config = ModelConfig( | ||
| outputs=frozenset(outputs), | ||
| autograd_outputs=frozenset({"forces", "stress"}), | ||
| autograd_inputs=frozenset({"positions"}), | ||
| required_inputs=frozenset({"charge"}), | ||
| optional_inputs=frozenset(), | ||
| supports_pbc=True, | ||
| needs_pbc=False, | ||
| neighbor_config=None, # AIMNet2 manages its own neighbor list | ||
| active_outputs=outputs, | ||
| ) |
There was a problem hiding this comment.
"stress" advertised in outputs but standalone stress computation silently does nothing
"stress" is included in model_config.outputs (and in autograd_outputs) and in active_outputs, meaning output_data() returns it and adapt_output will include a "stress" key in every return value. However, the forward() implementation at lines 502–506 does:
if compute_stresses:
# Stresses require a displacement tensor — not yet implemented
# for standalone AIMNet2.
passThe key is never populated, so callers receive output["stress"] = None with no warning. Any downstream code that checks if "stress" in outputs or iterates over the dict will silently consume a None tensor, which can cause cryptic failures later (e.g. in sum_outputs).
Either remove "stress" from outputs/autograd_outputs until standalone stress is implemented, or emit a warnings.warn when the key is requested but cannot be satisfied:
# Remove stress from advertised capabilities until implemented
outputs = {"energy", "forces", "charges"}| If the number of models doesn't match the saved config, or | ||
| if a group requires a derivative_fn that wasn't provided. | ||
| """ | ||
| checkpoint = torch.load(path, weights_only=False) |
There was a problem hiding this comment.
weights_only=False allows arbitrary pickle execution
torch.load(..., weights_only=False) deserializes arbitrary Python objects, making it possible for a malicious checkpoint to execute arbitrary code. Since the saved format only needs primitive Python types (dict, list, str, bool, int), weights_only=True can be used for the config/metadata portion:
checkpoint = torch.load(path, weights_only=True)If that is not sufficient for all saved dtypes, consider splitting into two loads — one weights_only=True for config metadata and one restricted call for weights — or using torch.serialization.add_safe_globals to allowlist only the needed types.
There was a problem hiding this comment.
Summary
This is a large, well-structured refactor that unifies the model subsystem around ModelConfig (replacing the ModelCard/ModelConfig split), replaces ComposableModelWrapper with a more powerful PipelineModelWrapper supporting autograd groups and output wiring, and adds first-class AIMNet2 support. The architecture is a clear improvement: frozen capability fields + mutable runtime fields in ModelConfig, the pipeline group/step abstraction for composing models, and shared autograd utilities. The test coverage for base.py and pipeline.py is excellent, including analytical force correctness checks. However, there are broken AIMNet2 tests and stale documentation that need fixing before merge.
Issues Found
-
[critical]
test/models/test_aimnet2.py:L121-129,L145-151-- All 14 AIMNet2 tests fail because_make_wrapper()andtest_import_guardreferencem._AIMNET_AVAILABLEwhich no longer exists. The module now uses@OptionalDependency.AIMNET.requireinstead of a module-level boolean flag.
Fix: Rewrite_make_wrapper()to monkey-patch theOptionalDependency.AIMNETavailability check instead of the removed_AIMNET_AVAILABLEflag. Fortest_import_guard, patch the decorator to simulate the import guard behavior. -
[critical]
docs/modules/models.rst:L18,L88-95-- ReferencesModelCard(line 18) andnvalchemi.models.composable.ComposableModelWrapper(lines 88-95), both of which are deleted in this PR. Sphinx build will produce broken cross-references.
Fix: ReplaceModelCardwithModelConfigon line 18. Replace the "Composition" section with entries forPipelineModelWrapper,PipelineStep,PipelineGroupundernvalchemi.models.pipeline. AddAIMNet2Wrapperunder a new ML Potentials subsection. -
[warning]
docs/userguide/about/intro.md:L114,L150-- Refers toModelCard(line 114) andComposableModelWrapper(line 150) in the intro guide table.
Fix: Update to referenceModelConfigandPipelineModelWrapper. -
[warning]
docs/index.md:L49-- "standardizedModelCardinterface" should be "standardizedModelConfiginterface". -
[warning]
nvalchemi/dynamics/base.py:L1394-1395-- The instance variable name remainsself.model_cardbut now points to aModelConfigobject. This naming mismatch is confusing.
Fix: Consider renaming toself.model_configwith a deprecation property formodel_card. -
[warning]
nvalchemi/models/pipeline.py_check_wiring()-- Accumulates available keys and logs them but performs no actual validation that downstream models' required inputs are satisfied by upstream outputs. The method name implies checking but it's purely informational.
Fix: Either rename to_log_wiring()or add actual validation.
Downstream Impact
- Dynamics:
compute()rewritten with output detachment and flexible key mapping. FusedStagemasked_updatefix removespost_updatecall. All 936 dynamics tests pass. - Neighbor hooks: Refactored to delegate to shared
_write_neighbor_data_to_batch(). Functionally equivalent. __init__.py: Lazy imports updated. No remaining references toModelCard,ComposableModelWrapper, orregistryin source code.- Examples:
07_composable_model_composition.pyupdated for+operator. New08_aimnet2_ewald_pipeline.pydemonstrates autograd pipeline. Both look correct.
Test and Example Results
- test_base.py: All PASS
- test_pipeline.py: All PASS (including analytical force verification)
- test_aimnet2.py: All 14 FAIL --
AttributeError: module has no attribute '_AIMNET_AVAILABLE' - test_composable.py: All 27 migration tests PASS
- test_ewald.py: All PASS
- test_pme.py: All PASS
- test/dynamics/: All 936 non-slow tests PASS
User Documentation
Must update before merge:
docs/modules/models.rst-- RemoveModelCardandComposableModelWrapper; add pipeline and AIMNet2 classesdocs/userguide/about/intro.mdL114, L150 -- ReplaceModelCardandComposableModelWrapperdocs/index.mdL49 -- ReplaceModelCardwithModelConfig- Consider adding a doc page for
nvalchemi.neighbors
Verdict
Request Changes -- The architecture and implementation quality are high, but all 14 AIMNet2 tests are broken and multiple doc pages reference deleted types. Fix the test helper, update the stale docs, and this is ready to merge.
laserkelvin
left a comment
laserkelvin
left a comment
There was a problem hiding this comment.
Left inline comments as well as the ones in my earlier comment
| .. model-capability-table:: | ||
| :category: ml | ||
| ```{list-table} | ||
| :header-rows: 1 |
There was a problem hiding this comment.
We can reinstate this later, but it was kind of nice to be able to auto-generate this table
| Combined output dict. | ||
| """ | ||
| additive = additive_keys or {"energy", "forces", "stress"} | ||
| result: ModelOutputs = OrderedDict() |
There was a problem hiding this comment.
Out of scope for this PR, but I imagine using TensorDict instead might be better than just the OrderedDict
|
Tip: Greploop — Automatically fix all review issues by running Use the Greptile plugin for Claude Code to query reviews, search comments, and manage custom context directly from your terminal. |
|
/ok to test 4f069de |
| # Enable grad on positions if any autograd output is active. | ||
| if self.model_config.autograd_outputs & self.model_config.active_outputs: | ||
| data.positions.requires_grad_(True) |
There was a problem hiding this comment.
requires_grad_(True) on non-leaf scaled_pos raises RuntimeError
When standalone stress computation is requested, forward() calls prepare_strain and replaces data["positions"] with scaled_pos before calling adapt_input. scaled_pos is a non-leaf tensor (derived from displacement.requires_grad_(True)), so calling requires_grad_(True) on it from inside adapt_input raises RuntimeError: you can only change requires_grad flags of leaf variables.
Reproduction path: standalone AIMNet2Wrapper with "stress" in active_outputs and a batch with a cell.
The fix is to enable gradients on positions before prepare_strain is applied — matching the pattern used by PipelineModelWrapper._run_autograd_group: enable requires_grad on positions first, then apply prepare_strain so strain is built on top of the leaf tensor, then remove the requires_grad_(True) call from adapt_input.
| def _configure_sub_models(self) -> None: | ||
| """Adjust sub-model configs based on group autograd strategy.""" | ||
| for group in self.groups: | ||
| if group.use_autograd: | ||
| for step in group.steps: | ||
| # Remove derivative keys from sub-model active_outputs — | ||
| # the group will compute these via autograd. | ||
| new_active = set(step.model.model_config.active_outputs) | ||
| new_active -= {"forces", "stress"} | ||
| step.model.model_config = step.model.model_config.model_copy( | ||
| update={"active_outputs": new_active} | ||
| ) |
There was a problem hiding this comment.
_configure_sub_models permanently mutates sub-model model_config
At pipeline construction time, _configure_sub_models replaces each sub-model's model_config with a copy that has "forces" and "stress" stripped from active_outputs. Because step.model.model_config is replaced directly, the sub-model is permanently changed for the lifetime of the object.
If the same model instance is later reused — in a second pipeline, standalone, or passed to a dynamics engine — it will be missing forces/stress in its active_outputs, silently producing only energies. The second pipeline's _build_model_config will then synthesize an active_outputs that lacks derivatives, with no warning.
The fix is to record the desired active outputs at pipeline build time without mutating the sub-model, for example by storing the per-step override in a dict keyed by step identity instead of patching step.model.model_config.
|
/ok to test a8987d4 |
|
/ok to test ff98e22 |
| neighbor_config = NeighborConfig( | ||
| cutoff=max_cutoff, | ||
| format=chosen_format, | ||
| half_list=sub_neighbor_configs[0].half_list, | ||
| max_neighbors=max_neighbors, | ||
| ) |
There was a problem hiding this comment.
skin dropped when synthesizing pipeline NeighborConfig
_build_model_config constructs the merged NeighborConfig without forwarding skin, so it always defaults to 0.0. make_neighbor_hooks() then calls NeighborListHook(nc, skin=nc.skin) — but nc.skin is 0.0 regardless of what any sub-model specified. Any pipeline whose sub-models set skin > 0 will rebuild the neighbor list every step, silently discarding the Verlet-skin optimization.
|
/ok to test c34cde0 |
|
/ok to test a290412 |
2 participants
ALCHEMI Toolkit Pull Request
Description
Refactor the model subsystem to unify
ModelCardandModelConfiginto a singleModelConfigwith capability (frozen) andruntime (mutable) fields, replace
ComposableModelWrapperwith a pipeline-based composition system (PipelineModelWrapper),remove the model registry, and add first-class AIMNet2 support.
Type of Change
Related Issues
Changes Made
ModelConfig: MergedModelCard(capability declaration) andModelConfig(runtime control) into a singleModelConfigclass. Capability fields usefrozensetto signal immutability; runtime fields (active_outputs,gradient_keys)remain mutable. All model wrappers updated to the new interface.
ComposableModelWrapperwithPipelineModelWrapper,PipelineGroup, andPipelineStep.Groups define mini-pipelines with independent derivative strategies; the
+operator provides simple additive composition.Supports dependent model chains (e.g., AIMNet2 charges wired into Ewald via
wiremappings).nvalchemi/models/registry.pyand associated tests.AIMNet2Wrapper: Addednvalchemi/models/aimnet2.pywith full AIMNet2 model support.nvalchemi.neighborsmodule: Standalonecompute_neighbors()function for one-shot neighbor list construction outsidedynamics loops.
nvalchemi/models/_utils.py: Extracted shared utilities (autograd derivative computation, stress/force helpers) usedacross model wrappers.
DemoModelWrapper,DFTD3ModelWrapper,EwaldModelWrapper,LennardJonesModelWrapper,MACEWrapper,PMEModelWrapper) to use the newModelConfiginterface andBaseModelMixincontract.NeighborListHookto work with the refactoredModelConfig.neighbor_config.test_physical_models.pyinto per-model test files (test_ewald.py,test_pme.py,test_dftd3.py,test_aimnet2.py,test_pipeline.py). Addedtest_base.pyforModelConfigandBaseModelMixinunit tests.Testing
make pytest)make lint)Checklist
Additional Notes
Breaking changes for downstream users:
ModelCardno longer exists — useModelConfigwithoutputs,autograd_outputs, etc. instead ofsupports_*/forces_via_autogradbooleans.ComposableModelWrapperremoved — migrate toPipelineModelWrapperor use themodel1 + model2operator.model_config.compute_forces,compute_stresses, etc. replaced bymodel_config.active_outputs = {"energy", "forces", ...}.register_model,list_foundation_models, etc.) removed.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.