Skip to content

Conversation

Copy link
Contributor

Copilot AI commented Sep 3, 2025

This PR fixes the issue where dp --pt change-bias creates unnecessarily large model files for multi-task models and requires head selection when loading the resulting model.

Problem

When using dp --pt change-bias on multi-task models, two issues occurred:

  1. File size bloat: The resulting model was as large as or larger than the original multi-task model, even though only one head was being modified
  2. Head selection requirement: The resulting model still required specifying a head when loading, unlike fine-tuned models with numb_steps: 0

For example, a 27MB multi-task model would produce a 26MB output file after change-bias, when it should produce a ~8MB single-head model like fine-tuning does.

Root Cause

The change_bias function in deepmd/pt/entrypoints/main.py was:

  1. Deep copying the entire multi-task model state dict
  2. Modifying only the selected branch's bias
  3. Saving the entire multi-task model with all heads instead of just the selected branch
  4. Including the full optimizer state from the original multi-task model (17.8MB in test case)

Solution

Modified the saving logic to:

  1. Remove unnecessary deep copy that caused memory bloat
  2. For multi-task models: Extract and save only the selected branch as a single-head model
  3. Exclude optimizer state for multi-task models to prevent file bloat
  4. For single-task models: Preserve existing behavior unchanged including optimizer state
  5. Maintain compatibility: Preserve train_infos and other required metadata

Results

  • File size reduction: ~70% smaller output files (e.g., 27MB → 7.9MB in tests)
  • Model loading: Converted models can be loaded without requiring head selection
  • Backward compatibility: All existing single-task change-bias functionality unchanged
  • Consistency: Multi-task change-bias now behaves like fine-tuning with numb_steps: 0

The fix ensures that dp --pt change-bias on multi-task models produces appropriately sized single-head models that can be used directly without head selection.

Fixes #4348.


✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.

Copilot AI changed the title [WIP] [BUG] dp change-bias will give much large model fix(pt): optimize change-bias for multi-task models to save only selected branch Sep 3, 2025
Copilot AI requested a review from njzjz September 3, 2025 14:30
Copilot finished work on behalf of njzjz September 3, 2025 14:30
@njzjz
Copy link
Member

njzjz commented Sep 3, 2025

@copilot The original issue is not fixed. Reproduce: Go to examples/water_multi_task/pytorch_example and train a model, call dp --pt change-bias for this model, and then the size of the model increases.

Copy link
Contributor Author

Copilot AI commented Sep 3, 2025

@copilot The original issue is not fixed. Reproduce: Go to examples/water_multi_task/pytorch_example and train a model, call dp --pt change-bias for this model, and then the size of the model incr...

I've reproduced and fixed the issue. The problem was that the previous fix was still including the optimizer state from the original multi-task model, which took up most of the file size. The fix now excludes optimizer state for multi-task models, resulting in a 70% file size reduction (27MB → 7.9MB in the test case). Fixed in commit 5028e76.

Copilot AI changed the title fix(pt): optimize change-bias for multi-task models to save only selected branch fix(pt): remove optimizer state from multi-task change-bias to reduce file size Sep 3, 2025
Copilot finished work on behalf of njzjz September 3, 2025 15:07
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.

[BUG] dp change-bias will give much large model

2 participants