fix(predict+pretrain): rdkit2D predict arg, issue #22, and atom/bond-mode DDP + float-loss fixes

[evasnow1992]

Summary

Two related correctness fixes in the prediction / evaluation path:

1. main.py predict rejects --rdkit2D_normalization_type — the predict command couldn't use rdkit-2D-normalized features.
2. Bond dropout is silently disabled after the first validation (and desyncs across DDP ranks) — fixes bond_drop_rate is silently disabled after the first validation (and diverges across ranks under DDP) #22.

Both are small, self-contained changes with no behavior change for existing correct runs.

Fix 1 — accept --rdkit2D_normalization_type on the predict parser

predict can build rdkit-2D-normalized features (rdkit_2d_normalized_cuik_molmaker), whose featurization reads args.rdkit2D_normalization_type (kermt/data/molgraph.py) and must match the normalization baked into the checkpoint. The flag was defined only on the finetune parser, so any predict run passing it failed with:

main.py: error: unrecognized arguments: --rdkit2D_normalization_type descriptastorus

Added the argument to add_predict_args, mirroring the finetune definition (choices fast/best/descriptastorus, default fast).

Fix 2 — don't disable bond dropout on the shared args (Fixes #22)

predict() set args.bond_drop_rate = 0 on the shared args instance. Training and evaluation reuse one args object, and graphs are rebuilt every epoch (caching off by default), so after the first validation pass bond dropout stayed disabled for every subsequent training epoch. Under DDP, only rank 0 evaluates, so rank 0 trained with bond_drop_rate=0 while the other ranks kept the configured rate, desyncing augmentation across ranks.

Fix: shallow-copy args before overriding bond_drop_rate, so the override is local to the evaluation call and never leaks back into training or other ranks. Applied the same fix to fingerprint.do_generate, which had the identical pattern (impact there is benign since it's a standalone command, fixed for consistency).

Changes

File	Change
kermt/util/parsing.py	Add --rdkit2D_normalization_type to the predict parser
task/predict.py	Shallow-copy args before disabling bond dropout (+ import copy)
task/fingerprint.py	Same shallow-copy fix in do_generate (+ import copy)

Testing

• tests/integration/test_pretrain_finetune.py::test_pretrain_ddp_finetune — now PASSES end-to-end (pretrain → finetune → predict). This test previously failed at the predict step on the unrecognized --rdkit2D_normalization_type argument.
• Verified the predict parser accepts --rdkit2D_normalization_type descriptastorus, and that both predict() and fingerprint.do_generate() no longer mutate the caller's args.

Fixes #22

[evasnow1992]

Additional fixes on this branch: atom/bond-mode pretraining

While validating the predict changes I found two pre-existing bugs that break multi-GPU pretraining with --embedding_output_type atom (or bond), and folded both fixes into this branch. Neither affects both-mode (the default).

1. DDP aborts: "parameters that were not used in producing loss" (commit 87bc5a5)

Symptom. Pretraining under DDP with --embedding_output_type atom aborts on the second iteration:

RuntimeError: Expected to have finished reduction in the prior iteration ...
parameters that were not used in producing loss.
Parameter indices which did not receive grad for rank N: 86 87 88 ...

Cause. In atom/bond mode the encoder produces only one aggregation level, so the opposite branch's FFNs and the bond/atom vocab + FG heads receive no gradient -- but all three trainers wrapped the model as DDP(self.model, device_ids=[gpu_id]) with the default find_unused_parameters=False. both mode is unaffected (every branch
participates in the loss).
Fix. find_unused_parameters=(self.args.embedding_output_type != 'both') in KERMTTrainer, KERMTCMIMTrainer, and KERMTHybridTrainer.
Why conditional, and not static_graph. Gating on != 'both' keeps the default path at zero overhead; and static_graph=True (the faster alternative for unused params) is not viable here because MTBlock's dynamic-depth sampling varies the graph shape every step.

2. AttributeError: 'float' object has no attribute 'item' (commit 0898359)

Symptom. Vocab-mode pretraining with --embedding_output_type atom crashes:

AttributeError: 'float' object has no attribute 'item'

Cause. In atom/bond mode the vocab loss returns a plain float 0.0 for the branch with no embeddings (e.g. bond-vocab + bond dist loss in atom mode), but KERMTTrainer.iter called .item() on the task/dist loss components unconditionally (train accumulation, eval branch, and wandb logging dict). KERMTHybridTrainer already guarded these, which is why hybrid mode was unaffected.
Fix. Guard every loss-component .item() with ... if not isinstance(x, float) else x, matching the hybrid trainer's idiom.

Verification (local, single-GPU DDP via WORLD_SIZE=1, kermt:latest container)

