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 bugs: the predict CLI now accepts --rdkit2D_normalization_type (previously unrecognized), and predict() / do_generate() no longer mutate the shared args instance when zeroing bond_drop_rate for evaluation — preventing silent dropout desync between DDP ranks and across training epochs. A third fix enables find_unused_parameters in the three DDP trainers when embedding_output_type != 'both', preventing DDP hangs in atom/bond-only pretrain mode.

• kermt/util/parsing.py: --rdkit2D_normalization_type added to add_predict_args with identical choices/default as the finetune parser, unblocking rdkit-2D-normalized predict runs.
• task/predict.py + task/fingerprint.py: copy.copy(args) before args.bond_drop_rate = 0 confines the override to the local call, fixing the DDP augmentation desync described in issue bond_drop_rate is silently disabled after the first validation (and diverges across ranks under DDP) #22.
• task/kermttrainer.py: Three trainers now conditionally set find_unused_parameters=True for atom/bond embedding modes, and main-loss accumulation guards are updated from bare .item() calls to float-safe equivalents.

Confidence Score: 4/5

All three fixes are targeted and correct; the copy.copy approach properly isolates the bond-dropout override, and the DDP parameter flag is gated on an argument that always has a default.

The changes are small and well-reasoned. The copy.copy(args) fix correctly stops args mutation, the --rdkit2D_normalization_type addition is a clean mirror of the finetune definition, and find_unused_parameters is safely conditional. A minor style inconsistency remains in the dist-loss accumulation guards across kermttrainer.py where the PR introduced isinstance in some places but left type() != float unchanged in others.

The dist-loss accumulation lines in task/kermttrainer.py (validation loop and iter loop) still use the old type() != float idiom while the newly touched metrics-dict lines in the same methods use isinstance.

Important Files Changed

Filename	Overview
kermt/util/parsing.py	Adds --rdkit2D_normalization_type to the predict argument parser, mirroring the finetune parser definition with identical choices and default. Straightforward fix with no issues.
task/predict.py	Introduces copy.copy(args) before setting bond_drop_rate = 0, preventing mutation of the shared args instance. The shallow copy is correct since bond_drop_rate is a scalar float.
task/fingerprint.py	Same shallow-copy fix as predict.py; benign impact for the standalone fingerprint command but consistent with the corrected pattern.
task/kermttrainer.py	Three trainers now set find_unused_parameters conditionally on embedding_output_type, and av_loss/bv_loss/fg_loss accumulation guards use isinstance(…, float). Minor inconsistency: dist-loss accumulation lines still use the old type() != float style while the newly added metrics-dict lines use isinstance.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Rank0 as DDP Rank 0
    participant RankN as DDP Rank N
    participant Predict as predict()
    participant Fingerprint as do_generate()

    Note over Rank0,RankN: Training epoch (before fix)
    Rank0->>Predict: evaluate(args)
    Predict->>Predict: "args.bond_drop_rate = 0 mutates shared args"
    Predict-->>Rank0: results
    Note over Rank0: Next epoch: bond_drop_rate=0 (wrong)
    Note over RankN: Next epoch: bond_drop_rate=configured (desync)

    Note over Rank0,RankN: Training epoch (after fix)
    Rank0->>Predict: evaluate(args)
    Predict->>Predict: "local_args = copy.copy(args)"
    Predict->>Predict: "local_args.bond_drop_rate = 0 local only"
    Predict-->>Rank0: results
    Note over Rank0: Next epoch: bond_drop_rate=configured (correct)
    Note over RankN: Next epoch: bond_drop_rate=configured (in sync)

    Note over Rank0: Fingerprint path (same fix)
    Rank0->>Fingerprint: do_generate(model, args, ...)
    Fingerprint->>Fingerprint: "local_args = copy.copy(args)"
    Fingerprint->>Fingerprint: "local_args.bond_drop_rate = 0 local only"

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Rank0 as DDP Rank 0
    participant RankN as DDP Rank N
    participant Predict as predict()
    participant Fingerprint as do_generate()

    Note over Rank0,RankN: Training epoch (before fix)
    Rank0->>Predict: evaluate(args)
    Predict->>Predict: "args.bond_drop_rate = 0 mutates shared args"
    Predict-->>Rank0: results
    Note over Rank0: Next epoch: bond_drop_rate=0 (wrong)
    Note over RankN: Next epoch: bond_drop_rate=configured (desync)

    Note over Rank0,RankN: Training epoch (after fix)
    Rank0->>Predict: evaluate(args)
    Predict->>Predict: "local_args = copy.copy(args)"
    Predict->>Predict: "local_args.bond_drop_rate = 0 local only"
    Predict-->>Rank0: results
    Note over Rank0: Next epoch: bond_drop_rate=configured (correct)
    Note over RankN: Next epoch: bond_drop_rate=configured (in sync)

    Note over Rank0: Fingerprint path (same fix)
    Rank0->>Fingerprint: do_generate(model, args, ...)
    Fingerprint->>Fingerprint: "local_args = copy.copy(args)"
    Fingerprint->>Fingerprint: "local_args.bond_drop_rate = 0 local only"

Loading

Comments Outside Diff (1)

1. task/kermttrainer.py, line 282-286 (link)

   The dist-loss accumulation lines still use type(x) != float while the metrics-dict entries added in the same PR use isinstance(x, float) for these same variables. Within a single method this split style is confusing — isinstance is the idiomatic Python check. Suggest aligning the accumulation lines to match.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "fix(pretrain-vocab): guard float loss co..." | Re-trigger Greptile}