Skip to content

Support JAX multitask training and finetuning - #234

Open
Romarin87 wants to merge 27 commits into
njzjz:parallelfrom
Romarin87:parallel
Open

Support JAX multitask training and finetuning#234
Romarin87 wants to merge 27 commits into
njzjz:parallelfrom
Romarin87:parallel

Conversation

@Romarin87

@Romarin87 Romarin87 commented Jun 26, 2026 •
edited
Loading

Copy link
Copy Markdown

Summary

This PR extends the JAX backend on the parallel branch with multitask training, finetuning, init/frozen-model loading fixes, mixed-system stat handling, and related serialization/freeze support.

The main goal is to bring the JAX training path closer to the PyTorch multitask/finetune behavior while preserving existing single-task training.

Base: njzjz/deepmd-kit:parallel
Head: Romarin87/deepmd-kit:parallel

Main Changes

JAX multitask training

  • Add JAX ModelWrapper support for model_dict multitask models.
  • Add shared parameter preprocessing for shared_dict.
  • Support shared type_map, descriptor, fitting net, and hybrid descriptor subcomponents.
  • Support case embedding assignment across task heads.
  • Add multitask data/stat handling in the JAX trainer.
  • Support task scheduling/probability handling during training.

JAX finetuning

  • Add single-task finetune support for JAX .jax / .hlo models.
  • Add multitask finetune support with branch selection.
  • Support model_branch / finetune_head behavior.
  • Support random fitting head initialization when requested.
  • Preserve target-side trainable flags while optionally loading compatible pretrained model parameters.
  • Add frozen-model merge logic for init-frz-model.
  • Fix finetune output-bias adjustment so sampled natoms are expanded per frame.
  • Use real_natoms_vec during finetune bias statistics when mixed/padded data provides true atom counts.

Serialization and freeze

  • Extend JAX serialization/deserialization for multitask model data.
  • Add branch selection helper for multitask checkpoints.
  • Allow freeze to export a selected multitask branch via head / model_branch.
  • Preserve Hessian export path when training from restart/init-model checkpoints.

JAX compatibility fixes

  • Add compatibility handling for older JAX mesh/sharding APIs.
  • Avoid eager-sharding failures on old JAX versions.
  • Avoid PartitionSpec device_put issues on older JAX.
  • Add shim for missing monitoring record_scalar.
  • Restore RepFlow execution and serialization behavior under JAX.
  • Add RepFlow rematerialization/checkpointing support to reduce EFH memory usage.
  • Keep array creation trace-safe for JAX export paths.

Mixed-system statistics and padding

  • Fix mixed-system stat sampling when using padded JAX batches with different atom counts.
  • Collect model statistics from the original per-system data rather than from padded mixed batches.
  • Preserve real_natoms_vec for mixed-type fitting statistics.
  • Fix mixed-batch Hessian labels to pad by (3 * Nmax)^2, not ordinary Nmax * ndof atomic-label width.

CI/dependency fixes

  • Bump cibuildwheel to v3.4.
  • Align C package TensorFlow wheel build with TensorFlow 2.20.
  • Pin protobuf to avoid current TensorFlow/protobuf incompatibility.
  • Remove the deprecated flake8 pre-commit hook and follow current ruff/pylint rules.

Tests Added

  • source/tests/jax/test_finetune.py
  • source/tests/jax/test_multitask.py
  • source/tests/common/test_model_stat.py

These cover:

  • JAX finetune rule generation.
  • Single-task and multitask finetune behavior.
  • Multitask shared parameter preprocessing.
  • Model branch selection and serialization.
  • Shared descriptor/fitting behavior.
  • Mixed-system stat collection with different atom counts.

Validation

  • pre-commit.ci - pr: passing
  • GitHub Actions checks are tracked on this PR

Romarin87 and others added 20 commits April 10, 2026 21:47
@Romarin87
fix(jax): preserve Hessian path when training from init_model/restart…
a06e6d7 
… checkpoints
@njzjz-bot @Romarin87
ci: align package_c TensorFlow image and wheel to 2.20 (deepmodeling#…
94026fc 
…5297)

