fix(train): share post-training bias adjustment#5745
Conversation
📝 WalkthroughWalkthroughAdds shared output-bias adjustment helpers, exports them publicly, and updates pt, pt_expt, tf2, and jax training flows to use them. JAX and pt_expt also add post-training bias adjustment paths, checkpoint saving, and distributed-state synchronization. ChangesShared bias-adjustment helper and backend integration
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/jax/train/trainer.py`:
- Around line 736-747: The JAX trainer’s run method should not default to
self.training_tasks because that attribute is never set in this path and can
raise AttributeError when tasks is omitted. Update run in the Trainer class to
require an explicit tasks value or otherwise source it from a valid initialized
field, and remove the fallback to self.training_tasks. Keep the existing train()
call site working by passing the task collection explicitly through the run
flow.
- Around line 749-756: The post-training bias adjustment in
_change_bias_after_training currently runs independently on every rank, which
can produce divergent checkpoint state. Update this method to execute
change_model_out_bias_by_task only on rank 0, then broadcast the recomputed
bias/state to all other ranks before saving. Use the existing trainer/JAX rank
utilities around _change_bias_after_training, self.models, and
change_model_out_bias_by_task so the final checkpoint is consistent across
hosts.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: a48da21d-7f31-4b57-ba98-ece2543894f2
📒 Files selected for processing (7)
deepmd/dpmodel/train/__init__.pydeepmd/dpmodel/train/trainer.pydeepmd/jax/train/trainer.pydeepmd/pt/train/training.pydeepmd/pt_expt/train/training.pydeepmd/tf2/train/trainer.pysource/tests/common/test_dpmodel_train.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5745 +/- ##
==========================================
- Coverage 80.87% 79.47% -1.41%
==========================================
Files 1003 1014 +11
Lines 112492 115392 +2900
Branches 4236 4274 +38
==========================================
+ Hits 90981 91709 +728
- Misses 19986 22139 +2153
- Partials 1525 1544 +19 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/jax/test_training.py (1)
354-405: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGood coverage of chief vs. peer broadcast paths; consider also covering the single-rank early return.
Both tests correctly exercise the
world_size=2branch of_change_bias_after_trainingThe JAX trainer optionally recomputes output bias on the chief rank, then broadcasts per-model pure state via multihost_utils.broadcast_one_to_all with is_source tied to is_chief and updates nnx models from the broadcasted state., but_bias_sync_traineralways setsworld_size=2, so the early-return branch (if self.rank_context.world_size <= 1: return) added in the same change is never exercised — a single-process run would skip broadcasting entirely and this isn't asserted anywhere.Consider adding a third test with
world_size=1assertingbroadcast_one_to_allis not called afterchange_model_out_bias_by_taskruns.✅ Suggested additional test
+def test_jax_change_bias_after_training_skips_broadcast_single_rank() -> None: + """Single-process runs recompute bias without attempting a broadcast.""" + trainer = _bias_sync_trainer(rank=0) + trainer.rank_context = RankContext(rank=0, world_size=1) + + with ( + patch( + "deepmd.jax.train.trainer.change_model_out_bias_by_task", + ) as change_model_out_bias_by_task, + patch( + "jax.experimental.multihost_utils.broadcast_one_to_all", + ) as broadcast_one_to_all, + ): + trainer._change_bias_after_training() + + change_model_out_bias_by_task.assert_called_once() + broadcast_one_to_all.assert_not_called()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/jax/test_training.py` around lines 354 - 405, Add a test for the single-rank early return in _change_bias_after_training: create a _bias_sync_trainer setup with world_size=1, call trainer._change_bias_after_training(), and assert the method returns without invoking jax.experimental.multihost_utils.broadcast_one_to_all. Keep the existing chief/peer tests unchanged, and reuse the same symbols (_bias_sync_trainer, _change_bias_after_training, change_model_out_bias_by_task, broadcast_one_to_all) so the new coverage targets the new early-return branch.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@source/tests/jax/test_training.py`:
- Around line 354-405: Add a test for the single-rank early return in
_change_bias_after_training: create a _bias_sync_trainer setup with
world_size=1, call trainer._change_bias_after_training(), and assert the method
returns without invoking jax.experimental.multihost_utils.broadcast_one_to_all.
Keep the existing chief/peer tests unchanged, and reuse the same symbols
(_bias_sync_trainer, _change_bias_after_training, change_model_out_bias_by_task,
broadcast_one_to_all) so the new coverage targets the new early-return branch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 770746c0-fb81-4b2d-b2c1-56df16865cf0
📒 Files selected for processing (2)
deepmd/jax/train/trainer.pysource/tests/jax/test_training.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/jax/train/trainer.py
Summary
Tests
Summary by CodeRabbit