From 60a77fd5b7ce91a8f928ac345447527b4bfdeb46 Mon Sep 17 00:00:00 2001 From: Eva Xue Date: Mon, 13 Jul 2026 13:51:00 -0700 Subject: [PATCH 1/4] fix(predict): accept --rdkit2D_normalization_type argument The predict command 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 value baked into the checkpoint. The flag was defined only on the finetune parser, so `main.py predict --rdkit2D_normalization_type ...` failed with "unrecognized arguments". Add it to the predict parser, mirroring the finetune definition (choices fast/best/descriptastorus, default fast). Signed-off-by: Eva Xue --- kermt/util/parsing.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/kermt/util/parsing.py b/kermt/util/parsing.py index 71e6886..3e5a67a 100644 --- a/kermt/util/parsing.py +++ b/kermt/util/parsing.py @@ -86,6 +86,9 @@ def add_predict_args(parser: ArgumentParser): parser.add_argument('--features_generator', type=str, nargs='*', choices=get_available_features_generators(), help='Method of generating additional features') + parser.add_argument('--rdkit2D_normalization_type', type=str, choices=("fast", "best", "descriptastorus"), default='fast', + help='Type of normalization for rdkit2D features. Choices: fast, best, descriptastorus. ' + 'Must match the value used at finetune time (it is baked into the features).') parser.add_argument('--features_path', type=str, nargs='*', help='Path to features to use in FNN (instead of features_generator)') parser.add_argument('--use_cuikmolmaker_featurization', action='store_true', default=False, From 4dbaef40f6fb3239c3614b20536eaafe064e65d1 Mon Sep 17 00:00:00 2001 From: Eva Xue Date: Mon, 13 Jul 2026 13:51:00 -0700 Subject: [PATCH 2/4] fix(predict): stop disabling bond dropout on the shared args (issue #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 other ranks kept the configured rate, desyncing augmentation across ranks. 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. Apply the same fix to fingerprint.do_generate, which had the identical pattern. Fixes #22 Signed-off-by: Eva Xue --- task/fingerprint.py | 4 ++++ task/predict.py | 8 ++++++++ 2 files changed, 12 insertions(+) diff --git a/task/fingerprint.py b/task/fingerprint.py index bfbc8cf..d66ed29 100644 --- a/task/fingerprint.py +++ b/task/fingerprint.py @@ -37,6 +37,7 @@ """ The fingerprint generation function. """ +import copy from argparse import Namespace from logging import Logger from typing import List @@ -63,6 +64,9 @@ def do_generate(model: nn.Module, :return: A list of fingerprints. """ model.eval() + # Disable bond dropout locally without mutating the shared args (same pattern + # as predict(); see issue #22). + args = copy.copy(args) args.bond_drop_rate = 0 preds = [] diff --git a/task/predict.py b/task/predict.py index caf71f4..1775743 100644 --- a/task/predict.py +++ b/task/predict.py @@ -37,6 +37,7 @@ """ The predict function using the finetuned model to make the prediction. . """ +import copy from argparse import Namespace from typing import List @@ -75,6 +76,13 @@ def predict(model: nn.Module, """ # debug = logger.debug if logger is not None else print model.eval() + # Evaluation must not apply bond dropout, but do NOT mutate the shared args: + # training and eval reuse one args instance and graphs are rebuilt every epoch + # (caching is off by default), so setting args.bond_drop_rate = 0 here would + # permanently disable bond dropout for all later training epochs -- and under + # DDP, where only rank 0 evaluates, desync augmentation across ranks. Use a + # shallow copy so the override is local to this call. See issue #22. + args = copy.copy(args) args.bond_drop_rate = 0 preds = [] From 87bc5a5f3d7d3869469e9cbc6a8254f0b1c17103 Mon Sep 17 00:00:00 2001 From: Eva Xue Date: Tue, 14 Jul 2026 10:46:52 -0700 Subject: [PATCH 3/4] fix(pretrain-ddp): set find_unused_parameters for non-both embedding modes The three pretrain trainers (KERMTTrainer, KERMTCMIMTrainer, KERMTHybridTrainer) wrapped the model in DDP without find_unused_parameters. In embedding_output_type=atom (or 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; DDP then aborts on the second iteration ("Expected to have finished reduction ... parameters that were not used in producing loss"). Enable find_unused_parameters only when embedding_output_type != 'both' (both exercises every branch, so keep it off to avoid the per-iteration graph-traversal overhead; static_graph is not viable because dynamic-depth sampling varies the graph across steps). Verified locally: hybrid + atom + DDP (WORLD_SIZE=1) crashed on iteration 2 before the change and trains a full epoch after it. Signed-off-by: Eva Xue --- task/kermttrainer.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/task/kermttrainer.py b/task/kermttrainer.py index 5fae8c2..cc9f023 100644 --- a/task/kermttrainer.py +++ b/task/kermttrainer.py @@ -215,7 +215,13 @@ def __init__(self, self.model.to(self.gpu_id) - self.model = DDP(self.model, device_ids=[gpu_id]) + # find_unused_parameters is required when embedding_output_type != 'both': + # 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. + # 'both' exercises every branch, so keep the flag off there to avoid DDP overhead. + # (static_graph is not an option: dynamic-depth sampling varies the graph per step.) + self.model = DDP(self.model, device_ids=[gpu_id], + find_unused_parameters=(self.args.embedding_output_type != 'both')) if self.args.tensorboard: self.writer = SummaryWriter(self.args.save_dir) @@ -564,7 +570,13 @@ def __init__(self, self.n_iter = 0 self.model.to(self.gpu_id) - self.model = DDP(self.model, device_ids=[gpu_id]) + # find_unused_parameters is required when embedding_output_type != 'both': + # 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. + # 'both' exercises every branch, so keep the flag off there to avoid DDP overhead. + # (static_graph is not an option: dynamic-depth sampling varies the graph per step.) + self.model = DDP(self.model, device_ids=[gpu_id], + find_unused_parameters=(self.args.embedding_output_type != 'both')) if self.args.tensorboard: self.writer = SummaryWriter(self.args.save_dir) @@ -871,7 +883,13 @@ def __init__(self, self.n_iter = 0 self.model.to(self.gpu_id) - self.model = DDP(self.model, device_ids=[gpu_id]) + # find_unused_parameters is required when embedding_output_type != 'both': + # 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. + # 'both' exercises every branch, so keep the flag off there to avoid DDP overhead. + # (static_graph is not an option: dynamic-depth sampling varies the graph per step.) + self.model = DDP(self.model, device_ids=[gpu_id], + find_unused_parameters=(self.args.embedding_output_type != 'both')) if self.args.tensorboard: self.writer = SummaryWriter(self.args.save_dir) From 08983595651eb3c0037e60b7b2ac49a94d8a9359 Mon Sep 17 00:00:00 2001 From: Eva Xue Date: Tue, 14 Jul 2026 11:00:11 -0700 Subject: [PATCH 4/4] fix(pretrain-vocab): guard float loss components in KERMTTrainer (atom/bond) In embedding_output_type=atom (or bond) mode the vocab loss returns a plain float 0.0 for the branch that has no embeddings (e.g. bond-vocab and bond dist loss in atom mode). KERMTTrainer.iter called .item() on the task and dist loss components unconditionally, raising "AttributeError: 'float' object has no attribute 'item'" once the affected branch was logged. Guard every loss-component .item() with `if not isinstance(x, float) else x`, matching the idiom KERMTHybridTrainer already uses (which is why hybrid mode was unaffected). Covers the train accumulation, eval branch, and the wandb logging dict. Verified: vocab + atom + DDP (WORLD_SIZE=1) trains a full epoch + validation. Signed-off-by: Eva Xue --- task/kermttrainer.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/task/kermttrainer.py b/task/kermttrainer.py index cc9f023..01b19ba 100644 --- a/task/kermttrainer.py +++ b/task/kermttrainer.py @@ -276,9 +276,9 @@ def validation(self, max_val_batches: int) -> List: loss_sum += loss.item() iter_count += self.args.batch_size - av_loss_sum += av_loss.item() - bv_loss_sum += bv_loss.item() - fg_loss_sum += fg_loss.item() + av_loss_sum += av_loss.item() if not isinstance(av_loss, float) else av_loss + bv_loss_sum += bv_loss.item() if not isinstance(bv_loss, float) else bv_loss + fg_loss_sum += fg_loss.item() if not isinstance(fg_loss, float) else fg_loss av_dist_loss_sum += av_dist_loss.item() if type(av_dist_loss) != float else av_dist_loss bv_dist_loss_sum += bv_dist_loss.item() if type(bv_dist_loss) != float else bv_dist_loss fg_dist_loss_sum += fg_dist_loss.item() if type(fg_dist_loss) != float else fg_dist_loss @@ -401,13 +401,13 @@ def iter(self, epoch, train=True) -> List: self.scheduler.step() else: # For eval model, only consider the loss of three task. - cum_loss_sum += av_loss.item() - cum_loss_sum += bv_loss.item() - cum_loss_sum += fg_loss.item() + cum_loss_sum += av_loss.item() if not isinstance(av_loss, float) else av_loss + cum_loss_sum += bv_loss.item() if not isinstance(bv_loss, float) else bv_loss + cum_loss_sum += fg_loss.item() if not isinstance(fg_loss, float) else fg_loss - av_loss_sum += av_loss.item() - bv_loss_sum += bv_loss.item() - fg_loss_sum += fg_loss.item() + av_loss_sum += av_loss.item() if not isinstance(av_loss, float) else av_loss + bv_loss_sum += bv_loss.item() if not isinstance(bv_loss, float) else bv_loss + fg_loss_sum += fg_loss.item() if not isinstance(fg_loss, float) else fg_loss av_dist_loss_sum += av_dist_loss.item() if type(av_dist_loss) != float else av_dist_loss bv_dist_loss_sum += bv_dist_loss.item() if type(bv_dist_loss) != float else bv_dist_loss fg_dist_loss_sum += fg_dist_loss.item() if type(fg_dist_loss) != float else fg_dist_loss @@ -438,12 +438,12 @@ def iter(self, epoch, train=True) -> List: if self.gpu_id == 0 and self.n_steps % train_log_interval == 0: train_metrics = { 'train/loss': loss.item(), - 'train/av_loss': av_loss.item(), - 'train/bv_loss': bv_loss.item(), - 'train/fg_loss': fg_loss.item(), - 'train/av_dist_loss': av_dist_loss.item(), - 'train/bv_dist_loss': bv_dist_loss.item(), - 'train/fg_dist_loss': fg_dist_loss.item(), + 'train/av_loss': av_loss.item() if not isinstance(av_loss, float) else av_loss, + 'train/bv_loss': bv_loss.item() if not isinstance(bv_loss, float) else bv_loss, + 'train/fg_loss': fg_loss.item() if not isinstance(fg_loss, float) else fg_loss, + 'train/av_dist_loss': av_dist_loss.item() if not isinstance(av_dist_loss, float) else av_dist_loss, + 'train/bv_dist_loss': bv_dist_loss.item() if not isinstance(bv_dist_loss, float) else bv_dist_loss, + 'train/fg_dist_loss': fg_dist_loss.item() if not isinstance(fg_dist_loss, float) else fg_dist_loss, 'train/lr': self.scheduler.get_lr()[0], 'train/epoch': epoch, 'train/batch_idx': ibatch + self.batch_idx_offset,