## Summary
- switch the `package_c` Docker build image from TensorFlow 2.18 to 2.20
- pin the installed TensorFlow wheel to `==2.20.*` via the existing
`TENSORFLOW_VERSION` env

## Why
The failing `Build C library` job appears to be in `package_c.yml`, not
in the `cibuildwheel` workflow changed by dependabot.

This change makes the Docker image and installed TensorFlow version move
together, so we can test whether the failure was caused by a TensorFlow
/ GLIBC mismatch from pulling a newer wheel into an older build image.

Authored by OpenClaw (model: gpt-5.4)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Chores**
* Updated C library build configuration to target TensorFlow 2.20,
aligning the build matrix with the newer TensorFlow runtime.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

(cherry picked from commit 084edd6)
@njzjz @Romarin87
fix(CI): pin protobuf<7.34 (deepmodeling#5272)
65e8dad 
See
tensorflow/tensorflow@23f7b26

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Chores**
* Updated TensorFlow versions and added protobuf constraints for
improved stability across CPU and GPU configurations.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
(cherry picked from commit 21d2b1a)
@dependabot @Romarin87
build(deps): bump pypa/cibuildwheel from 3.3 to 3.4 (deepmodeling#5290)
a6f6060 
Bumps [pypa/cibuildwheel](https://github.com/pypa/cibuildwheel) from 3.3
to 3.4.
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a
href="https://github.com/pypa/cibuildwheel/releases">pypa/cibuildwheel's
releases</a>.</em></p>
<blockquote>
<h2>v3.4.0</h2>
<ul>
<li>🌟 You can now build wheels using <code>uv</code> as a build
frontend. This should improve performance, especially if your project
has lots of build dependencies. To use, set <a
href="https://cibuildwheel.pypa.io/en/stable/options/#build-frontend"><code>build-frontend</code></a>
to <code>uv</code>. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2322">#2322</a>)</li>
<li>⚠️ We no longer support running on Travis CI. It may continue
working but we don't run tests there anymore so we can't be sure. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2682">#2682</a>)</li>
<li>✨ Improvements to building rust wheels on Android (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2650">#2650</a>)</li>
<li>🐛 Fix bug with the GitHub Action on Windows, where PATH was getting
unnecessarily changed, causing issues with meson builds. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2723">#2723</a>)</li>
<li>✨ Add support for quiet setting on <code>build</code> and
<code>uv</code> from the cibuildwheel <code>build-verbosity</code>
setting. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2737">#2737</a>)</li>
<li>📚 Docs updates, including guidance on using Meson on Windows (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2718">#2718</a>)</li>
</ul>
<h2>v3.3.1</h2>
<ul>
<li>🛠 Update dependencies and container pins, including updating to
CPython 3.14.2. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2708">#2708</a>)</li>
</ul>
</blockquote>
</details>
<details>
<summary>Changelog</summary>
<p><em>Sourced from <a
href="https://github.com/pypa/cibuildwheel/blob/main/docs/changelog.md">pypa/cibuildwheel's
changelog</a>.</em></p>
<blockquote>
<hr />
<h2>title: Changelog</h2>
<h1>Changelog</h1>
<h3>v3.4.0</h3>
<p><em>5 March 2026</em></p>
<ul>
<li>🌟 You can now build wheels using <code>uv</code> as a build
frontend. This should improve performance, especially if your project
has lots of build dependencies. To use, set <a
href="https://cibuildwheel.pypa.io/en/stable/options/#build-frontend"><code>build-frontend</code></a>
to <code>uv</code>. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2322">#2322</a>)</li>
<li>⚠️ We no longer support running on Travis CI. It may continue
working but we don't run tests there anymore so we can't be sure. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2682">#2682</a>)</li>
<li>✨ Improvements to building rust wheels on Android (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2650">#2650</a>)</li>
<li>🐛 Fix bug with the GitHub Action on Windows, where PATH was getting
unnecessarily changed, causing issues with meson builds. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2723">#2723</a>)</li>
<li>✨ Add support for quiet setting on <code>build</code> and
<code>uv</code> from the cibuildwheel <code>build-verbosity</code>
setting. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2737">#2737</a>)</li>
<li>📚 Docs updates, including guidance on using Meson on Windows (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2718">#2718</a>)</li>
</ul>
<h3>v3.3.1</h3>
<p><em>5 January 2026</em></p>
<ul>
<li>🛠 Update dependencies and container pins, including updating to
CPython 3.14.2. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2708">#2708</a>)</li>
</ul>
<h3>v3.3.0</h3>
<p><em>12 November 2025</em></p>
<ul>
<li>🐛 Fix an incompatibility with Docker v29 (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2660">#2660</a>)</li>
<li>✨ Adds <code>test-runtime</code> option, to customise how tests on
simulated/emulated environments are run (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2636">#2636</a>)</li>
<li>✨ Adds support for new <code>manylinux_2_35</code> images on 32-bit
ARM <code>armv7l</code>, offering better C++20 compatibility (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2656">#2656</a>)</li>
<li>✨ <code>build[uv]</code> is now supported on Android (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2587">#2587</a>)</li>
<li>✨ You can now install extras (such as <code>uv</code>) with a simple
option on the GitHub Action (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2630">#2630</a>)</li>
<li>✨ <code>{project}</code> and <code>{package}</code> placeholders are
now supported in <code>repair-wheel-command</code> (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2589">#2589</a>)</li>
<li>🛠 The versions set with <code>dependency-versions</code> no longer
constrain packages specified by your <code>build-system.requires</code>.
Previously, on platforms other than Linux, the constraints in this
option would remain in the environment during the build. This has been
tidied up make behaviour more consistent between platforms, and to
prevent version conflicts. (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2583">#2583</a>)</li>
<li>🛠 Improve the handling of <code>test-command</code> on Android,
enabling more options to be passed (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2590">#2590</a>)</li>
<li>📚 Docs improvements (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2618">#2618</a>)</li>
</ul>
<h3>v3.2.1</h3>
<p><em>12 October 2025</em></p>
<ul>
<li>🛠 Update to CPython 3.14.0 final (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2614">#2614</a>)</li>
<li>🐛 Fix the default MACOSX_DEPLOYMENT_TARGET on Python 3.14 (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2613">#2613</a>)</li>
<li>📚 Docs improvements (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2617">#2617</a>)</li>
</ul>
<h3>v3.2.0</h3>
<p><em>22 September 2025</em></p>
<ul>
<li>✨ Adds GraalPy v25 (Python 3.12) support (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2597">#2597</a>)</li>
</ul>
<!-- raw HTML omitted -->
</blockquote>
<p>... (truncated)</p>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/ee02a1537ce3071a004a6b08c41e72f0fdc42d9a"><code>ee02a15</code></a>
Bump version: v3.4.0</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/f08ce70d7891b3a6b3d41f7909a69816eeea643b"><code>f08ce70</code></a>
chore: match copyright to BSD-2-clause template (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2758">#2758</a>)</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/95b4b79b96bae59effea03a5f297b0cb5aa31b6b"><code>95b4b79</code></a>
Fix incorrect document regarding pyodide auditwheel (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2752">#2752</a>)</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/f046d0a94dc550b5dccb93ae748704599fa1f9ae"><code>f046d0a</code></a>
Bump to Pyodide v0.29.3 (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2743">#2743</a>)</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/f5a31ee79d1dfd39a9a0a87d7cd412a29afb8fe1"><code>f5a31ee</code></a>
chore(deps): bump the actions group across 1 directory with 4 updates
(<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2755">#2755</a>)</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/35ec4565f5ca3474c283a0852fdf512d5817850f"><code>35ec456</code></a>
[pre-commit.ci] pre-commit autoupdate (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2756">#2756</a>)</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/51fed90270069c1f2ae3c075bcd6b0dad1d739c2"><code>51fed90</code></a>
[Bot] Update dependencies (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2734">#2734</a>)</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/b70562f37a2c65580ef235073974d067233fa82a"><code>b70562f</code></a>
Debug GraalPy/uv test failures on Windows (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2741">#2741</a>)</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/0a85e29549d1fd12d1ee07727cb55c39e16173a5"><code>0a85e29</code></a>
Work around CDN rate limiting on Python.org in bin/update_pythons.py (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2753">#2753</a>)</li>
<li><a
href="https://github.com/pypa/cibuildwheel/commit/2d33864bafd35e54234e57c036c38cd59288543c"><code>2d33864</code></a>
Pin the actions/setup-python for the externally used action.yml (<a
href="https://redirect.github.com/pypa/cibuildwheel/issues/2749">#2749</a>)</li>
<li>Additional commits viewable in <a
href="https://github.com/pypa/cibuildwheel/compare/v3.3...v3.4">compare
view</a></li>
</ul>
</details>
<br />

[![Dependabot compatibility
score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=pypa/cibuildwheel&package-manager=github_actions&previous-version=3.3&new-version=3.4)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores)

Dependabot will resolve any conflicts with this PR as long as you don't
alter it yourself. You can also trigger a rebase manually by commenting
`@dependabot rebase`.

[//]: # (dependabot-automerge-start)
[//]: # (dependabot-automerge-end)

---

<details>
<summary>Dependabot commands and options</summary>
<br />

You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits
that have been made to it
- `@dependabot show <dependency name> ignore conditions` will show all
of the ignore conditions of the specified dependency
- `@dependabot ignore this major version` will close this PR and stop
Dependabot creating any more for this major version (unless you reopen
the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop
Dependabot creating any more for this minor version (unless you reopen
the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop
Dependabot creating any more for this dependency (unless you reopen the
PR or upgrade to it yourself)

</details>

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
(cherry picked from commit 24e54bf)
@Romarin87
Backport devel CI dependency fixes to parallel (#2)
b11b495 
Backports CI baseline fixes from devel into parallel:

- align package_c TensorFlow build image and wheel to 2.20
- pin protobuf<7.34 for TensorFlow CI compatibility
- bump pypa/cibuildwheel from 3.3 to 3.4
@Romarin87
perf(jax): checkpoint RepFlow blocks for EFH memory
10821f2
@Romarin87
feat(jax): add single-task finetune support
bdc8765
@Romarin87
Implement JAX multitask training and finetune fixes
3a1edd3
@Romarin87
feat(jax): align finetune and multitask PT parity
52bc6d2
@Romarin87
chore(jax): clean up PR review nits
4b78c5b
@Romarin87
fix(jax): restore descriptor-safe repflow execution
73f64db
@Romarin87
fix(jax): keep repflow residual deserialize backend-neutral
6ec03c2
@Romarin87
fix(jax): guard optional nnx eager sharding API
cfe0920
@Romarin87
fix(jax): accept legacy no-mesh sharding errors
c11ac80
@Romarin87
fix(jax): support mesh setup on older jax
b4d68aa
@Romarin87
fix(jax): avoid codeql mesh context warning
dea77ee
@Romarin87
fix(jax): avoid PartitionSpec device_put on old jax
ebffb39
@Romarin87
fix(jax): shim missing monitoring record_scalar
81f4bd3
@Romarin87
Merge pull request #1 from Romarin87/codex/jax-repflow-remat-memory
587353b 
feat(jax): align finetune and multitask PT parity
@Romarin87
Fix mixed-system stat sampling
f8c0730
njzjz-bot and others added 7 commits June 26, 2026 13:56
@njzjz-bot @Romarin87
chore(pre-commit): remove flake8 hook (deepmodeling#5324)
c0ec0a6 
Problem
- flake8 is redundant with ruff in pre-commit and adds extra
dependency/latency.

Change
- Remove the flake8 repo hook from .pre-commit-config.yaml.

Authored by OpenClaw (model: gpt-5.2)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Chores**
* Removed Flake8 linting tool from the development pre-commit pipeline.
This change reduces external dependencies and maintenance overhead for
developers. The pre-commit pipeline continues to function normally with
other configured quality checks remaining active.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
7da00ff 
for more information, see https://pre-commit.ci
@Romarin87
style(jax): satisfy ruff on parallel PR
e5235ed
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
167fe58 
for more information, see https://pre-commit.ci
@Romarin87
fix(dpmodel): satisfy explicit array creation checker
53fe672
@Romarin87
fix(dpmodel): keep arange trace-safe for jax
8f3bc9e
@Romarin87
fix(jax): pack finetune stats with real natoms
44ac6c1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

Python size/XXL

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@Romarin87 @njzjz-bot @njzjz