Add MatterGen crystal diffusion score-core export - #731
Open
justinchuby wants to merge 2 commits into
Open
Conversation
Validate the official mp_20_base score core against pinned MatterGen source and a source-derived golden fixture. Cover PBC invariance and the dft_band_gap adapter's conditional and unconditional score paths. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 332c98c9-cf56-4bab-8552-ac1b6d730556 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Export the pinned official MatterGen GemNet-T score core with Hydra configuration parsing, exhaustive checkpoint routing, a documented dynamic PBC graph ABI, and CLI integration. Add a source-faithful CPU host sampler with the released D3PM/VE/VP schedule and a real ONNX sampling golden while explicitly retaining host orchestration in the partial export report. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 332c98c9-cf56-4bab-8552-ac1b6d730556 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| @@ -0,0 +1,183 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,187 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,187 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| # Copyright (c) Microsoft Corporation. | ||
| # Licensed under the MIT License. | ||
|
|
||
| from __future__ import annotations |
| @@ -0,0 +1,323 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,1206 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| self.div_term = nn.Parameter([hidden_dim // 2]) | ||
| # Source registers div_term as a float32 buffer. It must not be | ||
| # demoted with model weights when the exporter requests fp16/bf16. | ||
| setattr(self.div_term, "_keep_float32", True) |
| @@ -0,0 +1,208 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
| @@ -0,0 +1,208 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
|
|
||
| """Task wiring for the host-orchestrated MatterGen GemNet-T score core.""" | ||
|
|
||
| from __future__ import annotations |
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
|
|
||
| def run( | ||
| self, output_names: Sequence[str] | None, input_feed: Mapping[str, np.ndarray] | ||
| ) -> Sequence[np.ndarray]: ... |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
There are a couple of concrete robustness issues (artifact validation in the parity fixture and CLI error handling for save-time failures) that should be fixed before merging.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR adds a first-class Mobius integration for the pinned microsoft/mattergen release by exporting the deterministic float32 GemNet‑T score core to standard ONNX, while explicitly keeping periodic graph construction and the diffusion/sampling loop as host-owned responsibilities (with a provided CPU host runtime).
Changes:
- Add MatterGen score-core model +
MatterGenScoreTaskexporting an explicit periodic-graph + condition-port ABI, plus strict Lightning checkpoint routing. - Add a source-faithful CPU host runtime (
MatterGenHostSampler) including periodic graph/triplet construction and the full 1,000-step schedule around a pluggable score callback. - Add CLI entrypoint support (
--mattergen-checkpoint), documentation, and L3/L4/L5 parity + golden tests.
File summaries
| File | Description |
|---|---|
| tests/integration/mattergen_parity_test.py | Adds pinned-source L3/L4 parity and an L5 host-sampling golden against real weights. |
| tests/cli_test.py | Adds CLI routing/validation tests for MatterGen-specific build behavior. |
| testdata/golden/diffusion/mattergen-mp20-score.json | Adds one-step real-weight score-core golden provenance + outputs. |
| testdata/golden/diffusion/mattergen-mp20-host-sample.json | Adds deterministic 1,000-step host-sampling golden artifact provenance. |
| src/mobius/tasks/_mattergen.py | Introduces MatterGenScoreTask wiring the explicit host ABI and partial export report. |
| src/mobius/tasks/init.py | Registers MatterGenScoreTask and task name mattergen-score. |
| src/mobius/models/mattergen.py | Declares the GemNet‑T score-core graph and strict preprocess/weight-routing behavior. |
| src/mobius/models/init.py | Exposes MatterGen models from the public models package. |
| src/mobius/integrations/mattergen/_weights.py | Adds fail-closed Lightning checkpoint loading + exact initializer routing. |
| src/mobius/integrations/mattergen/_weights_test.py | Unit tests for strict weight loading, routing, and reporting. |
| src/mobius/integrations/mattergen/_runtime.py | Adds the host-owned periodic graph + sampler runtime and ORT callback adapter. |
| src/mobius/integrations/mattergen/_runtime_test.py | Unit tests for periodic graph ordering, CFG call order, and sampler determinism. |
| src/mobius/integrations/mattergen/_contract.py | Defines pinned hub/source commits, host-owned steps, allowlists, and validation helpers. |
| src/mobius/integrations/mattergen/_contract_test.py | Validates the pinned contract constants and helper semantics. |
| src/mobius/integrations/mattergen/_configs.py | Parses expanded Hydra configs into a typed MatterGenConfig without Hydra deps. |
| src/mobius/integrations/mattergen/_builder.py | Adds pinned builder (build_mattergen) supporting Hub or safe local checkpoint roots. |
| src/mobius/integrations/mattergen/_builder_test.py | Tests config parsing, port ABI, pinned revision enforcement, and family coverage. |
| src/mobius/integrations/mattergen/init.py | Adds a lazy public API surface for MatterGen integration symbols. |
| src/mobius/main.py | Adds MatterGen CLI routing and a new --mattergen-checkpoint flag. |
| src/mobius/init.py | Exposes build_mattergen at the package top level. |
| pyproject.toml | Adds PyYAML dependency used for loading pinned Hydra configs. |
| docs/model-catalog.md | Documents MatterGen under periodic crystal diffusion exports. |
| docs/mattergen.md | Adds detailed ABI, host/runtime boundary, and build instructions for MatterGen. |
| docs/index.md | Adds MatterGen docs page to the documentation index. |
Review details
- Files reviewed: 24/24 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+384
to
+396
| try: | ||
| pkg = build_mattergen( | ||
| mattergen_source, | ||
| checkpoint=args.mattergen_checkpoint or "mp_20_base", | ||
| revision=revision, | ||
| dtype=dtype_override, | ||
| load_weights=load_weights, | ||
| execution_provider=execution_provider, | ||
| ) | ||
| except ValueError as error: | ||
| raise SystemExit(f"Error: {error}") from error | ||
| _save_package(pkg, output_dir, args, optimize, component_filter) | ||
| return |
Comment on lines
+156
to
+162
| path = Path(raw_path) | ||
| if not path.is_file() and environment_variable == _SOURCE_DIR_ENV: | ||
| if not path.is_dir(): | ||
| pytest.skip(f"{environment_variable} is not a readable source directory: {path}") | ||
| elif not path.is_file(): | ||
| pytest.skip(f"{environment_variable} is not a readable checkpoint: {path}") | ||
| return path |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
microsoft/mattergen@5244495dd9a979ff71abc7548a0b14b9deb0069asupport for all released checkpoint families.Runtime boundary
ModelPackageremains deliberately partial (end_to_end_runnable: false): periodic neighbor/triplet construction, D3PM/VE/VP scheduling, classifier-free guidance, and crystal validation are host-owned.MatterGenHostSamplersupplies the source-faithful CPU host loop around an explicit ONNX Runtime callback; ONNX Runtime GenAI is not presented as an end-to-end runtime.Evidence
842ffe735f7d06cec89d56aa23d9f001e1124b30.dft_band_gap.mp_20_baseONNX weights and validates the committed deterministic crystal artifact.Limitations
Only CPU/default float32 score-core execution is assessed. f16, bf16, CUDA, and GenAI are explicitly refused. Pymatgen Structure/CIF serialization remains application-owned; the host validates a dependency-free structural artifact.