Three center integration and numerical stabilities of trinity - #366
Open
floatingCatty wants to merge 2 commits into
Open
Three center integration and numerical stabilities of trinity#366floatingCatty wants to merge 2 commits into
floatingCatty wants to merge 2 commits into
Conversation
|
This advisory review plan was generated from changed file names using trusted base-branch code. DeePTB PR Review Plan / DeePTB PR 审查计划Risk / 风险等级: High (高) · Changed files / 变更文件: 14 Why / 风险来源
Recommended Review / 建议审查重点
Detailed risk areas
Human review focus
Local commands and hold conditionsSuggested local commands:
Hold conditions:
Advisory only. / 仅作为审查辅助。 |
Contributor
📝 WalkthroughWalkthroughThis PR adds opt-in spectral balancing to Trinity, trainer hardening with EMA, gradient clipping, grouped learning rates, and warmup-cosine scheduling. It also improves device-safe tensor handling, adds a Huber Hamiltonian loss, updates checkpoint behavior, and introduces regression and workflow tests. ChangesStability and training enhancements
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TrainingLoop
participant Optimizer
participant EMA
participant Validation
participant Checkpoint
TrainingLoop->>Optimizer: compute gradients and optimizer.step()
TrainingLoop->>EMA: update shadow parameters
Validation->>EMA: average_parameters()
EMA-->>Validation: EMA weights
Checkpoint->>EMA: average_parameters()
EMA-->>Checkpoint: EMA model state
Checkpoint-->>TrainingLoop: raw and EMA checkpoint state
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches⚔️ Resolve merge conflicts
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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (1)
dptb/nn/embedding/trinity.py (1)
1013-1020: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
node_irreps_infor the node feature LayerNorm.
sln_nnormalizes node features, but it is currently initialized with the edge irreps (self.irreps_in) instead of the node irreps (self.node_irreps_in). SinceLayercurrently passes the exact same irreps for both arguments, this doesn't crash at runtime, but it introduces a latent bug that will surface if node and edge irreps ever diverge.♻️ Proposed refactor
- self.sln_n = SeperableLayerNorm( - irreps=self.irreps_in, + self.sln_n = SeperableLayerNorm( + irreps=self.node_irreps_in, eps=5e-3, affine=True, normalization='component', std_balance_degrees=True, per_l=spectral_balance,🤖 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 `@dptb/nn/embedding/trinity.py` around lines 1013 - 1020, Update the SeperableLayerNorm initialization for sln_n to use self.node_irreps_in instead of self.irreps_in, while leaving the edge-feature normalization configuration unchanged.
🤖 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 `@docs/adr/0005-trinity-sota-enhancements.md`:
- Around line 188-210: Revise the compatibility statement introducing the
train_options flags so it excludes allow_tf32 from the claim that all defaults
preserve previous behavior. Keep the allow_tf32 entry explicit that its default
is false and therefore changes the prior TF32 behavior.
In `@dptb/nn/threecenter.py`:
- Around line 177-181: Update ThreeCenterFactorized so compute_P, to_feature,
to_reduced, and forward derive dtype and device from a moved tensor such as the
er_max buffer or an existing parameter, rather than cached
self.dtype/self.device attributes; update the affected test in
dptb/tests/test_threecenter.py lines 487-506 only as needed to verify behavior
after module device or dtype moves. In dptb/nn/threecenter.py lines 177-181,
remove or stop relying on stale plain attributes while preserving the registered
er_max buffer.
In `@dptb/nnops/ema.py`:
- Around line 41-45: Update EMA.copy_to to validate that an explicitly supplied
parameters collection has the same count as the EMA shadow parameters before
copying; reject mismatches rather than allowing zip to silently truncate, while
preserving the existing filtering and copy behavior for matching collections.
In `@dptb/nnops/loss.py`:
- Around line 410-425: Add the missing _elem_loss method to EigHamLoss before
forward, matching the loss computation used by HamilLossAbs: combine loss1 and
the square root of loss2 with the 0.5 weighting, returning a torch.Tensor so all
existing _elem_loss calls in forward work.
- Around line 768-770: Update the loss module initialization to register
onsite_weight and hopping_weight via PyTorch’s register_buffer instead of plain
tensor attributes, preserving their existing initialization values and device
while ensuring model.to(device) moves both buffers with the module.
In `@dptb/nnops/trainer.py`:
- Around line 224-233: Update the checkpoint restore logic around trainer.ema
and raw_model_state_dict so raw training weights are loaded whenever the
checkpoint contains them, regardless of whether trainer.ema is enabled. After an
EMA state restore fails in the existing ValueError handler, reinitialize the EMA
shadow from the restored raw model weights, preserving deployment-weight
restoration for successful EMA loads.
- Around line 214-222: Update the checkpoint restoration loop in Trainer so
optimizer and scheduler state are restored as a single atomic operation: stage
or validate both states before applying either, and if either load_state_dict
call raises ValueError, RuntimeError, or KeyError, discard both restored states
and retain fresh optimizer and scheduler instances. Preserve the existing
warning behavior while ensuring no partially restored pairing remains.
In `@dptb/tests/test_threecenter.py`:
- Line 553: In the test setup around sln0, split the chained assignment
statement into separate statements so the code complies with Ruff E702, while
preserving the existing values of y and off.
In `@dptb/tests/test_trainer_hardening.py`:
- Around line 283-296: Correct the restart assertion in the test around
Saver._save and Trainer.restart so it expects the next unexecuted iteration
without validating an extra increment. Use the actual iteration-plugin save path
if needed, or assert that the restarted trainer.iter equals the value stored by
_save; keep the EMA restoration checks unchanged.
- Around line 100-114: Update test_ema_state_dict_roundtrip_and_mismatch to use
a model exposing a different number of parameter tensors for bad, rather than
another Linear layer with two parameters. Keep the existing ValueError assertion
so the test exercises the EMA parameter-count mismatch path.
In `@dptb/utils/argcheck.py`:
- Line 360: The warmup_cos scheduler requires per-iteration learning-rate
updates but can currently be used with the trainer’s per-epoch default. In
dptb/utils/argcheck.py lines 360-360, reject warmup_cos unless
update_lr_per_iter=True or enable that setting automatically; in
dptb/utils/tools.py lines 231-248, make the scheduler/trainer stepping contract
explicit and add coverage for the default configuration.
In `@dptb/utils/tools.py`:
- Around line 168-172: The group_key function must preserve the layer index for
embedding layer names. Update its embedding branch to return a key containing
the first three components, such as embedding.layers.0, while retaining the
existing fallback behavior for other names.
---
Nitpick comments:
In `@dptb/nn/embedding/trinity.py`:
- Around line 1013-1020: Update the SeperableLayerNorm initialization for sln_n
to use self.node_irreps_in instead of self.irreps_in, while leaving the
edge-feature normalization configuration unchanged.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: aeaa15dd-e7bd-4ce0-b161-91306f39cf04
📒 Files selected for processing (14)
docs/adr/0005-trinity-sota-enhancements.mddptb/nn/cutoff.pydptb/nn/embedding/trinity.pydptb/nn/norm.pydptb/nn/tensor_product.pydptb/nn/threecenter.pydptb/nnops/ema.pydptb/nnops/loss.pydptb/nnops/trainer.pydptb/plugins/saver.pydptb/tests/test_threecenter.pydptb/tests/test_trainer_hardening.pydptb/utils/argcheck.pydptb/utils/tools.py
Comment on lines
+188
to
+210
| Shipped as `train_options` flags, all defaulting to the previous behaviour so existing configs and | ||
| checkpoints are unaffected: | ||
|
|
||
| - `per_group_lr: true` (**P2**) — `build_wrms_param_groups` (`dptb/utils/tools.py`) splits the | ||
| optimizer into per-block groups with lr scaled by each block's weight RMS (trust ratio at init, | ||
| clamped to [0.02, 1.0], referenced to the median group RMS). On a full-mode Trinity this puts the | ||
| small-init AtomicResNet heads (`edge_prediction_h2`, `edge_prediction_s`, |w|_rms≈0.05) at lr scale | ||
| ≈0.077 while the O(1) embedding stack stays at 1.0 — a **13× measured spread**, exactly the | ||
| imbalance the audit flagged. Scales compose correctly with any scheduler (each group's base_lr is | ||
| scaled independently). Verified in `test_wrms_param_groups_*`. | ||
| - `grad_clip_norm: <float>` (**P3**) — global-norm gradient clipping each step | ||
| (`clip_grad_norm_`); off at 0.0. Tames the loss spikes typical of batch_size=1 Hamiltonian fitting. | ||
| - `ema_decay: <float>` (**P3**) — `ExponentialMovingAverage` (`dptb/nnops/ema.py`) of the weights, | ||
| with the standard `(1+t)/(10+t)` warmup on the effective decay. **Validation scores and Saver | ||
| checkpoints use the averaged weights** (`model_state_dict` = EMA), while the raw training weights | ||
| are stored under `raw_model_state_dict` for exact restart. This removes the single-iteration noise | ||
| from best-checkpoint selection (the failure mode where stage-2 `best.pth` was honestly worse than | ||
| `latest`). Verified end-to-end in `test_trainer_engages_hardening_and_checkpoints_ema`. | ||
| - `lr_scheduler.type: warmup_cos` (**P3**) — linear warmup then cosine to `eta_min` | ||
| (`get_lr_scheduler`), meant to be stepped per-iteration (`update_lr_per_iter: true`). Replaces the | ||
| RoP-to-floor collapse. Verified in `test_warmup_cosine_shape`. | ||
| - `allow_tf32: false` (**P3**, default) — TF32 matmuls are now explicitly disabled unless requested; | ||
| Hamiltonian targets need the precision. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Clarify that the TF32 default changes prior behavior.
“All defaulting to the previous behaviour” conflicts with Lines 209-210, which say TF32 is now explicitly disabled. Exclude allow_tf32 from that compatibility claim.
🤖 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 `@docs/adr/0005-trinity-sota-enhancements.md` around lines 188 - 210, Revise
the compatibility statement introducing the train_options flags so it excludes
allow_tf32 from the claim that all defaults preserve previous behavior. Keep the
allow_tf32 entry explicit that its default is false and therefore changes the
prior TF32 behavior.
Comment on lines
+177
to
+181
| self.bessel = BesselBasis(r_max=torch.tensor(er_max, dtype=dtype, device=device), num_basis=n_radial_basis, trainable=True) | ||
| # register as a buffer (with device) so model.to(device) moves it; a plain tensor attribute | ||
| # stays on CPU and breaks the polynomial_cutoff below when the model runs on GPU. Non-persistent: | ||
| # it is fully determined by the `er_max` config, so it stays out of the checkpoint (no layout change). | ||
| self.register_buffer("er_max", torch.tensor(er_max, dtype=dtype, device=device), persistent=False) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
sed -n '1,260p' dptb/nn/threecenter.pyRepository: deepmodeling/DeePTB
Length of output: 14855
🏁 Script executed:
sed -n '430,560p' dptb/tests/test_threecenter.pyRepository: deepmodeling/DeePTB
Length of output: 7509
🏁 Script executed:
rg -n "class BesselBasis|def forward|self\._cg|reduced_values\(|compute_P\(" dptb/nn/threecenter.py dptb/nn/radial_basis.pyRepository: deepmodeling/DeePTB
Length of output: 1219
🏁 Script executed:
sed -n '1,260p' dptb/nn/radial_basis.pyRepository: deepmodeling/DeePTB
Length of output: 4273
🏁 Script executed:
sed -n '260,520p' dptb/nn/threecenter.pyRepository: deepmodeling/DeePTB
Length of output: 7686
🏁 Script executed:
rg -n "self\.device\s*=|self\.dtype\s*=|def to\(" dptb/nn/threecenter.py dptb/nn/*.pyRepository: deepmodeling/DeePTB
Length of output: 2149
🏁 Script executed:
python3 - <<'PY'
from pathlib import Path
text = Path("dptb/nn/threecenter.py").read_text()
for needle in ["self.dtype, self.device =", "self._cg =", "dtype=self.dtype, device=self.device", "self.device)", "self.dtype)"]:
print(needle, text.count(needle))
PYRepository: deepmodeling/DeePTB
Length of output: 268
ThreeCenterFactorized still caches dtype/device as plain attributes. nn.Module.to(...) moves the buffers and parameters, but it will not update self.device/self.dtype, so later allocations in compute_P, to_feature, to_reduced, and forward can still use the stale device after a move. Derive them from an existing tensor instead.
📍 Affects 2 files
dptb/nn/threecenter.py#L177-L181(this comment)dptb/tests/test_threecenter.py#L487-L506
🤖 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 `@dptb/nn/threecenter.py` around lines 177 - 181, Update ThreeCenterFactorized
so compute_P, to_feature, to_reduced, and forward derive dtype and device from a
moved tensor such as the er_max buffer or an existing parameter, rather than
cached self.dtype/self.device attributes; update the affected test in
dptb/tests/test_threecenter.py lines 487-506 only as needed to verify behavior
after module device or dtype moves. In dptb/nn/threecenter.py lines 177-181,
remove or stop relying on stale plain attributes while preserving the registered
er_max buffer.
Comment on lines
+41
to
+45
| def copy_to(self, parameters=None): | ||
| parameters = self._params if parameters is None else [p for p in parameters if p.requires_grad] | ||
| with torch.no_grad(): | ||
| for s, p in zip(self.shadow, parameters): | ||
| p.copy_(s) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Reject partial EMA copies when parameter counts differ.
An explicit parameters collection with a different length is silently truncated by zip, leaving only part of the target model updated.
Proposed fix
def copy_to(self, parameters=None):
parameters = self._params if parameters is None else [p for p in parameters if p.requires_grad]
+ if len(parameters) != len(self.shadow):
+ raise ValueError(
+ f"EMA tracks {len(self.shadow)} parameters, but received {len(parameters)}"
+ )
with torch.no_grad():
for s, p in zip(self.shadow, parameters):
p.copy_(s)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def copy_to(self, parameters=None): | |
| parameters = self._params if parameters is None else [p for p in parameters if p.requires_grad] | |
| with torch.no_grad(): | |
| for s, p in zip(self.shadow, parameters): | |
| p.copy_(s) | |
| def copy_to(self, parameters=None): | |
| parameters = self._params if parameters is None else [p for p in parameters if p.requires_grad] | |
| if len(parameters) != len(self.shadow): | |
| raise ValueError( | |
| f"EMA tracks {len(self.shadow)} parameters, but received {len(parameters)}" | |
| ) | |
| with torch.no_grad(): | |
| for s, p in zip(self.shadow, parameters): | |
| p.copy_(s) |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 44-44: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
🤖 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 `@dptb/nnops/ema.py` around lines 41 - 45, Update EMA.copy_to to validate that
an explicitly supplied parameters collection has the same count as the EMA
shadow parameters before copying; reject mismatches rather than allowing zip to
silently truncate, while preserving the existing filtering and copy behavior for
matching collections.
Source: Linters/SAST tools
Comment on lines
410
to
+425
| pre = data[AtomicDataDict.NODE_FEATURES_KEY][self.idp.mask_to_nrme[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()]] | ||
| tgt = ref_data[AtomicDataDict.NODE_FEATURES_KEY][self.idp.mask_to_nrme[ref_data[AtomicDataDict.ATOM_TYPE_KEY].flatten()]] | ||
| onsite_loss = 0.5*(self.loss1(pre, tgt) + torch.sqrt(self.loss2(pre, tgt))) | ||
| onsite_loss = self._elem_loss(pre, tgt) | ||
|
|
||
| pre = data[AtomicDataDict.EDGE_FEATURES_KEY][self.idp.mask_to_erme[data[AtomicDataDict.EDGE_TYPE_KEY].flatten()]] | ||
| tgt = ref_data[AtomicDataDict.EDGE_FEATURES_KEY][self.idp.mask_to_erme[ref_data[AtomicDataDict.EDGE_TYPE_KEY].flatten()]] | ||
| hopping_loss = 0.5*(self.loss1(pre, tgt) + torch.sqrt(self.loss2(pre, tgt))) | ||
| hopping_loss = self._elem_loss(pre, tgt) | ||
|
|
||
| if self.overlap: | ||
| pre = data[AtomicDataDict.EDGE_OVERLAP_KEY][self.idp.mask_to_erme[data[AtomicDataDict.EDGE_TYPE_KEY].flatten()]] | ||
| tgt = ref_data[AtomicDataDict.EDGE_OVERLAP_KEY][self.idp.mask_to_erme[ref_data[AtomicDataDict.EDGE_TYPE_KEY].flatten()]] | ||
| overlap_loss = 0.5*(self.loss1(pre, tgt) + torch.sqrt(self.loss2(pre, tgt))) | ||
| overlap_loss = self._elem_loss(pre, tgt) | ||
|
|
||
| pre = data[AtomicDataDict.NODE_OVERLAP_KEY][self.idp.mask_to_nrme[data[AtomicDataDict.ATOM_TYPE_KEY].flatten()]] | ||
| tgt = ref_data[AtomicDataDict.NODE_OVERLAP_KEY][self.idp.mask_to_nrme[ref_data[AtomicDataDict.ATOM_TYPE_KEY].flatten()]] | ||
| overlap_loss += 0.5*(self.loss1(pre, tgt) + torch.sqrt(self.loss2(pre, tgt))) | ||
| overlap_loss += self._elem_loss(pre, tgt) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Missing _elem_loss method in EigHamLoss.
The forward method has been refactored to call self._elem_loss(pre, tgt), but this method was never added to EigHamLoss (unlike HamilLossAbs). This will cause an AttributeError during training.
🐛 Proposed fix
Define the missing _elem_loss method inside the EigHamLoss class before forward:
def _elem_loss(self, pre: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
return 0.5 * (self.loss1(pre, tgt) + torch.sqrt(self.loss2(pre, tgt)))🤖 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 `@dptb/nnops/loss.py` around lines 410 - 425, Add the missing _elem_loss method
to EigHamLoss before forward, matching the loss computation used by
HamilLossAbs: combine loss1 and the square root of loss2 with the 0.5 weighting,
returning a torch.Tensor so all existing _elem_loss calls in forward work.
Comment on lines
+768
to
770
| self.onsite_weight = torch.ones(idp.num_types, device=self.device) | ||
| self.hopping_weight = torch.ones(len(idp.bond_types), device=self.device) | ||
| if isinstance(onsite_weight, float) or isinstance(onsite_weight, int): |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Register stateful tensors as buffers.
onsite_weight and hopping_weight are initialized as standard attributes rather than registered buffers. If the loss module is moved to a different device (e.g., via model.to(device)) after initialization, these tensors will remain on the original device, leading to a device mismatch crash during the forward pass.
🔧 Proposed fix
Use register_buffer so PyTorch properly manages their device placement:
- self.onsite_weight = torch.ones(idp.num_types, device=self.device)
- self.hopping_weight = torch.ones(len(idp.bond_types), device=self.device)
+ self.register_buffer("onsite_weight", torch.ones(idp.num_types, device=self.device))
+ self.register_buffer("hopping_weight", torch.ones(len(idp.bond_types), device=self.device))(Note: PyTorch's __setattr__ behavior ensures that subsequent in-place modifications and reassignments to these attributes in this __init__ will safely update the registered buffer.)
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| self.onsite_weight = torch.ones(idp.num_types, device=self.device) | |
| self.hopping_weight = torch.ones(len(idp.bond_types), device=self.device) | |
| if isinstance(onsite_weight, float) or isinstance(onsite_weight, int): | |
| self.register_buffer("onsite_weight", torch.ones(idp.num_types, device=self.device)) | |
| self.register_buffer("hopping_weight", torch.ones(len(idp.bond_types), device=self.device)) | |
| if isinstance(onsite_weight, float) or isinstance(onsite_weight, int): |
🤖 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 `@dptb/nnops/loss.py` around lines 768 - 770, Update the loss module
initialization to register onsite_weight and hopping_weight via PyTorch’s
register_buffer instead of plain tensor attributes, preserving their existing
initialization values and device while ensuring model.to(device) moves both
buffers with the module.
|
|
||
| # unaffine per-l norm sends each l to ~unit RMS independently | ||
| sln0 = SeperableLayerNorm(irreps, eps=1e-6, affine=False, per_l=True, dtype=torch.float64) | ||
| y = sln0(x); off = 0 |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Split the statements to satisfy Ruff E702.
Proposed fix
- y = sln0(x); off = 0
+ y = sln0(x)
+ off = 0📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| y = sln0(x); off = 0 | |
| y = sln0(x) | |
| off = 0 |
🧰 Tools
🪛 Ruff (0.15.21)
[error] 553-553: Multiple statements on one line (semicolon)
(E702)
🤖 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 `@dptb/tests/test_threecenter.py` at line 553, In the test setup around sln0,
split the chained assignment statement into separate statements so the code
complies with Ruff E702, while preserving the existing values of y and off.
Source: Linters/SAST tools
Comment on lines
+100
to
+114
| def test_ema_state_dict_roundtrip_and_mismatch(): | ||
| torch.manual_seed(0) | ||
| m = torch.nn.Linear(3, 3) | ||
| ema = ExponentialMovingAverage(m.parameters(), decay=0.9) | ||
| ema.update() | ||
| sd = ema.state_dict() | ||
| ema2 = ExponentialMovingAverage(m.parameters(), decay=0.9) | ||
| ema2.load_state_dict(sd) | ||
| for a, b in zip(ema.shadow, ema2.shadow): | ||
| assert torch.allclose(a, b) | ||
| # loading into a model with a different #params must error, not silently corrupt | ||
| other = torch.nn.Linear(5, 5) | ||
| bad = ExponentialMovingAverage(other.parameters(), decay=0.9) | ||
| with pytest.raises(ValueError): | ||
| bad.load_state_dict(sd) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Exercise an actual EMA parameter-count mismatch.
Linear(3, 3) and Linear(5, 5) both expose two parameters, so this reaches the shape check rather than the asserted count check. Use a model with a different number of parameter tensors.
Proposed fix
- other = torch.nn.Linear(5, 5)
+ other = torch.nn.Sequential(
+ torch.nn.Linear(3, 3),
+ torch.nn.Linear(3, 3),
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def test_ema_state_dict_roundtrip_and_mismatch(): | |
| torch.manual_seed(0) | |
| m = torch.nn.Linear(3, 3) | |
| ema = ExponentialMovingAverage(m.parameters(), decay=0.9) | |
| ema.update() | |
| sd = ema.state_dict() | |
| ema2 = ExponentialMovingAverage(m.parameters(), decay=0.9) | |
| ema2.load_state_dict(sd) | |
| for a, b in zip(ema.shadow, ema2.shadow): | |
| assert torch.allclose(a, b) | |
| # loading into a model with a different #params must error, not silently corrupt | |
| other = torch.nn.Linear(5, 5) | |
| bad = ExponentialMovingAverage(other.parameters(), decay=0.9) | |
| with pytest.raises(ValueError): | |
| bad.load_state_dict(sd) | |
| def test_ema_state_dict_roundtrip_and_mismatch(): | |
| torch.manual_seed(0) | |
| m = torch.nn.Linear(3, 3) | |
| ema = ExponentialMovingAverage(m.parameters(), decay=0.9) | |
| ema.update() | |
| sd = ema.state_dict() | |
| ema2 = ExponentialMovingAverage(m.parameters(), decay=0.9) | |
| ema2.load_state_dict(sd) | |
| for a, b in zip(ema.shadow, ema2.shadow): | |
| assert torch.allclose(a, b) | |
| # loading into a model with a different `#params` must error, not silently corrupt | |
| other = torch.nn.Sequential( | |
| torch.nn.Linear(3, 3), | |
| torch.nn.Linear(3, 3), | |
| ) | |
| bad = ExponentialMovingAverage(other.parameters(), decay=0.9) | |
| with pytest.raises(ValueError): | |
| bad.load_state_dict(sd) |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 108-108: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
🤖 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 `@dptb/tests/test_trainer_hardening.py` around lines 100 - 114, Update
test_ema_state_dict_roundtrip_and_mismatch to use a model exposing a different
number of parameter tensors for bad, rather than another Linear layer with two
parameters. Keep the existing ValueError assertion so the test exercises the EMA
parameter-count mismatch path.
Comment on lines
+283
to
+296
| ckpt_dir = tmp_path / "ckpt"; ckpt_dir.mkdir() | ||
| saver = Saver(); saver.register(trainer, str(ckpt_dir)) | ||
| saver._save("trinity.iter", trainer.model, trainer.model.model_options, | ||
| trainer.common_options, trainer.train_options) | ||
| ckpt = str(ckpt_dir / "trinity.iter.pth") | ||
| iter_at_save = trainer.iter | ||
| ema_shadow_at_save = [s.clone() for s in trainer.ema.shadow] | ||
|
|
||
| # ---- restart: exact resume ---- (train.py always passes common_options) | ||
| r = Trainer.restart(checkpoint=ckpt, train_datasets=make_ds(), common_options=dict(common)) | ||
| assert r.iter == iter_at_save + 1 # resumes at the next iteration | ||
| assert r.ema is not None | ||
| for a, b in zip(r.ema.shadow, ema_shadow_at_save): | ||
| assert torch.allclose(a, b), "EMA shadow not restored on restart" |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Do not validate an extra iteration increment as an exact restart.
After two completed iterations, trainer.iter is already the next iteration. _save() stores that value, while restart adds another one; the assertion therefore codifies a skipped counter value. Exercise the real iteration-plugin save path or require the restarted value to equal the next unexecuted iteration.
🧰 Tools
🪛 Ruff (0.15.21)
[error] 283-283: Multiple statements on one line (semicolon)
(E702)
[error] 284-284: Multiple statements on one line (semicolon)
(E702)
[warning] 295-295: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
🤖 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 `@dptb/tests/test_trainer_hardening.py` around lines 283 - 296, Correct the
restart assertion in the test around Saver._save and Trainer.restart so it
expects the next unexecuted iteration without validating an extra increment. Use
the actual iteration-plugin save path if needed, or assert that the restarted
trainer.iter equals the value stored by _save; keep the EMA restoration checks
unchanged.
| Argument("linear", dict, LinearLR()), | ||
| Argument("rop", dict, ReduceOnPlateau(), doc="rop: reduce on plateau"), | ||
| Argument("cos", dict, CosineAnnealingLR(), doc="cos: cosine annealing"), | ||
| Argument("warmup_cos", dict, WarmupCosineLR(), doc="warmup_cos: linear warmup then cosine annealing (step per-iteration)"), |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
warmup_cos can silently run on the wrong time scale. The scheduler counts iterations, while the independently configured trainer defaults to stepping schedulers per epoch.
dptb/utils/argcheck.py#L360-L360: rejectwarmup_cosunlessupdate_lr_per_iter=True, or make the selection enable it automatically.dptb/utils/tools.py#L231-L248: make the stepping requirement explicit in the scheduler/trainer contract and cover the default configuration with a test.
📍 Affects 2 files
dptb/utils/argcheck.py#L360-L360(this comment)dptb/utils/tools.py#L231-L248
🤖 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 `@dptb/utils/argcheck.py` at line 360, The warmup_cos scheduler requires
per-iteration learning-rate updates but can currently be used with the trainer’s
per-epoch default. In dptb/utils/argcheck.py lines 360-360, reject warmup_cos
unless update_lr_per_iter=True or enable that setting automatically; in
dptb/utils/tools.py lines 231-248, make the scheduler/trainer stepping contract
explicit and add coverage for the default configuration.
Comment on lines
+168
to
+172
| def group_key(name): | ||
| parts = name.split(".") | ||
| if len(parts) >= 3 and parts[0] == "embedding": | ||
| return ".".join(parts[:2]) | ||
| return parts[0] |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Include the layer index in WRMS group keys.
For names such as embedding.layers.0.*, parts[:2] produces embedding.layers, combining every message-passing layer despite the documented per-layer grouping.
Proposed fix
def group_key(name):
parts = name.split(".")
- if len(parts) >= 3 and parts[0] == "embedding":
+ if len(parts) >= 3 and parts[:2] == ["embedding", "layers"]:
+ return ".".join(parts[:3])
+ if len(parts) >= 2 and parts[0] == "embedding":
return ".".join(parts[:2])
return parts[0]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def group_key(name): | |
| parts = name.split(".") | |
| if len(parts) >= 3 and parts[0] == "embedding": | |
| return ".".join(parts[:2]) | |
| return parts[0] | |
| def group_key(name): | |
| parts = name.split(".") | |
| if len(parts) >= 3 and parts[:2] == ["embedding", "layers"]: | |
| return ".".join(parts[:3]) | |
| if len(parts) >= 2 and parts[0] == "embedding": | |
| return ".".join(parts[:2]) | |
| return parts[0] |
🤖 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 `@dptb/utils/tools.py` around lines 168 - 172, The group_key function must
preserve the layer index for embedding layer names. Update its embedding branch
to return a key containing the first three components, such as
embedding.layers.0, while retaining the existing fallback behavior for other
names.
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
1 participant
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.
Scope
Briefly describe what this PR changes and why.
DeePTB Impact Area
Check every area that may be affected, even indirectly.
Risk And Compatibility
If any answer is "yes", explain the intended compatibility behavior.
Tests
List the tests you ran and the behavior they cover.
AI Assistance
Notes:
Merge Decision
For maintainers. Fill this before merging.
Summary by CodeRabbit
New Features
Bug Fixes
Documentation