`bond_drop_rate` is silently disabled after the first validation (and diverges across ranks under DDP)

[sveccham]

Summary

predict() mutates the shared args object by setting args.bond_drop_rate = 0 as a side effect of evaluation. Because training and evaluation share the same args instance and graphs are re-featurized every epoch (caching is off by default), this permanently disables bond dropout for the rest of the run after the first validation pass — regardless of what the user requested via --bond_drop_rate.

Root cause

Three things combine:

1. predict() unconditionally zeroes the value on the shared args:

   # task/predict.py:76-77
   model.eval()
   args.bond_drop_rate = 0
   

2. Caching is off by default (--no_cache defaults to True), so MolGraph is rebuilt every batch/epoch and re-reads args.bond_drop_rate at featurization time:

   # kermt/data/molgraph.py:472-474
   mol_graph = MolGraph(smiles, args=args)   # reads args.bond_drop_rate (molgraph.py:207-208)
   if not args.no_cache:                      # no_cache defaults True (parsing.py:53) -> never cached
       shared_dict[smiles] = mol_graph
   

3. The epoch loop runs train() then evaluate() (which calls predict()), so the mutation lands after the first epoch's training.

Impact

Single-process finetune (pre-existing, main.py finetune): for --bond_drop_rate 0.1:

Epoch	train() uses	after evaluate()
0	0.1	args set to 0.0
1..N	0.0	0.0

Only the first training epoch applies bond dropout; all later epochs train with it disabled. With --num-folds/--ensemble-size > 1, the mutation persists across the process, so only the first fold's first epoch ever uses the requested rate.

DDP finetune (PR #21): validation runs on rank 0 only, so only rank 0 calls predict() and zeroes the value. Non-main ranks never do, so from epoch 1 onward:

• rank 0 trains with bond_drop_rate = 0.0
• ranks 1..N-1 train with bond_drop_rate = 0.1

i.e. ranks apply inconsistent bond-dropout augmentation to their local batches.

Expected behavior

bond_drop_rate should stay at the user-requested value for the entire training run (or, if there's an intentional reason to disable it, that should be explicit and consistent across all epochs, folds, and ranks). Evaluation must not silently alter training hyperparameters.

Suggested fix

Stop mutating args in predict() — use a local value (e.g. featurize eval batches with a bond_drop_rate=0 passed explicitly / a shallow-copied args) so evaluation leaves training's bond_drop_rate intact. This fixes both the single-process "first-epoch-only" behavior and the DDP cross-rank inconsistency in one change.

Notes

Pre-existing bug in the single-process path; surfaced during review of #21, which turns the (consistent-but-wrong) behavior into a rank-inconsistent one.