• hybrid + atom + DDP: crashed on iteration 2 before fix Docker container build error #1; trains a full epoch after.
• vocab + atom + DDP: crashed on the float .item() before fix make cuik-molmaker optional #2; trains a full
  epoch + validation after both fixes (bond-vocab loss logs as 0.0000, as expected).

[greptile-apps]

Greptile Summary

This PR fixes two correctness issues in the prediction and pre-training paths: a DDP augmentation desync caused by mutating the shared args object before evaluation (issue #22), and a crash in atom/bond-only DDP training caused by unused parameters in the opposite encoder branch.

• Bond-dropout desync (Fix make cuik-molmaker optional #2): predict() and do_generate() now shallow-copy args before zeroing bond_drop_rate, so training epochs and non-rank-0 DDP workers are no longer silently affected.
• DDP unused-parameter crash (kermttrainer.py): All three trainer classes now set find_unused_parameters=(embedding_output_type != 'both') on DDP, which is required when atom- or bond-only mode leaves the opposite branch's heads without gradients.
• Supporting fixes: load_checkpoint_for_prediction gains a guard for old checkpoints missing newer arg keys and fixes a swapped Finetune/Predict label in its error message; make_predictions gains an early features_size mismatch check; run_evaluation.py and fingerprint.py are updated for the (model, extras) tuple return shape.

Confidence Score: 5/5

• The changes are targeted and self-contained; the bond-dropout and DDP fixes address well-understood, reproducible bugs without altering any model architecture or training logic.
• All three trainer classes receive consistent DDP fixes, the shared-args mutation is cleanly isolated with a shallow copy, the consistency-check guard in load_checkpoint_for_prediction correctly handles old checkpoints, and the integration test now passes end-to-end. No new logic paths introduce regressions.
• No files require special attention; the one pre-existing isinstance vs type() inconsistency in kermttrainer.py for dist-loss variables is already tracked in a prior review comment.

Important Files Changed

Filename	Overview
kermt/util/parsing.py	Adds a code comment documenting why rdkit2D_normalization_type is deliberately absent from the predict parser; no functional change, correct approach.
kermt/util/utils.py	Improves load_checkpoint_for_prediction: adds a key not in vars(loaded_args) guard for old checkpoints, fixes a swapped Finetune/Predict label in the error message, and skips the features_generator consistency check when the caller supplies precomputed features via --features_path.
task/predict.py	Core fix for issue #22: copy.copy(args) before setting bond_drop_rate=0 prevents the shared args object from being mutated; also adds a features_size mismatch check and switches to load_checkpoint_for_prediction for strict loading.
task/fingerprint.py	Applies the same copy.copy(args) pattern as predict.py to do_generate; also fixes a call-site mismatch — load_checkpoint now returns a tuple, so the unpacking model, _ is added.
task/kermttrainer.py	Adds find_unused_parameters=(embedding_output_type != 'both') to all three DDP wraps; fixes crash when atom/bond-only mode leaves the other branch's heads without gradients. Also guards av/bv/fg loss accumulation and logging with isinstance(x, float) to tolerate float placeholders. Dist-loss accumulation still uses the older type(x) != float idiom (pre-existing inconsistency noted in an earlier review comment).
task/run_evaluation.py	One-line fix: load_checkpoint now returns (model, extras), so the call-site is updated to unpack the tuple with model, _.
tests/integration/test_pretrain_finetune.py	Removes --rdkit2D_normalization_type descriptastorus from the predict test call; the value is now correctly inherited from the checkpoint via make_predictions, which is the intended design.

Sequence Diagram

sequenceDiagram
    participant CLI as main.py predict
    participant MP as make_predictions()
    participant LC as load_checkpoint_for_prediction()
    participant P as predict()

    CLI->>MP: args (no rdkit2D_normalization_type)
    MP->>MP: load_args(ckpt[0]) → train_args
    MP->>MP: "copy missing keys from train_args → args<br/>(incl. rdkit2D_normalization_type)"
    MP->>MP: "args.features_size = test_data.features_size()"
    MP->>MP: check features_size vs ckpt_features_size
    loop each checkpoint_path
        MP->>LC: "current_args=args"
        LC->>LC: consistency check (rdkit2D_normalization_type, features_generator)
        LC->>LC: copy model-related args → current_args
        LC-->>MP: model
        MP->>P: args (original, shared)
        P->>P: "args_copy = copy.copy(args)"
        P->>P: "args_copy.bond_drop_rate = 0"
        P->>P: "MolCollator(args=args_copy)"
        P-->>MP: preds, loss_avg
    end
    MP-->>CLI: predictions

Loading

_{Reviews (4): Last reviewed commit: "fix tests/integration/test_pretrain_fine..." | Re-trigger Greptile}

[sveccham]

Thanks for the fixes and sorry for the delay.