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 in the prediction/pretrain path: bond dropout being permanently disabled (and desynced across DDP ranks) after the first validation pass, and a KeyError crash in the checkpoint consistency check when a predict arg key was absent from an older checkpoint. Several pre-existing bugs are also cleaned up along the way.

• Bond dropout / DDP desync (issue bond_drop_rate is silently disabled after the first validation (and diverges across ranks under DDP) #22): predict() and fingerprint.do_generate() now shallow-copy args before zeroing bond_drop_rate, so the override is scoped to the evaluation call and never leaks back into training or across DDP ranks. All three trainer classes also gain find_unused_parameters=(embedding_output_type != 'both') so atom/bond-only DDP runs no longer hang on unused parameter detection.
• Consistency check and return-value fixes: load_checkpoint_for_prediction gains a guard for keys absent from older checkpoints and a features_generator bypass when features_path is supplied; fingerprint.generate_fingerprints and run_evaluation.run_evaluation fix a pre-existing bug where load_checkpoint's tuple return was assigned directly to model instead of being unpacked.

Confidence Score: 5/5

• This PR is safe to merge; all changes are targeted bug fixes with clear rationale and no behavior change for correctly-configured runs.
• Every change either restores correct behavior (bond dropout scoping, DDP unused-parameter flag, tuple unpacking) or adds a defensive guard (missing checkpoint key, features_size mismatch, features_generator bypass). The shallow-copy approach is the minimal correct fix for the shared-args mutation, and the DDP flag is conditioned precisely on the modes that need it.
• No files require special attention; the changes are self-contained and well-commented.

Important Files Changed

Filename	Overview
kermt/util/parsing.py	Adds a comment explaining that rdkit2D_normalization_type is intentionally absent from the predict parser so the checkpoint's baked-in value flows through via make_predictions() without being shadowed by an argparse default.
kermt/util/utils.py	Fixes two real bugs in load_checkpoint_for_prediction: a KeyError when a finetune_predict_consistency_args key is absent from an older checkpoint, and a swapped Finetune/Predict label in the error message. Also adds a features_generator bypass when features_path is supplied, preventing the consistency check from incorrectly rejecting the normal precomputed-features workflow.
task/predict.py	Fixes the DDP bond-dropout desync (issue #22) with a shallow copy of args before setting bond_drop_rate=0; switches from load_checkpoint to load_checkpoint_for_prediction; adds an early features_size mismatch check with a human-readable hint before the loader would emit a cryptic tensor-shape error.
task/fingerprint.py	Applies the same shallow-copy fix for bond_drop_rate as predict.py; also fixes a pre-existing bug where load_checkpoint's tuple return was assigned directly to model instead of being properly unpacked.
task/kermttrainer.py	Adds find_unused_parameters=(embedding_output_type != 'both') to DDP wrapping in all three trainer classes, fixing a DDP hang/error in atom- or bond-only embedding mode where the opposite branch's parameters receive no gradient. Also upgrades av/bv/fg loss accumulation to use isinstance() guards, fixing crashes when those losses are returned as Python floats rather than tensors.
task/run_evaluation.py	Fixes a pre-existing bug where load_checkpoint's tuple return was being assigned to model directly instead of unpacked; the old code would have crashed immediately when the tuple was used as a nn.Module.

Sequence Diagram

sequenceDiagram
    participant Train as Training Loop (rank 0..N)
    participant Eval as evaluate() / predict()
    participant DDP as DDP Model
    participant Args as shared args

    Note over Train,Args: Before fix: args.bond_drop_rate mutated globally
    Train->>DDP: "forward pass (bond_drop_rate=configured)"
    Train->>Eval: call evaluate()
    Eval->>Args: "args.bond_drop_rate = 0  ❌ mutates shared object"
    Eval->>DDP: "predict with bond_drop_rate=0"
    Train->>DDP: "next epoch forward (bond_drop_rate=0 ❌ leaked)"

    Note over Train,Args: After fix: args shallow-copied locally
    Train->>DDP: "forward pass (bond_drop_rate=configured ✅)"
    Train->>Eval: call evaluate()
    Eval->>Eval: "args_copy = copy.copy(args)"
    Eval->>Eval: "args_copy.bond_drop_rate = 0"
    Eval->>DDP: "predict with args_copy (bond_drop_rate=0 ✅)"
    Train->>DDP: next epoch forward (original bond_drop_rate preserved ✅)

Loading

_{Reviews (3): Last reviewed commit: "fix(fingerprint+eval): unpack the (model..." | Re-trigger Greptile}

[sveccham]

Thanks for the fixes and sorry for the delay.

[sveccham]

This this line still needed for DDP finetuning?

[evasnow1992]

I think so. DDP itself doesn't move the model to the target device; it only assumes the model is already on the specified device_id and operates from there.

[sveccham]

I think rdkit2D_normalization_type in predict should be the same as the checkpoint. Doesn't this give the user an option to use the finetuned checkpoint incorrectly?

[evasnow1992]

I dug into this a bit further and realized this is actually a bigger issue.

A couple of months ago, you checked in a commit that addressed this by introducing a stricter model loading function for prediction, with consistency checks on both rdkit2D_normalization_type and feature_generator. That change also switched main.py in prediction mode to use the new model loader.

Later, when I checked in the larger batch of changes for Contrastive KERMT, I switched back to the default model loader because of compatibility concerns, which inadvertently removed these consistency checks.

To resolve the rdkit2D_normalization_type issue, I'll revert my change here and switch prediction back to using the stricter model loader.

While investigating this, I also found a related issue with feature_generator. It should remain consistent with the fine-tuned checkpoint, but it also shouldn't be inherited blindly, since doing so can cause the feature generation step to be skipped. I'll add another commit to address this properly as well.