Skip to content

Checkpoint interaction neighbor normalization#1538

Draft
darthjaja6 wants to merge 1 commit into
ACEsuit:developfrom
darthjaja6:fix/checkpoint-avg-neighbors-1415
Draft

Checkpoint interaction neighbor normalization#1538
darthjaja6 wants to merge 1 commit into
ACEsuit:developfrom
darthjaja6:fix/checkpoint-avg-neighbors-1415

Conversation

@darthjaja6

Copy link
Copy Markdown

Fixes #1415.

Summary

  • register each interaction block avg_num_neighbors as a persistent buffer, so it follows checkpoints and dtype/device conversions
  • keep strict loading compatible with legacy checkpoints that do not contain the new buffer
  • update model conversion and fine-tuning transfers to copy legacy float values safely

Validation

  • pytest tests/unit/test_finetuning_utils.py tests/unit/test_models.py -q (15 passed)
  • pytest tests/unit -q -k "not benchmark" (247 passed, 16 deselected)
  • MACE_REQUIRE_CAPS=gpu,cueq,oeq pytest tests/backends -q (50 passed)
  • converted the released legacy MH1 model and verified the float normalization becomes a checkpointed tensor
  • black, isort, pylint, and git diff --check

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant