deepmodeling / deepmodeling/unimol_tools
[Code scan] Guarantee masked atoms or skip empty pretraining masked losses
- Dominant language
- Python
- Stars
- 34
- Forks
- 9
- PR merge metrics
- No merged PRs in 30d
Description
This issue is a result of a Codex global repository scan.
## Summary
Pretraining computes the number of masked atoms as int(sz * mask_prob + random). For small molecules this can be zero. The loss then computes mean token, coordinate, and distance losses over empty tensors, which can produce NaN and poison training.
## Code references
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/dataset.py#L267-L274
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/loss.py#L44-L68
https://github.com/deepmodeling/unimol_tools/blob/4596596aa8f73eb462d5cc5a921d79966d0465da/unimol_tools/pretrain/loss.py#L79-L99
## Impact
A batch containing molecules with zero selected masked atoms can generate NaN loss values. Once NaNs enter gradients or checkpoints, the pretraining run may become unrecoverable.
## Suggested fix
When mask_prob > 0, clamp num_mask to at least one and at most the atom count. Alternatively, explicitly skip masked objectives for samples or batches with zero masked atoms and avoid mean reductions over empty tensors.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the referenced masking logic in unimol_tools/pretrain/dataset.py around lines 267-274, then trace the masked reductions in unimol_tools/pretrain/loss.py around lines 44-68 and 79-99. Check the behavior for small molecules that select zero atoms and verify the chosen handling prevents NaN losses while preserving valid pretraining objectives.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100