P3-06: the SLURM LoRA fine-tune of the RiNALMo-giga Stage-2 re-ranker (production + §11 sweep) - #99
Merged
Merged
Conversation
…oke gate (P3-06)
The Stage-2 RiNALMo-giga LoRA fine-tune the P1-15 harness config was waiting for.
Nothing is trained here: the §9.3 submit-ack is open, so no checkpoint, no dvc add.
Inputs: data/processed/stage2_dataset.parquet (P3-01); stage2/{model,losses,heads,
tokenizer}.py (P3-03/04/05); train/lora_harness.py (P1-15 §10.3 contract);
reports/p1/lora_vram_smoke.json (P1-16 sizing).
Outputs: src/tbox_finder/stage2/train.py; src/tbox_finder/train/ddp.py;
conf/train/stage2.yaml; conf/optim/stage2.yaml;
slurm/p3/stage2_lora_finetune.sbatch; tests/ml/test_stage2_train_smoke.py.
The §7 stop imp.md hands this step — the train-eligibility policy — is settled by
user sign-off (AskUserQuestion 2026-08-03) and encoded in row_eligibility():
admit iff fold_random == train AND NOT calib AND (nested_train is True OR
fold_basis == decoy_pool_random). D5 because PRD §12 grades leave-one-order-out on
the TWO-STAGE system; the scheme-A intersection so P3-07 fits its temperature on the
disjoint calib split and P3-10 grades GATE-2 on THIS checkpoint rather than an
ADR-0004 A6 eval twin. Measured: 9,059 train rows (5,199 corpus + 3,860 parentless
decoy; 4,795 pos / 4,264 neg), 2,070 val, row-id overlap 0, partition exact.
The sweep is a SLURM --array, not a literal Hydra --multirun (the P2-06 reasoning:
one process per point, so an OOM in point k cannot poison point k+1). Six points =
loss.aux_weight {0.0,0.5,1.0} x optim.lr {1e-4,3e-4}; aux_weight=0 is P3-08's no-aux
arm. LoRA rank is NOT swept — §10.3 pins r=16 as a drift-guarded constant.
Two hand-offs resolved: boundary_use_crf is False and True is REFUSED at construction;
pooling stays masked-mean. DDP uses find_unused_parameters=True (a pure-decoy batch
leaves the record-level heads without a gradient) with use_reentrant=False
checkpointing, and the module count carrying the flag is measured, not assumed.
train/ddp.py promotes the rank helpers, the PYTHONHASHSEED verifier and ShardedSampler
out of train_stage1.py, which now imports them: the two trainers run in different
conda envs (ADR-0002 A4), so a fork would ship the bug in one while fixing the other.
Validation: 80 pass torch-free (tbox-finder-data, hydra + real-corpus tiers armed),
87 pass under tbox-ml-rna (torch tier, real LoRA wrap + backward); test_sbatch_overrides
composes all ten override tokens (22 pass); test_sbatch_rm_targets clean; ruff 0.15.15
+ black 25.11.0 clean; bash -n clean. 10 source sabotages, each RED against its NAMED
test, each restored byte-identically. A CPU end-to-end run with max_records=32 failed
its own gate on exactly one clause, full_population — the completeness clause.
Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
… gate for every inline python -c block
Caught by the §9.3 pre-ack check that extracts the command's OWN bytes from the shipped
sbatch and executes them on the cluster, rather than running a retyped equivalent.
The block wrote an escaped quote inside an f-string expression:
f"attn={getattr(model.config, \"_attn_implementation\", None)}"
The body is a single-quoted shell string, so the backslash survives into the Python
source and the block does not parse. `bash -n` passes it — the shell string is
well-formed — so the defect was invisible to every local check. The job would have
died seconds in, after the queue wait; that is job 789's failure shape.
Fix: bind `attn` to a name before the f-string, plus a header note saying why no
backslash-escaped quote may appear in that block.
Gate: tests/ml/test_stage2_train_smoke.py::
test_every_inline_python_block_in_every_sbatch_actually_compiles compiles the body of
every `PYTHONPATH=... python -c '...'` block in every shipped sbatch (7 found, all
parse) and carries an emptiness guard, so a regex matching nothing cannot pass it
vacuously.
Validation: bash -n clean; ruff 0.15.15 + black 25.11.0 clean; 81 pass / 0 fail
torch-free with the hydra + real-corpus tiers armed; test_sbatch_overrides 22 pass.
Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…nd, not the run's Running the shipped bytes on the cluster surfaced `attn=eager`, which is correct for what the probe loads and WRONG as a description of the run. `load_rinalmo_backbone()` called bare passes attn_implementation=None to from_pretrained, so transformers picks its default; the training run goes through build_peft_model, which SELECTS the backend from the measured sm_86 evidence in reports/p1/kernel_smoke.json. Verified on the cluster: flash_attn_importable=True, is_sm86=True, model_supports_flash_attn=True -> the run selects flash_attention_2. Left as it was, every job log would have advertised eager attention for a run using FA-2 — the "a report that names a backend the model never used" defect P1-15 wrote its own read-back guard against. The probe now prints both, labelled: the warm-load's value marked NOT the run backend, and the selection the run will actually make, with its reason. Validation: bash -n clean; the inline-python compile gate still passes (81 pass / 0 fail torch-free with the hydra + real-corpus tiers armed). Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…kground, and test real rows Job 1036 died on all six points. This fixes the one cause that is mine, and closes the test gap that let it reach the cluster. All 7,007 decoy rows carry label_string = None (all four pools; 0 of 23,535 corpus rows), because a decoy has no T-box element annotation. Stage2SequenceDataset.__getitem__ applied its alignment guard unconditionally and raised on the first decoy. Decoys are now supervised as ALL-BACKGROUND, not ignored — the Stage-1 convention for the same shared SegmentationHead (data/negatives.py: "every real nucleotide labelled background, with no IGNORE_INDEX ... That is not a trick"), so head (b) gets real anti-mimicry signal from 7,007 rows in the dense-background regime PRD §11's focal loss targets. User decision 2026-08-03. The index is DERIVED via labels.CLASS_CODE -> CODE_TO_CLASS -> the spec's own vocabulary, not a literal 0. The branch is fail-closed: source == decoy AND is_tbox is False, so an unlabelled corpus row still raises. THE TEST GAP, which is the real defect: _decoy() inherited a 16-char label_string from _row() that no real decoy has, and the real-corpus tier only COUNTED eligibility — nothing ever built a dataset ITEM from a real row. Closed with: one real row of every pool through __getitem__; a walk of the entire 9,059-row admitted population; and a fail-closed check that an unlabelled POSITIVE still raises. Both new branches sabotage-bitten. Also, from cause (B): the sbatch now fails fast on an unhealthy node. Node two returned "Failed to initialize NVML: Driver/library version mismatch (595.84)" and killed three points inside init_process_group. nvidia-smi + a real torch.cuda allocation now run BEFORE the 2.5 GB download, turning a 2-min opaque NCCL death into a 5-second STAGE2_NODE_UNHEALTHY exit. It asks the node whether it works rather than assuming which node is good, so §13's never-pin rule is untouched. Array serialised (%1) while one of two nodes is unusable. Validation: 92 pass / 0 fail with every tier armed; test_sbatch_overrides 22 pass; ruff + black clean; bash -n clean; 12 sabotages each RED against its NAMED test. The CPU end-to-end run now completes 16/16 steps over decoy-containing batches, all six terms contributing, failing its gate on only full_population. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…ord the root cause Node `two` traced from /var/log/dpkg.log and /proc/driver/nvidia/version: it runs kernel 6.8.0-110-generic (booted 2026-07-14) with a RESIDENT nvidia module of 595.71.05, while userspace is libnvidia-ml.so.595.84. The 595 series was installed 2026-07-14 14:52 and upgraded 595.71.05 -> 595.84 on 2026-07-31 06:26; DKMS built 595.84 only for kernels -134/-136, not for the running -110 (whose nvidia.ko.zst no longer exists), and removed the old build while the loaded module stayed resident. The node has been GPU-dead since 2026-07-31 06:27 — three days before job 1036. A reboot fixes it (GRUB_DEFAULT=0 selects -136, which has the matching module); a module reload cannot. This also REFUTES the tidy explanation for cause (C): the upgrade predates the run by three days, so "a driver upgrade rolled mid-run" is not merely unsupported, it is contradicted. Point 0's SIGSEGV stays unattributed. --exclude=two is added, dated 2026-08-03 and scoped to this fault, with a test that refuses the flag unless the sbatch also carries its removal condition — so the carve-out cannot outlive the reboot and silently halve the cluster. This is --exclude, not the --nodelist §13 forbids: §13 bans pinning TO a preferred node because `one` is often down; this excludes a node MEASURED broken and leaves `one` an ordinary scheduler choice. User-approved. Validation: 92 pass / 0 fail all tiers; test_sbatch_overrides 22 pass; test_sbatch_rm_targets clean for this file (the 8 failures there are the pre-existing _KNOWN_DEFECTIVE xfails); ruff + black clean; bash -n clean. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…1044 OOM'd Job 1044: all 6 points CUDA-OOM in the backbone forward, 13.97 GiB allocated of 15.60, ~16 min in, 0 checkpoints. The three job-1036 causes are gone — (A) fixed, every point reached training over decoy batches; (B) neutralised, all six placed on `one`; (C) did NOT recur, no SIGSEGV anywhere including the aux_weight=0.0 arm. ROOT CAUSE OF THE MIS-SIZING, and it is methodological: both submits were sized from reports/p1/lora_vram_smoke.json (P1-16, batch 8 -> 2.014 GiB). That report's own config says loss_is_placeholder: true, data_is_synthetic: true, is_science: false, seq_len_nt: 350, 5 steps — a bare backbone under a placeholder loss, no heads, no objective. It was extrapolated to a seven-head objective on 552-token sequences for 1,420 steps. The wall-clock was labelled an extrapolation; the VRAM headroom was never questioned. imp.md's Compute line had asked for the 100-record smoke that was skipped. The correction, in two parts: 1. train.forward_backward is EXTRACTED from the training loop, so the sizing harness measures the step the trainer runs rather than a copy of it. A harness that reimplements the step can drift from it exactly as P1-16 did; `real_objective` is a gate clause reading step_function == the trainer's own, beside data_is_synthetic/loss_is_placeholder false, and a test asserts by INTERCEPTION that measure_batch really calls it. 2. stage2/sizing.py + slurm/p3/stage2_sizing_smoke.sbatch: 1 GPU, no DDP, real rows from the real admitted population. Descending batch sweep (8,4,2,1) with OOM caught per point; WORST-CASE and typical batches (59-552 tokens, collator pads to batch max — sizing on the median is how this failed); SIX optimiser steps because AdamW allocates its state on the first .step() and a growing footprint only shows as a series (growth_ratio); gradient-checkpointing effectiveness as an on/off comparison. The gate grades the MEASUREMENT, not the answer: "batch 8 does not fit" is a successful run and exits 0. A first draft of `swept` demanded a point that ran steps, which failed an all-OOM sweep — a gate pressuring the harness toward a convenient number, the very failure being corrected. It now requires a DEFINITE outcome per point (ran, or OOM'd). Validation: test_stage2_sizing 17 pass / 0 fail (torch + real-corpus tiers armed); the trainer suite still 92 pass after the extraction; ruff + black clean; bash -n clean; 4 further sabotages each RED against its named test (16 across the step). Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
… a checkpointing no-op Job 1051 GREEN. Measured on one A4000 over 200 real admitted rows through the trainer's own forward_backward: batch 8 worst_case (551 tok) OOM batch 8 typical (417 tok) 11.6641 GiB 857 ms batch 4 worst_case (551 tok) 8.1827 GiB 596 ms <- chosen batch 2 worst_case 4.8775 GiB 335 ms batch 1 worst_case 3.1640 GiB 202 ms Job 1044 is fully explained: batch 8 TYPICAL fits, batch 8 WORST CASE does not — a run that trains ~16 min then dies on the first long-sequence batch. growth_ratio 1.01-1.04 across six steps, i.e. FLAT: a per-step footprint, not an accumulation, so batch size is the whole fix. THE FINDING THE GATES COULD NOT SEE: gradient checkpointing is a NO-OP on this backbone. 3.1640 GiB on vs 3.1595 GiB off — saving ratio 0.9986 — across 36 modules carrying the flag. The source confirms it: modeling_rinalmo.py advertises supports_gradient_checkpointing = True (:82) and stores self.gradient_checkpointing (:118), but the encoder loop (:704) never calls _gradient_checkpointing_func. P1-16's report records bf16_and_gradient_checkpointing as a PASSING gate clause; this step's gradient_checkpointing_verified counts flagged modules and build_model raises when the count is zero. All three are satisfied by an attribute nothing reads — flag-counting where only effect-measuring would do. The sizing harness caught it by measuring a difference instead of reading a setting. Response: gradient_checkpointing ships FALSE with the measurement and source lines recorded, plus a drift-guard test asserting the port still lacks the hook, so a future multimolecule that adds it makes the test fail rather than the project silently foregoing the optimisation. Hand-wiring torch.utils.checkpoint is deferred, not assumed. Re-sized: batch_size=4, gradient_accumulation_steps=2 (effective batch unchanged at 8), eval_batch_size=8. ~283 steps/epoch/rank at ~500 ms => ~25-30 min/point against a 2 h limit. The stale P1-16 rationale is DELETED from conf/train/stage2.yaml, not left beside the new one. Validation: test_stage2_sizing 18 pass, trainer suite 92 pass, test_sbatch_overrides 22 pass; ruff + black clean; bash -n clean. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
Job 1053 point 0 TRAINED SUCCESSFULLY — 10 epochs, best val 0.001677 at epoch 4, rank 0's report overall_pass: true, 51.4 MB adapter, 13,040,268 saved params — and wrote no .OK, because ranks 1-7 raised on `checkpoint_written, provenance_recorded`. Only the primary rank saves the adapter and takes the git snapshot, but every rank built the report AND judged it. On ranks 1-7 those two clauses were correctly False, and raising on them turned a healthy run into a non-zero exit: torchrun propagated it, the sbatch's rc check skipped the cp out of node-local scratch, and the EXIT trap deleted the checkpoint. ~45 GPU-minutes destroyed at the finish line. Third variant of one theme in this step, worth naming: a clause is only as good as the SCOPE it is evaluated in. P2-10d'-c read state the run mutated; the checkpointing clause counted a flag instead of measuring an effect; this grades artifacts on processes never meant to produce them. Each passed its unit tests, because each was tested where its evidence existed. Fix: non-primary ranks build the report and RETURN WITHOUT JUDGING. Every rank still raises on its own OOM/NaN/collate failure — those propagate out of the loop and never reach the gate, and that path caught jobs 1036 and 1044. No collective, so no barrier to get wrong. Guarded by a clause-level test (a non-primary report derives exactly those two False while every OTHER clause still holds) and an ast/source check that the early return precedes validate_report. Sabotage-bitten by deleting the return. The checkpoint-copy behaviour is left as is, deliberately: a run that fails its gate must not leave a checkpoint. The bug was the gate, not the copy. Also: node two is FIXED (verified — kernel 6.8.0-134-generic, resident module 595.84 == userspace 595.84, nvidia-smi lists 8 A4000s). The cluster is now heterogeneous (one on 590.48.01), so --exclude=two is RETAINED for a different, weaker reason — holding the driver constant across six points meant to be compared. Its rationale is REWRITTEN, not left: the fault it cited no longer exists. The guard test now requires a dated rationale plus a removal condition rather than one specific sentence. Validation: trainer 94 pass, sizing 18 pass, overrides 22 pass; ruff + black clean; bash -n clean; 17 sabotages across the step. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…ually used
User decision 2026-08-03: --exclude=two and --array=0-5%1 both removed. The node fault they
guarded is fixed, and the %1 comment had already committed to its own removal ("Restore 0-5
once both nodes are healthy"), so this executes a recorded decision. Trigger: cryosparc holds
2 of 8 GPUs on EACH node, so a whole-node request cannot start on either until those clear
(the P2-06 starvation pattern, ~4 h), and excluding a node halved the chance of one freeing.
sbatch --test-only reporting a 2027 start is SLURM's infinite-TIMELIMIT placeholder.
ACCEPTED CAVEAT: the nodes now run different drivers — one on 590.48.01, two on 595.84 —
so points may land on either. Chosen over an open-ended wait. Expected immaterial (same CUDA
12.x, same torch build, same pinned wheels), but expected is not controlled, and six arms
that exist to be compared should not be uniform on an axis only by assumption.
So the axis is RECORDED: train.device_record() writes GPU name, capability, visible-device
count, torch version, CUDA runtime AND the kernel driver version into every report. The
driver comes from nvidia-smi, not torch — torch.version.cuda is the CUDA RUNTIME, which is
not what diverged; reporting it would have looked like provenance while recording the wrong
number.
It is provenance, NOT a verdict: a run on either driver is legitimate, so the block is not a
gate clause, and a test asserts that stripping it leaves the derived clauses unchanged — a
device block leaking into the gate would fail a good run for landing on the other node. A
second test asserts device_record() survives a missing nvidia-smi and records None.
Validation: trainer 96 pass, sizing 18 pass, overrides 22 pass; ruff + black clean; bash -n
clean. The sbatch guard now requires any future carve-out to carry a date and a removal
condition, and requires the accepted caveat to stay written down.
Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…ctory
Job 1064's first two points trained, passed their gate, and copied 48 MB checkpoints to the
repo path — then crashed writing provenance.json with IsADirectoryError. The sbatch declared
`{ckpt_dir}/lora_adapter` as a provenance output; that is PEFT's save_pretrained layout, a
DIRECTORY, and provenance.sha256_file opens each declared output as a file.
sha256_file is NOT weakened: raising on a non-file is correct fail-loud behaviour for a
shared helper used across the repo. The defect was the declaration.
Fix: train.checkpoint_output_files(ckpt_dir) enumerates every FILE under the checkpoint,
recursively and sorted, and refuses both silent-nothing cases (not a directory; a directory
with no files). Shared by the sbatch and the pending backfill so the two cannot disagree
about what a checkpoint's outputs are. Enumerated rather than hard-coded so a PEFT version
writing a different file set is hashed in full instead of silently under-recorded.
Notably this failure was CHEAP where the previous one was not: the copy out of node-local
scratch precedes the sidecar write, so the checkpoints survived — the same bug class as job
1053's, at ~0 GPU-h instead of ~45 GPU-min, because that ordering was fixed deliberately.
Tests: the enumerator yields only files, sorted, never the adapter dir; every declared output
is hashable by the real sha256_file; the directory itself still raises IsADirectoryError so
the contract is asserted rather than assumed; both refusal paths covered. Sabotage-bitten by
dropping the is_file() filter. The inline-python compile gate now also covers heredoc blocks
(`python - <<'PY'`), which is where the provenance writer lives — an unparsed heredoc is as
dead as an unparsed -c block, though note compiling would NOT have caught this one.
Validation: trainer 97 pass, overrides 22 pass; ruff + black clean; bash -n clean.
Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…ld not write Every point of job 1064 trains, passes its gate and copies its checkpoint, then dies writing provenance.json. The checkpoints are intact; only the sidecars are missing. This rebuilds them from each run's OWN report — the artifact the producing node wrote, carrying git_sha, env_lock_sha256, seed, config, population census and the device record. Nothing is invented: every field is either copied from that report or computed by hashing the artifact. It is WEAKER than a native sidecar and says so in its own output: `reconstructed: true`, the reason, the report it came from, and an explicit limitation field. A reader must be able to tell a backfilled record from a native one without reading the script. Two guards found by testing the write path rather than assuming it: 1. build_provenance derives the TOP-LEVEL git_sha from the working tree it runs in and takes no override — so a backfill from a moved checkout would silently attribute the run to the wrong commit, in the very field CLAUDE.md §11 names. Now VERIFIED: if the tree is not on the run's own commit the sidecar is deleted and the point refused, making "back-fill before re-syncing" an enforced precondition rather than something to remember. Confirmed by watching it refuse and unlink. 2. A report whose gate did not pass is refused outright — a sidecar must never describe a failed run. Confirmed by flipping overall_pass in a real report. Also: the tally counts actual writes, so a write-time refusal cannot be reported as a write (it said "wrote 1" while writing none). Exercised end-to-end against a REAL cluster report: 2 inputs and 3 outputs hashed, all 64-hex, git_sha 5a215d7 carried from the run, driver 595.84, best_val 0.00037977. The write path surfaced two missing-input failures in my mock before it worked — provenance hashes inputs too, which is why the backfill must run where the dataset and vocab live. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…ssing gates Artifact-based verification (sacct off): array gone from squeue, six reports, six 48 MB checkpoints (4 files each, stage2_heads.pt present), zero real failures. The only .err error was 6x IsADirectoryError, uniformly the known sidecar defect. point aux_w lr best_val epoch driver aux0.0_lr1e-4 0.0 1e-4 0.000380 6 595.84 aux0.0_lr3e-4 0.0 3e-4 0.000599 8 590.48.01 aux0.5_lr1e-4 0.5 1e-4 0.023589 8 590.48.01 These totals are NOT comparable across arms: at aux_weight=0 the objective IS the binary term alone, at 0.5 it is that plus five weighted aux terms. Arm comparison is P3-08's job, on the calibrated binary head with matched metrics; the reports carry the per-term breakdown it needs. Driver split 3/3 across the two nodes, recorded per point in each report's device block — the accepted caveat made auditable rather than assumed. Provenance: all six sidecars were lost to the IsADirectoryError and are RECONSTRUCTED from each run's own report, each carrying reconstructed: true, the reason, the source report and an explicit limitation. Two guards earned their place: the git_sha check (deletes the sidecar and refuses on a moved checkout), and the constraint it exposed — the script had imported a helper added AFTER the run's commit, while the guard requires running AT it. A repair script cannot depend on code newer than the tree it repairs, so the enumeration is now defined in the script with tests/unit/test_backfill_provenance.py pinning it to the shared helper, both refusal paths included, so the deliberate fork cannot drift. Retrieval verified: all 24 files match their recorded sha256. dvc add -> stage2_rinalmo.dvc (md5 b3d1c09c…dir, 313,437,285 B, 30 files) into the SHARED main-repo cache via .dvc/config.local, since a worktree's own .dvc/ is empty and the object would land where phase-exit dvc push never looks. Still owed: the six .OK markers were never written (the sidecar crash preceded them), so the §9.3 verification stands on reports + checkpoints + .err scan — a weaker signal, stated as such. The corrected sbatch (cfb8971) writes both natively. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…ked-path deleters CI went red on the RESULT commit. tests/unit/test_sbatch_rm_targets.py compares every `rm` target against `git ls-files`, and both p3 sbatch files did `rm -f "$REPORT"`. That was safe while reports/p3/**.json were UNTRACKED — which is why it passed locally and in CI up to 0f968fd. The RESULT commit committed all six sweep reports plus stage2_sizing.json, which retroactively made those `rm`s delete tracked paths. The test's verdict changed because the repo state changed, not because the test is flaky, and the rule it enforces is real: deleting a tracked path dirties the tree outside _DATA_STAGING_PREFIXES, which in the Stage-1 design makes _provenance_complete derive FALSE and a run fail its own gate AFTER training. Fix is the P2-09 pattern the failure message names: do not rm a tracked report, fingerprint it. Both files now `rm -f` only their untracked marker ($OK / $DONE) and take an md5 before the run, then assert after it that the report CHANGED. Both failure directions are explicit — an uncomputable fingerprint must not pass by being unequal to the old one, and an unchanged one must not pass at all — so "a report exists" can no longer be mistaken for "this run wrote it". Validation: test_sbatch_rm_targets clean for both p3 files (the 8 remaining failures there are the pre-existing _KNOWN_DEFECTIVE xfails); trainer 97, sizing 18, overrides 22, backfill 3; ruff + black clean; bash -n clean on both. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 40 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (8)
📝 WalkthroughWalkthroughAdded a complete Stage-2 RiNALMo LoRA training pipeline with Hydra and Slurm workflows, distributed utilities, GPU sizing, validated reports, checkpoint provenance, six sweep results, and smoke-test coverage. ChangesStage-2 training pipeline
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant SlurmJob
participant Stage2Training
participant Dataset
participant Checkpoint
participant Provenance
SlurmJob->>Stage2Training: Launch Hydra Stage-2 run
Stage2Training->>Dataset: Select and load admitted rows
Dataset-->>Stage2Training: Provide training and validation batches
Stage2Training->>Checkpoint: Save LoRA and Stage-2 head artifacts
Stage2Training->>Provenance: Write validated report and provenance
Provenance-->>SlurmJob: Confirm artifacts before success marker
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
Owner
Author
|
Please review this PR. Context for what to prioritise: This step landed a real GPU training run after four submits, three of which failed on my own bugs. The corrections are in the diff; I would value adversarial attention on whether the guards actually guard. Highest-value targets:
Known gaps, stated rather than hidden: the six |
|
| Filename | Overview |
|---|---|
| src/tbox_finder/stage2/train.py | Core training entrypoint: eligibility, DDP loop, gate, and checkpoint. Rank-scoping of the gate (early return on non-primary ranks) is correctly placed before validate_report. Eligibility is fail-closed with proper ordering of guards. Three minor issues: _bool_or_none error message labels calib as nested_train; three ternary clauses in derive_clauses are correct but precedence-sensitive; losses_finite cannot observe NaN gradients on non-primary ranks in the final optimizer step. |
| src/tbox_finder/train/ddp.py | DDP helpers promoted from train_stage1 to share one implementation. ShardedSampler correctly truncates to the largest multiple of world_size before striding to prevent DDP deadlock from ragged shards. check_pythonhashseed verifies (not sets) the env var, correctly documenting why in-process setting is a no-op. |
| scripts/backfill_stage2_provenance.py | One-shot repair for job-1064 sidecars. The checkpoint_output_files fork is justified (script must run at the run's commit which predates the shared helper) and pinned by test_backfill_provenance.py. The git_sha guard is sound: write_provenance records the current checkout SHA, which is then compared to the run's SHA; mismatch deletes the sidecar and refuses. |
| src/tbox_finder/stage2/sizing.py | Sizing harness that calls the real forward_backward function the trainer uses. Correctly measures both typical and worst-case batches, includes multiple optimizer steps to capture AdamW lazy allocation, and measures gradient checkpointing effectiveness rather than assuming it. |
| tests/ml/test_stage2_train_smoke.py | Four-tier test suite covering eligibility (pure), Hydra composition, torch step, and committed report. The decoy fixture correctly has label_string=None (matching real data). The AST-level test verifying the non-primary early-return precedes validate_report is a notable guard against refactor regressions. |
| tests/unit/test_backfill_provenance.py | Pins the forked checkpoint_output_files in the backfill script against the shared helper in train.py. Tests both the happy path and both refusal paths (empty dir, non-directory input) to prevent the fork from diverging silently. |
| src/tbox_finder/train/train_stage1.py | Imports DDP helpers from the new shared ddp.py module rather than defining them locally. No functional changes — the promotion consolidates one implementation. |
Sequence Diagram
sequenceDiagram
participant SB as sbatch (SLURM array)
participant R0 as rank 0 (primary)
participant RN as ranks 1-7
participant FS as filesystem
SB->>R0: "torchrun --nproc_per_node=8"
SB->>RN: "torchrun --nproc_per_node=8"
R0->>R0: git snapshot (once, before training)
R0->>R0: "select_rows(rung=train)"
RN->>RN: "select_rows(rung=train)"
loop DDP training epochs
R0->>R0: forward_backward(ddp_model, loss_fn, batch)
RN->>RN: forward_backward(ddp_model, loss_fn, batch)
R0-->>RN: NCCL all-reduce gradients
end
R0->>FS: save checkpoint (adapter + heads)
R0->>FS: write report JSON
R0->>R0: "validate_report -> derive_clauses"
Note over R0: gate checks checkpoint_written,<br/>provenance_recorded (rank-0 artifacts only)
RN->>RN: build_report (not validated)
Note over RN: early return before validate_report
alt gate passes
R0->>SB: exit 0
else gate fails
R0->>SB: "raise RuntimeError, exit != 0"
end
Prompt To Fix All With AI
### Issue 1
src/tbox_finder/stage2/train.py:201
**Misleading `nested_train` label in `_bool_or_none` error message**
`_bool_or_none` was written for `nested_train`, but line 229 calls it for `calib` as well. If a `calib` column ever carries an unexpected value (e.g., `"yes"` from a schema migration), the resulting error reads `nested_train='yes' is neither missing nor boolean`, pointing at the wrong field entirely. The function works correctly — it is just the message that misdirects debugging.
### Issue 2
src/tbox_finder/stage2/train.py:611-617
**Ternary precedence is correct but easy to misread in three clauses**
Python parses `A and B if C else D` as `(A and B) if C else D`, which is the intended semantics in all three conditional clauses (`gradient_checkpointing_verified`, `val_disjoint_from_train`, `checkpoint_written`). However, the expression is visually easy to read as `A and (B if C else D)`, which gives a different result in the `else` branch: for example, `gradient_checkpointing_verified` when checkpointing is *off* would become `_pos_int(n_modules) and (n_modules == 0)` = always False, rather than just `n_modules == 0`. The tests cover the current behaviour, but a future maintainer editing any of these three lines without knowing the rule is likely to introduce a regression.
### Issue 3
src/tbox_finder/stage2/train.py:587-592
**`losses_finite` gate is blind to NaN gradients on non-primary ranks in the final optimizer step**
`n_nonfinite` counts forward-pass loss *values* on rank 0 only. With DDP + NCCL, a NaN gradient produced by a non-primary rank all-reduces into every rank's gradient tensor, poisoning rank 0's parameters after `optimizer.step()`. Normally rank 0's subsequent forward pass would then produce a NaN loss and increment `n_nonfinite`—but if the NaN occurs in the very last batch of the final epoch there is no subsequent forward pass. Rank 0's `final_train_total` was computed before the step, `n_nonfinite_steps` remains 0, `losses_finite` passes, and the saved checkpoint carries NaN-infected weights.
This is a narrow timing window and all six job-1064 points appear numerically clean, so this is not a present defect. The structural gap is that the gate measures rank-0 loss stability but cannot observe gradient-level pathologies on other ranks at the training horizon.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Reviews (1): Last reviewed commit: "fix(p3-06): committing the reports turne..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # ([[pandas-3-nan-truthy-in-training-env]]). `H.NULL_TOKENS` is reused rather than | ||
| # respelled so the two null vocabularies cannot drift apart. | ||
| return None | ||
| raise ValueError(f"nested_train={value!r} is neither missing nor boolean") |
There was a problem hiding this comment.
nested_train label in _bool_or_none error message
_bool_or_none was written for nested_train, but line 229 calls it for calib as well. If a calib column ever carries an unexpected value (e.g., "yes" from a schema migration), the resulting error reads nested_train='yes' is neither missing nor boolean, pointing at the wrong field entirely. The function works correctly — it is just the message that misdirects debugging.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/tbox_finder/stage2/train.py
Line: 201
Comment:
**Misleading `nested_train` label in `_bool_or_none` error message**
`_bool_or_none` was written for `nested_train`, but line 229 calls it for `calib` as well. If a `calib` column ever carries an unexpected value (e.g., `"yes"` from a schema migration), the resulting error reads `nested_train='yes' is neither missing nor boolean`, pointing at the wrong field entirely. The function works correctly — it is just the message that misdirects debugging.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.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!
Comment on lines
+611
to
+617
| # works ([[in-process-no-ops-look-like-compliance]]). | ||
| clauses["gradient_checkpointing_verified"] = ( | ||
| _pos_int(wrap.get("n_modules_with_checkpointing")) | ||
| and wrap.get("checkpoint_use_reentrant") is False | ||
| if wrap.get("gradient_checkpointing") | ||
| else wrap.get("n_modules_with_checkpointing") == 0 | ||
| ) |
There was a problem hiding this comment.
Python parses A and B if C else D as (A and B) if C else D, which is the intended semantics in all three conditional clauses (gradient_checkpointing_verified, val_disjoint_from_train, checkpoint_written). However, the expression is visually easy to read as A and (B if C else D), which gives a different result in the else branch: for example, gradient_checkpointing_verified when checkpointing is off would become _pos_int(n_modules) and (n_modules == 0) = always False, rather than just n_modules == 0. The tests cover the current behaviour, but a future maintainer editing any of these three lines without knowing the rule is likely to introduce a regression.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/tbox_finder/stage2/train.py
Line: 611-617
Comment:
**Ternary precedence is correct but easy to misread in three clauses**
Python parses `A and B if C else D` as `(A and B) if C else D`, which is the intended semantics in all three conditional clauses (`gradient_checkpointing_verified`, `val_disjoint_from_train`, `checkpoint_written`). However, the expression is visually easy to read as `A and (B if C else D)`, which gives a different result in the `else` branch: for example, `gradient_checkpointing_verified` when checkpointing is *off* would become `_pos_int(n_modules) and (n_modules == 0)` = always False, rather than just `n_modules == 0`. The tests cover the current behaviour, but a future maintainer editing any of these three lines without knowing the rule is likely to introduce a regression.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.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!
Comment on lines
+587
to
+592
| clauses["losses_finite"] = ( | ||
| _finite(losses.get("final_train_total")) | ||
| and bool(losses.get("per_term_final")) | ||
| and all(_finite(v) for v in (losses.get("per_term_final") or {}).values()) | ||
| and losses.get("n_nonfinite_steps") == 0 | ||
| ) |
There was a problem hiding this comment.
losses_finite gate is blind to NaN gradients on non-primary ranks in the final optimizer step
n_nonfinite counts forward-pass loss values on rank 0 only. With DDP + NCCL, a NaN gradient produced by a non-primary rank all-reduces into every rank's gradient tensor, poisoning rank 0's parameters after optimizer.step(). Normally rank 0's subsequent forward pass would then produce a NaN loss and increment n_nonfinite—but if the NaN occurs in the very last batch of the final epoch there is no subsequent forward pass. Rank 0's final_train_total was computed before the step, n_nonfinite_steps remains 0, losses_finite passes, and the saved checkpoint carries NaN-infected weights.
This is a narrow timing window and all six job-1064 points appear numerically clean, so this is not a present defect. The structural gap is that the gate measures rank-0 loss stability but cannot observe gradient-level pathologies on other ranks at the training horizon.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/tbox_finder/stage2/train.py
Line: 587-592
Comment:
**`losses_finite` gate is blind to NaN gradients on non-primary ranks in the final optimizer step**
`n_nonfinite` counts forward-pass loss *values* on rank 0 only. With DDP + NCCL, a NaN gradient produced by a non-primary rank all-reduces into every rank's gradient tensor, poisoning rank 0's parameters after `optimizer.step()`. Normally rank 0's subsequent forward pass would then produce a NaN loss and increment `n_nonfinite`—but if the NaN occurs in the very last batch of the final epoch there is no subsequent forward pass. Rank 0's `final_train_total` was computed before the step, `n_nonfinite_steps` remains 0, `losses_finite` passes, and the saved checkpoint carries NaN-infected weights.
This is a narrow timing window and all six job-1064 points appear numerically clean, so this is not a present defect. The structural gap is that the gate measures rank-0 loss stability but cannot observe gradient-level pathologies on other ranks at the training horizon.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.…ptim guard, +5 Review gate met via the CodeRabbit CLI. The GitHub app returned only its fair-usage notice and Copilot returned "quota limit reached"; Greptile was invoked and stayed SILENT for a third consecutive time (PR #98 x2, PR #99 x1), spending a credit for nothing. 11 findings, each verified by execution before being accepted or refused. THE ONE THAT REACHES THE SHIPPED CHECKPOINTS: the LR scheduler's domain was micro-batches while scheduler.step() fires per OPTIMISER step. With accumulation 2, job 1064's cosine covered exactly 50% — 1,415 of 2,830 — ending at 0.5501x base instead of ~0. All six arms trained under a schedule their config does not describe. User decision: fix the code, KEEP the checkpoints, record the defect. Nothing downstream is invalidated (P3-08 compares arms that all share it), but they are not optimally-annealed models and any absolute number derived from them carries that. A CLAIM I MADE WAS WRONG: I reported growth_ratio 1.01-1.04 as a flat PER-STEP footprint. reset_peak_memory_stats() ran once before the loop, so the series was a RUNNING MAXIMUM, monotonic by construction. The no-leak conclusion survives — a leak would still raise a running max — but the instrument was mislabelled and I oversold it. Fixed; the sbatch header repeating the claim is corrected too. Also fixed, each sabotage-bitten: unknown optim.* keys raise as loss.* ones do; accumulation counts per-epoch and flushes the trailing group (283 odd micro-batches made the cumulative phase shift every epoch); the sbatch clears the checkpoint dir only AFTER the health check; WORLD_SIZE=0 raises; the backfill reads optim.lr from its real group and refuses null seed/env_lock/input; sizing COMPUTES headroom instead of asserting it. TWO OF MY OWN NEW TESTS WERE VACUOUS. The LR-domain test asserted arithmetic and never touched train_stage2's wiring — reverting the trainer left it green. Now asserts the shipped construction. (The sbatch test did bite; my sabotage for it was wrong.) Refused with reasons: device_record()'s index 0 is correct — every rank calls set_device(local_rank) first. Validation: trainer 102, sizing 18, overrides 22, backfill 3; ruff + black clean; bash -n clean; 23 sabotages across the step. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (6)
src/tbox_finder/stage2/train.py (3)
1189-1190: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winThe process group is never destroyed.
train_stage2callsdist.init_process_grouphere and no code path callsdist.destroy_process_group(). NCCL then tears down during interpreter shutdown. That produces teardown warnings and can leave a rank blocked at exit, which torchrun reports as a non-zero exit code.The consequence is specific to this job: the sbatch gates the copy out of node-local scratch on
rc, so a teardown-only failure destroys the checkpoint in the same way job 1053 did.Add an explicit teardown that runs on both the success and failure paths, for example a
try/finallyaround the body that callsdist.destroy_process_group()whendist.is_initialized().🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tbox_finder/stage2/train.py` around lines 1189 - 1190, Wrap the distributed training body in train_stage2 with a try/finally so every success and failure path explicitly tears down the process group. In the finally block, check dist.is_initialized() and call dist.destroy_process_group(), preserving the existing initialization behavior and avoiding teardown-time job failures.
1422-1436: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftThe objective and loss evidence is rank-local while the gate is rank 0's.
terms_seen,n_nonfinite,final_totalandper_term_finalare accumulated per rank and are never reduced across ranks. Rank 0 then grades them inobjective_terms_matchandlosses_finite.Two consequences:
- A weighted aux term that finds supervision on other shards but never on rank 0's shard fails
objective_terms_matchafter a healthy run.- A non-finite loss confined to another rank is not counted in
n_nonfinite_steps.This is the same shape as the job-1053 defect the comment at lines 1453-1470 describes: the gate asks one rank about a property of the whole run. An
all_reduceon the non-finite count and anall_gather_objectonterms_seenbeforebuild_reportwould make the recorded evidence match the scope of the verdict.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tbox_finder/stage2/train.py` around lines 1422 - 1436, Aggregate rank-local evidence before constructing the report: all-reduce n_nonfinite across ranks and all-gather terms_seen, merging the gathered terms according to L.TERMS order. Ensure final_total and per_term_final are also reduced or otherwise made globally representative before build_report, so objective_terms_match and losses_finite evaluate whole-run evidence rather than rank 0’s local state.
1343-1346: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winEvery rank recomputes the whole validation pass.
evaluatebuildsorder = list(range(len(dataset)))at line 1097 with no rank sharding, and this call site runs on all ranks. With 8 ranks and 2,070 validation rows, the node performs 8 identical forward passes over the same rows every epoch.evaluateuses the unwrappedmodel, so no collective is involved and no rank depends on another's result.Two options: shard the validation order with
ddp.ShardedSamplerand reduce the totals, or run the pass on the primary rank and broadcastbest_val. Sharding keeps the ranks in lockstep and cuts the epoch tail by roughlyworld_size.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tbox_finder/stage2/train.py` around lines 1343 - 1346, Update the validation flow around evaluate so distributed ranks do not process the full validation dataset independently. Prefer sharding validation samples across ranks, reducing aggregate loss/count metrics, and ensuring every rank receives the same final metrics while remaining synchronized; preserve the existing single-rank behavior when distributed evaluation is disabled.src/tbox_finder/stage2/sizing.py (2)
190-197: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: keep the peak observed before an OOM.
An OOM point discards
overall_peakeven when steps completed. The committed report shows this: the batch-8 worst case ran one step at 14.9686 GiB and recordspeak_vram_gib: null. The per-step series retains the value, so no data is lost, but the "how close did it get" number needs a reader to open the series.Consider recording it under a separate key so
peak_vram_gibkeeps meaning "a batch that completed the sweep".♻️ Proposed optional change
except Exception as exc: # noqa: BLE001 - an OOM here is a RESULT, not a crash is_oom = "OutOfMemoryError" in type(exc).__name__ or "out of memory" in str(exc).lower() record["oom"] = bool(is_oom) record["error"] = f"{type(exc).__name__}: {str(exc)[:300]}" + if is_oom and record["peak_vram_gib_per_step"]: + record["peak_vram_gib_before_oom"] = max(record["peak_vram_gib_per_step"]) if not is_oom: raise🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tbox_finder/stage2/sizing.py` around lines 190 - 197, Preserve the last observed VRAM peak when an OOM occurs by recording overall_peak under a separate OOM-specific key in the exception handler, while leaving peak_vram_gib unset so it continues to represent a completed sweep. Update the handler around overall_peak and the existing record["oom"] assignment without changing non-OOM error behavior.
399-423: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winMake the swept checkpointing regime explicit and aligned with the training config.
Line 407 hardcodes
gradient_checkpointing=Truefor every sweep point, solargest_fitting_batch_worst_caseandrecommendationdescribe the ON regime.conf/train/stage2.yamlships the flag FALSE, per the drift guard in tests/ml/test_stage2_sizing.py lines 301-331. The two regimes match today only because the hook is a no-op, which that same test says may change.Expose the regime as a parameter and record it in
config, so the recommendation states which regime produced it. Also consider running the on/off comparison at the recommended batch size rather than atmin(batch_sweep), because batch 1 gives the weakest signal for the batch production will use.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/tbox_finder/stage2/sizing.py` around lines 399 - 423, Update the Stage 2 sizing flow around the batch sweep and checkpointing comparison to accept and use the configured gradient-checkpointing regime instead of hardcoding True; expose that regime in the returned config so largest_fitting_batch_worst_case and recommendation identify whether they are ON or OFF. Run the checkpointing on/off comparison at the recommended batch size rather than min(batch_sweep), while preserving the existing sweep measurements and reporting behavior.tests/unit/test_backfill_provenance.py (1)
44-77: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the refusal paths, not only the enumerator.
The three tests pin
checkpoint_output_filesagainsttrain.checkpoint_output_files. Nothing exercises_sidecarormain, so the fail-closed behaviour the script exists to guarantee is untested:
_sidecarraisesValueErrorwhenprovenance.seed,provenance.env_lock,provenance.git_sha, ordata.dataset_parquetis absent. A regression that dropped one field would write a sidecar asserting provenance over nothing.mainrefuses a point whose report hasgate.overall_passfalse.maindeletes the sidecar and refuses when the checkoutgit_shadiffers from the run's. This is the guard the PR description asks reviewers to confirm, and it depends onwrite_provenanceemitting a top-levelgit_sha. A test would settle that dependency.- Repeated
--writeagainst the same checkpoint should produce the sameoutputslist. That case exposes the sidecar-self-reference defect noted onscripts/backfill_stage2_provenance.py.Each case needs only a temporary directory, a small report dict, and a monkeypatch of
SWEEP_DIRandCKPT_ROOT.Do you want me to write these tests?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/test_backfill_provenance.py` around lines 44 - 77, Extend tests in test_backfill_provenance.py beyond checkpoint_output_files: add refusal-path coverage for _sidecar when each required provenance.seed, provenance.env_lock, provenance.git_sha, or data.dataset_parquet field is missing; verify main rejects reports with gate.overall_pass false; verify git_sha mismatch removes the sidecar and raises; and verify repeated --write produces an identical outputs list without self-referencing the sidecar. Use temporary directories with small report dictionaries and monkeypatch SWEEP_DIR and CKPT_ROOT.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@reports/p3/sweep/aux0.0_lr1e-4.json`:
- Around line 301-319: Add route-specific binary accuracy metrics to evaluate()
by separating validation records into the nested_train and parentless_decoy
routes before P3-08 selection. Report each route’s accuracy alongside the
existing aggregate accuracy, preserving the current aggregate metric and
ensuring both validation subsets are measured independently.
In `@reports/p3/sweep/aux1.0_lr1e-4.json`:
- Around line 3-12: Update the report writer’s checkpoint metadata so
best_val_epoch and best_val_total are not presented as properties of final-epoch
weights; either persist the best-epoch weights, or rename those fields to
explicitly identify them as validation-history values. Also address the six
committed reports consistently, documenting or correcting their metadata
according to the chosen interpretation.
In `@reports/p3/sweep/aux1.0_lr3e-4.json`:
- Around line 209-227: Rename the report gate clause currently emitted as
provenance_recorded to report_provenance_block_complete in the report writer,
including any references used to compute or serialize that clause. Preserve its
existing in-process provenance-block validation and leave the separate sidecar
verification in stage2_lora_finetune.sbatch unchanged.
In `@scripts/backfill_stage2_provenance.py`:
- Around line 104-105: Update the “outputs” construction in the backfill flow
around write_provenance to exclude the provenance.json sidecar from
checkpoint_output_files(ckpt) at the call site. Keep checkpoint_output_files
unchanged so its existing contract and test remain intact, while ensuring
repeated --write runs produce records that do not attest to the sidecar itself.
In `@slurm/p3/stage2_lora_finetune.sbatch`:
- Around line 319-330: Move the stale-checkpoint removal block containing the
KEY validation, rm -rf, and mkdir from its current pre-training position to
immediately before the successful copy from $SCRATCH_CKPT to $CKPT_DIR. Ensure
it runs only after torchrun, validation gates, and report freshness checks
confirm a valid artifact, so any failed run preserves the existing checkpoint.
In `@src/tbox_finder/stage2/sizing.py`:
- Around line 279-284: Update validate_report in
src/tbox_finder/stage2/sizing.py at lines 279-284 to re-derive recommendation
via _recommend using measurements and device, then append a validation problem
when it differs from report["recommendation"]. Regenerate
reports/p3/stage2_sizing.json at lines 279-282 with the current harness so
recommendation includes worst_case_peak_gib, device_total_gib, headroom_gib, and
the basis documenting DDP exclusion.
- Around line 158-165: Update the record initialization and batching logic
around Stage2SequenceDataset and collate_stage2 so record["n_examples"] stores
the actual number of elements passed to collate_stage2, using the truncated
count rather than the requested batch_size, while keeping the key present when
no batch is available.
In `@src/tbox_finder/stage2/train.py`:
- Line 381: The num_workers setting is unused because training batches are built
through _batched rather than a DataLoader. In
src/tbox_finder/stage2/train.py:381, remove the num_workers dataclass field and
its diagnostics output; in conf/train/stage2.yaml:71, remove the corresponding
num_workers configuration key.
- Around line 924-926: Replace the attribute-based logic in
_count_checkpointed_modules and its consumers build_model and derive_clauses
with evidence that checkpointing actually reduces peak allocation, using the
existing fixed-batch sizing comparison produced by stage2_sizing.json;
alternatively, explicitly downgrade both consumers’ claims so the count is not
treated as proof of effective checkpointing.
- Around line 581-586: Extend the steps_ran condition in the training-stage
validation to require measured n_optimizer_steps to equal
expected_n_optimizer_steps, alongside the existing positive-value and n_steps
checks. Update the steps fixture in tests/ml/test_stage2_train_smoke.py to
provide both optimizer-step keys so the runtime comparison is exercised, while
preserving the existing scheduler and micro-batch validations.
In `@tests/ml/test_stage2_train_smoke.py`:
- Around line 1100-1102: Remove the dead conditional expression in the AST
parsing setup around T.train_stage2, passing source.lstrip() directly to
ast.parse. Retain inspect.getsource and its import because they are still
required to obtain the source text.
- Around line 405-414: Update the steps fixture used by the stage-2 smoke test
to match the shipped train_stage2 report: replace total_scheduled_steps with
expected_n_optimizer_steps, optimizer_steps_per_epoch_per_rank,
gradient_accumulation_steps, and total_scheduled_optimizer_steps, using values
consistent with the accumulation-2 regime. Adjust n_steps and n_optimizer_steps
so they reflect gradient accumulation rather than being identical, and preserve
the existing derive_clauses expectations.
---
Nitpick comments:
In `@src/tbox_finder/stage2/sizing.py`:
- Around line 190-197: Preserve the last observed VRAM peak when an OOM occurs
by recording overall_peak under a separate OOM-specific key in the exception
handler, while leaving peak_vram_gib unset so it continues to represent a
completed sweep. Update the handler around overall_peak and the existing
record["oom"] assignment without changing non-OOM error behavior.
- Around line 399-423: Update the Stage 2 sizing flow around the batch sweep and
checkpointing comparison to accept and use the configured gradient-checkpointing
regime instead of hardcoding True; expose that regime in the returned config so
largest_fitting_batch_worst_case and recommendation identify whether they are ON
or OFF. Run the checkpointing on/off comparison at the recommended batch size
rather than min(batch_sweep), while preserving the existing sweep measurements
and reporting behavior.
In `@src/tbox_finder/stage2/train.py`:
- Around line 1189-1190: Wrap the distributed training body in train_stage2 with
a try/finally so every success and failure path explicitly tears down the
process group. In the finally block, check dist.is_initialized() and call
dist.destroy_process_group(), preserving the existing initialization behavior
and avoiding teardown-time job failures.
- Around line 1422-1436: Aggregate rank-local evidence before constructing the
report: all-reduce n_nonfinite across ranks and all-gather terms_seen, merging
the gathered terms according to L.TERMS order. Ensure final_total and
per_term_final are also reduced or otherwise made globally representative before
build_report, so objective_terms_match and losses_finite evaluate whole-run
evidence rather than rank 0’s local state.
- Around line 1343-1346: Update the validation flow around evaluate so
distributed ranks do not process the full validation dataset independently.
Prefer sharding validation samples across ranks, reducing aggregate loss/count
metrics, and ensuring every rank receives the same final metrics while remaining
synchronized; preserve the existing single-rank behavior when distributed
evaluation is disabled.
In `@tests/unit/test_backfill_provenance.py`:
- Around line 44-77: Extend tests in test_backfill_provenance.py beyond
checkpoint_output_files: add refusal-path coverage for _sidecar when each
required provenance.seed, provenance.env_lock, provenance.git_sha, or
data.dataset_parquet field is missing; verify main rejects reports with
gate.overall_pass false; verify git_sha mismatch removes the sidecar and raises;
and verify repeated --write produces an identical outputs list without
self-referencing the sidecar. Use temporary directories with small report
dictionaries and monkeypatch SWEEP_DIR and CKPT_ROOT.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: e1a3972d-c114-481e-b23b-3fda7669a2c1
⛔ Files ignored due to path filters (3)
analyses/phase3_log.qmdis excluded by!**/*.qmddata/processed/checkpoints/.gitignoreis excluded by!data/**data/processed/checkpoints/stage2_rinalmo.dvcis excluded by!data/**
📒 Files selected for processing (19)
conf/optim/stage2.yamlconf/train/stage2.yamlreports/p3/stage2_sizing.jsonreports/p3/sweep/aux0.0_lr1e-4.jsonreports/p3/sweep/aux0.0_lr3e-4.jsonreports/p3/sweep/aux0.5_lr1e-4.jsonreports/p3/sweep/aux0.5_lr3e-4.jsonreports/p3/sweep/aux1.0_lr1e-4.jsonreports/p3/sweep/aux1.0_lr3e-4.jsonscripts/backfill_stage2_provenance.pyslurm/p3/stage2_lora_finetune.sbatchslurm/p3/stage2_sizing_smoke.sbatchsrc/tbox_finder/stage2/sizing.pysrc/tbox_finder/stage2/train.pysrc/tbox_finder/train/ddp.pysrc/tbox_finder/train/train_stage1.pytests/ml/test_stage2_sizing.pytests/ml/test_stage2_train_smoke.pytests/unit/test_backfill_provenance.py
Comment on lines
+301
to
+319
| "val": { | ||
| "best_epoch": 6, | ||
| "best_total": 0.00037977641439557516, | ||
| "history": [ | ||
| { | ||
| "binary_accuracy": 0.996135265700483, | ||
| "epoch": 0, | ||
| "n_batches": 259, | ||
| "n_records": 2070, | ||
| "per_term": { | ||
| "binary": 0.017270480773605417, | ||
| "boundary": 1.4936640283315799, | ||
| "cognate_aa": 2.249292553161562, | ||
| "regulatory_mode": 1.0713178186342984, | ||
| "specifier_codon": 2.9770189583531677, | ||
| "trna_family": 2.9287205672172045 | ||
| }, | ||
| "total": 0.017270480773605417 | ||
| }, |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Look for per-route or per-decoy-class validation metrics in the Stage-2 trainer.
set -euo pipefail
fd -t f 'train.py' src/tbox_finder/stage2 --exec ast-grep outline {} --items all
rg -n -C4 'binary_accuracy|admitted_by_route|parentless_decoy|fold_basis' \
--glob 'src/**/*.py' --glob 'tests/**/*.py'Repository: bioedca/tbox-finder
Length of output: 4629
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- trainer selection, evaluation, and report construction ---'
sed -n '200,290p;760,900p;1071,1138p;1158,1245p;1245,1335p' src/tbox_finder/stage2/train.py
printf '%s\n' '--- route and dataset definitions ---'
sed -n '1,280p' src/tbox_finder/stage2/dataset.py
printf '%s\n' '--- report files and route-related fields ---'
fd -t f -e json reports | head -80
rg -n -C3 '"binary_accuracy"|"admit_route"|"admitted_by_route"|"parentless_decoy"|"fold_basis"|"aux_weight"|"lr"' reports src tests \
--glob '*.json' --glob '*.py' | head -300Repository: bioedca/tbox-finder
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import json
from pathlib import Path
p = Path("reports/p3/sweep/aux0.0_lr1e-4.json")
obj = json.loads(p.read_text())
print("top-level keys:", sorted(obj))
for k, v in obj.items():
if isinstance(v, dict):
print(k, "keys:", sorted(v))
if k in {"train", "val"}:
print(k, "history length:", len(v.get("history", [])))
print(k, "sample history keys:", sorted(v["history"][0]) if v.get("history") else [])
PY
printf '%s\n' '--- all sweep report names ---'
fd -t f -e json reports/p3/sweep | sort
printf '%s\n' '--- route-related source references ---'
rg -n -C5 'def _admit_route|ADMIT_|SOURCE_DECOY|fold_basis|binary_accuracy|evaluate\(' src/tbox_finder/stage2 testsRepository: bioedca/tbox-finder
Length of output: 2875
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '200,290p;760,900p;1071,1138p;1158,1335p' src/tbox_finder/stage2/train.py
sed -n '1,280p' src/tbox_finder/stage2/dataset.py
rg -n -C3 'binary_accuracy|admitted_by_route|parentless_decoy|fold_basis|aux_weight|lr' reports src tests \
--glob '*.json' --glob '*.py'Repository: bioedca/tbox-finder
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import json
from pathlib import Path
paths = sorted(Path("reports/p3/sweep").glob("*.json"))
print("p3 sweep reports:", len(paths))
for p in paths:
r = json.loads(p.read_text())
d = r["data"]
cfg = r["config"]
tr = d["train_census"]
va = d["val_census"]
hist = r.get("val", {}).get("history", [])
first = hist[0] if hist else {}
last = hist[-1] if hist else {}
print(p.name)
print(" config:", {"aux_weight": cfg.get("loss", {}).get("aux_weight"), "lr": cfg.get("lr"),
"train_rung": cfg.get("train_rung"), "val_rung": cfg.get("val_rung"),
"admit_parentless_decoys": cfg.get("admit_parentless_decoys")})
print(" train:", {"n_rows": d.get("n_train_rows"), "n_admitted": tr.get("n_admitted"),
"routes": tr.get("admitted_by_route"), "positive": tr.get("n_admitted_positive"),
"negative": tr.get("n_admitted_negative")})
print(" val:", {"n_rows": d.get("n_val_rows"), "n_admitted": va.get("n_admitted"),
"routes": va.get("admitted_by_route"), "positive": va.get("n_admitted_positive"),
"negative": va.get("n_admitted_negative")})
print(" val binary:", {"epoch0": first.get("binary_accuracy"), "last": last.get("binary_accuracy")})
PY
printf '%s\n' '--- exact report schema checks ---'
python3 - <<'PY'
import json
from pathlib import Path
for p in sorted(Path("reports/p3/sweep").glob("*.json")):
r=json.loads(p.read_text())
print(p.name, "route metrics in val history:",
[k for k in (r["val"]["history"][0] if r["val"]["history"] else {}) if "route" in k or "decoy" in k])
PYRepository: bioedca/tbox-finder
Length of output: 3442
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
p = Path("src/tbox_finder/stage2/train.py")
tree = ast.parse(p.read_text())
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "evaluate":
print("evaluate range:", node.lineno, node.end_lineno)
print("local names:", sorted({n.id for n in ast.walk(node) if isinstance(n, ast.Name)}))
print("route-related names:", sorted({n.id for n in ast.walk(node) if isinstance(n, ast.Name)
and ("route" in n.id.lower() or "decoy" in n.id.lower())}))
print("return keys:")
for sub in ast.walk(node):
if isinstance(sub, ast.Dict):
keys = [k.value for k in sub.keys if isinstance(k, ast.Constant)]
if keys:
print(keys)
PY
python3 - <<'PY'
import json
from pathlib import Path
for p in sorted(Path("reports/p3/sweep").glob("*.json")):
r=json.loads(p.read_text())
print(p.name, "gate:", r["gate"]["overall_pass"], "failed:", r["gate"]["failed"])
print(" clauses:", r["gate"]["clauses"])
PYRepository: bioedca/tbox-finder
Length of output: 3276
Add route-specific binary metrics before P3-08 selection.
evaluate() reports only aggregate accuracy. It does not separate nested_train from parentless_decoy, although validation contains 1,581 and 489 rows from these routes. Aggregate epoch-0 accuracy is already 0.9937–0.9966 across the six arms. Add route-specific accuracy metrics before selecting on these results.
🧰 Tools
🪛 OpenGrep (1.26.0)
[ERROR] 312-312: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@reports/p3/sweep/aux0.0_lr1e-4.json` around lines 301 - 319, Add
route-specific binary accuracy metrics to evaluate() by separating validation
records into the nested_train and parentless_decoy routes before P3-08
selection. Report each route’s accuracy alongside the existing aggregate
accuracy, preserving the current aggregate metric and ensuring both validation
subsets are measured independently.
Comment on lines
+209
to
+227
| "gate": { | ||
| "clauses": { | ||
| "both_admit_routes_fired": true, | ||
| "calib_refused": true, | ||
| "checkpoint_written": true, | ||
| "full_population": true, | ||
| "gradient_checkpointing_verified": true, | ||
| "heads_outside_wrapper": true, | ||
| "lora_contract_held": true, | ||
| "losses_finite": true, | ||
| "no_d5_holdout_trained": true, | ||
| "objective_terms_match": true, | ||
| "provenance_recorded": true, | ||
| "steps_ran": true, | ||
| "val_disjoint_from_train": true | ||
| }, | ||
| "failed": [], | ||
| "overall_pass": true | ||
| }, |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
provenance_recorded passed on the run whose sidecar write failed.
This arm is the job-1064 point named in scripts/backfill_stage2_provenance.py. The clause reads true and gate.overall_pass reads true, yet no provenance.json was written for this checkpoint. The clause therefore measures the report's own provenance block, which train.py fills in-process, and not the checkpoint sidecar.
That is defensible behaviour, because train.py cannot observe a write the sbatch performs afterwards. The name is what misleads: a reader auditing the sweep sees provenance_recorded: true and concludes the checkpoint carries provenance. Rename the clause in the report writer to name what it checks, for example report_provenance_block_complete.
The verification procedure in slurm/p3/stage2_lora_finetune.sbatch Lines 92-96 already checks the sidecar separately, so no coverage is lost.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@reports/p3/sweep/aux1.0_lr3e-4.json` around lines 209 - 227, Rename the
report gate clause currently emitted as provenance_recorded to
report_provenance_block_complete in the report writer, including any references
used to compute or serialize that clause. Preserve its existing in-process
provenance-block validation and leave the separate sidecar verification in
stage2_lora_finetune.sbatch unchanged.
Owner
Author
|
@coderabbitai review r2 over the fixes applied from r1. Eleven findings were reviewed; each was verified by execution before being accepted or refused, and the changes are in Applied:
Refused, with reason: Please prioritise:
|
|
I will prioritize the optimizer-step count at epoch boundaries, the scheduler-step domain, and whether the new tests fail when the corrected behavior is removed. ✅ Action performedReview finished.
|
…ts never achieved 12 findings, each verified by execution. r2 was sharper than r1 because it reviewed r1's fixes. THE ONE THAT CORRECTS A REPORTED NUMBER: the checkpoint block carried best_val_total while saved_from_epoch is 9 in every arm — there is no best-checkpoint restore. So best_val_* described epochs whose weights were discarded: point saved best reported SAVED weights score aux0.0_lr1e-4 9 6 0.000380 0.008841 (23x) aux0.0_lr3e-4 9 8 0.000599 0.002500 (4x) aux1.0_lr1e-4 9 9 0.039954 0.039954 (best == final) The RESULT table I wrote quoted the best-epoch column. Fixed: `saved_val_total` is now in the checkpoint block with a `note` saying the weights are the final epoch's and arms compare on it; best_val_* is renamed *_observed_during_training. AND r1 HAD POINTED AT THIS — it flagged the config comment calling the val rung a "best-checkpoint signal", a minor I filed as wording and folded into the schedule fix instead of following to the artifact. A comment claiming behaviour the code lacks is a defect report about the CODE; editing the comment is the wrong fix when the comment describes what a reader would want to be true. Other majors, each sabotage-bitten: steps_ran now grades the optimiser domain too (it could not have seen its own regression return); gradient_checkpointing_verified renamed gradient_checkpointing_flag_consistent (counting an attribute on 36 modules while the hook is never called is not verification); the sizing recommendation is re-derived by the validator; the sbatch checkpoint clear moved again, to immediately before the copy — past the health check still let an OOM end a point with nothing and the previous arm deleted; the backfill excludes its own sidecar so --write is idempotent; num_workers deleted (no DataLoader exists, so reports claimed a setting never applied). THREE MORE OF MY TESTS WERE WEAKER THAN THEY LOOKED, all found by sabotage: steps_ran passed with either comparison removed (they covered for each other), the sbatch test asserted only "after the health check" rather than the property r2 established, and one anchor was stale. Five such tests across this step now. Validation: trainer 110, sizing 20, overrides 22, backfill 4; ruff + black clean; bash -n clean; 28 sabotages across the step. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…aim I put in CLAUDE.md Greptile DID review PR #99 — ~18 minutes after the trigger, as greptile-apps[bot], with 3 P2 findings. A 6-minute poll had concluded "silent", reported the credit wasted, and WROTE INTO CLAUDE.md §5.1a that the trigger path was dead. That was a false conclusion from an impatient poll, promoted to documentation where a future session would have obeyed it. §5.1a corrected: it works, budget ~25 min, and look in pulls/<n>/reviews for greptile-apps[bot] — the reply is a REVIEW, not an issue comment, so a poll reading only issue comments misses it. The prior PR #98 "silences" are now suspect for the same reason. Findings, all real: (P2) losses_finite was blind to a NaN gradient on any non-primary rank. n_nonfinite counts the RANK-0 FORWARD LOSS; under DDP a NaN gradient from another rank all-reduces into every rank's gradients and poisons the weights while rank 0's loss stays finite. Fixed by reading clip_grad_norm_'s pre-clip total norm — computed after the all-reduce, non-finite iff some gradient is — as n_nonfinite_grad_steps, now a second term in the clause. This is the THIRD gate in this step that measured the wrong side of what it named, after the checkpointing flag-count and the micro-batch-only step check. (P2) _bool_or_none blamed the wrong column: written for nested_train, also called for calib and is_tbox, so a "yes" in calib would raise an error naming nested_train. Now takes the field name; every call site passes its own. (P2, advisory) Ternary precedence in three clauses — semantics confirmed correct, readability only. Left as-is and recorded rather than silently dropped. Also corrected in §5.1: CodeRabbit is INCREMENTAL and AUTOMATIC. @coderabbitai review on an already-reviewed commit is a no-op. Both rounds I attributed to manual triggers were the automatic reviews of the pushes; coverage reads from original_commit_id, not from whether a command was sent. And a green CodeRabbit CHECK does not mean a review ran — it read pass while the body was only the Fair-Usage notice. Validation: trainer 113, sizing 20, overrides 22, backfill 4; ruff + black clean; bash -n clean. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
Owner
Author
|
@coderabbitai review Automatic review on push is now disabled for this repo, so this is a manual request for the commits since your last round. Reviewed so far:
Please prioritise:
|
🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
echo '=== Commit summaries ==='
git show --stat --oneline --decorate 2b3bf6d c6fc074
git diff --unified=80 6424a51..HEAD -- \
src/tbox_finder/stage2/train.py \
src/tbox_finder/stage2/sizing.py \
tests/ml/test_stage2_train_smoke.py \
tests/ml/test_stage2_sizing.py \
'reports/p3/sweep/*.json'
echo '=== Stage-2 training outline ==='
ast-grep outline src/tbox_finder/stage2/train.py --items all --type function,class --view expanded
echo '=== Targeted implementations ==='
rg -n -A100 -B25 \
'def (derive_clauses|validate_report|forward_backward|train_stage2|_bool_or_none)\b|clip_grad_norm_|saved_val_total|n_nonfinite_grad_steps|best_val_|steps_ran|total_opt_steps' \
src/tbox_finder/stage2/train.py
echo '=== Targeted test assertions ==='
rg -n -A35 -B12 \
'nonfinite|grad_clip|clip_grad_norm_|saved_val_total|best_val_|total_opt_steps|scheduler|vacuous|mutation' \
tests/ml/test_stage2_train_smoke.py tests/ml/test_stage2_sizing.py
echo '=== Report metric references ==='
rg -n -C4 \
'"(saved_val_total|best_val|best_val_total_observed_during_training|n_nonfinite_grad_steps|steps_ran|total_opt_steps)"' \
reports/p3/sweepLength of output: 215263 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo '=== PyTorch dependency pins ==='
rg -n -C3 '(torch|pytorch)' pyproject.toml requirements.txt environment.yml conda*.yml envs 2>/dev/null || true
echo '=== Committed sweep report checkpoint and validation fields ==='
for report in reports/p3/sweep/*.json; do
echo "--- $report"
jq '{
checkpoint,
val,
losses,
steps,
gate: .gate.clauses
}' "$report"
done
echo '=== All saved/best metric identifiers in committed reports and analysis ==='
rg -n -C2 '(saved_val_total|saved_from_epoch|best_val|best_total|best_epoch|final.*epoch|epoch.*final)' \
reports/p3/sweep analyses/phase3_log.qmd
echo '=== Exact changed test regions ==='
sed -n '1315,1435p' tests/ml/test_stage2_train_smoke.py
sed -n '1000,1060p' tests/ml/test_stage2_train_smoke.py
sed -n '1170,1205p' tests/ml/test_stage2_train_smoke.py
echo '=== DDP configuration and backward call context ==='
sed -n '915,940p' src/tbox_finder/stage2/train.py
sed -n '1265,1375p' src/tbox_finder/stage2/train.py
echo '=== Git changes to reports since reviewed commit ==='
git diff --name-status 6424a51..HEAD -- reports/p3/sweep
git diff --unified=20 6424a51..HEAD -- reports/p3/sweepLength of output: 59345 🤖 web_search executed:
Length of output: 11898 🤖 web_search executed:
Length of output: 10623 🏁 Script executed: #!/bin/bash
set -u
echo '=== Report files changed after 6424a51 ==='
git diff --name-status 6424a51..HEAD -- reports/p3/sweep || true
git diff --unified=8 6424a51..HEAD -- reports/p3/sweep || true
echo '=== One-line summary for all committed sweep reports ==='
for report in reports/p3/sweep/*.json; do
jq -c '{
file: input_filename,
checkpoint: {
saved_from_epoch,
saved_val_total,
best_val_epoch,
best_val_total,
best_val_epoch_observed_during_training,
best_val_total_observed_during_training
},
val: {best_epoch: .val.best_epoch, best_total: .val.best_total,
final_epoch: (.val.history[-1].epoch), final_total: (.val.history[-1].total)},
loss_fields: (.losses | keys),
step_fields: (.steps | keys),
gate_fields: (.gate.clauses | keys)
}' "$report"
done
echo '=== Complete new regression tests ==='
sed -n '1195,1435p' tests/ml/test_stage2_train_smoke.py
echo '=== Locate runtime training tests and mocking seams ==='
rg -n -A45 -B12 \
'train_stage2\(|monkeypatch.*(forward_backward|clip_grad_norm_|optimizer|scheduler)|gradient_accumulation|grad_clip' \
tests/ml/test_stage2_train_smoke.py
echo '=== Current PyTorch 2.7 clipping implementation availability ==='
git ls-files | rg '(^|/)(ml-rna\.conda-lock\.yml|ml-rna\.yml)$'Length of output: 30192 🤖 web_search executed:
Length of output: 19498
The DDP signal itself is correct. DDP waits for all bucket all-reduces before ✅ Action performedReview finished.
|
… excusing it
All seven committed reports failed their own validators after r2/r3 strengthened the clause
sets — found by running the validators against the committed artifacts rather than assuming
they still applied. SCHEMA_VERSION bumped 1 -> 2 on both modules: a clause set is part of a
report's shape, so adding one makes earlier reports a different schema, not a corrupt one.
The two failures are categorically different and must not be conflated:
steps_ran — a REAL defect. Job 1064 sized the cosine in micro-batches; the optimiser
advanced 1,415 of 2,830 scheduled steps. Schema 2 refuses those runs
correctly. The habit on a version bump is to excuse a clause the old schema
lacks — right only when the old report re-derives TRUE. This re-derives
FALSE for a real reason, so excusing it would launder a known defect
through a compatibility shim. Asserted as failing instead.
losses_finite — NOT a defect, an absence. Schema 1 predates n_nonfinite_grad_steps (added
for Greptile's NaN-gradient finding), so the post-all-reduce check was never
measured. Unmeasured is neither passed nor failed; excusing it into TRUE
would assert those runs had finite gradients when nobody looked.
Pinned by test: six reports present (emptiness guard first), all schema 1, failing EXACTLY
["losses_finite", "steps_ran"] — a new entry means a second defect, a missing one means a
clause stopped biting or the reports were regenerated. The distinction is checked, not
described: the scheduler clause has its evidence and fails on it, the gradient clause has no
evidence at all, and what schema 1 did measure was clean.
The reports are kept as written — they are job 1064's own record, and the decision was to
keep the checkpoints and RECORD the defect. Regenerating them to pass is the opposite.
Validation: trainer 116, sizing 20, overrides 22, backfill 4; ruff + black clean.
Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
tests/ml/test_stage2_train_smoke.py (1)
409-418: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winThe optimizer-step count in the fixture contradicts the trailing-flush arithmetic this file asserts elsewhere.
With
batches_per_epoch_per_rank = 283andgradient_accumulation_steps = 2, the shipped loop flushes the trailing partial group, so each epoch performs_n_batches(283, 2) = 142optimizer steps. That matchesoptimizer_steps_per_epoch_per_rank, and over 10 epochs it gives 1420. The fixture instead records 1415, which is2830 // 2and therefore drops one flush per epoch.test_the_predicted_optimiser_step_count_equals_what_the_loop_actually_doesproves 1420 is the correct value for this shape.
derive_clausesdoes not comparen_optimizer_stepsagainstoptimizer_steps_per_epoch_per_rank, so the inconsistency stays silent. That is the exact drift the comment on Lines 406-408 says the fixture must prevent.🔧 Proposed fixture correction
"n_steps": 2830, "expected_n_steps": 2830, - "n_optimizer_steps": 1415, - "expected_n_optimizer_steps": 1415, + "n_optimizer_steps": 1420, + "expected_n_optimizer_steps": 1420, "batches_per_epoch_per_rank": 283, "optimizer_steps_per_epoch_per_rank": 142, "gradient_accumulation_steps": 2, "world_size": 8, - "warmup_steps": 85, - "total_scheduled_optimizer_steps": 1415, + "warmup_steps": 85, + "total_scheduled_optimizer_steps": 1420,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ml/test_stage2_train_smoke.py` around lines 409 - 418, Update the fixture values for the 283-batch, 2-step accumulation case so trailing partial groups are counted: change n_optimizer_steps, expected_n_optimizer_steps, and total_scheduled_optimizer_steps from 1415 to 1420. Keep optimizer_steps_per_epoch_per_rank at 142 and ensure the related expected values remain internally consistent with the loop’s flush behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/tbox_finder/stage2/sizing.py`:
- Around line 261-263: Update the sizing completion logic around the
with_numbers-derived clauses to evaluate every non-OOM record, including those
with peak_vram_gib=None, so enough_steps_for_optimizer_state and
measured_the_requested_batch cannot pass vacuously. Make swept require a
recorded peak_vram_gib for each non-OOM record, and add a regression case
covering oom=False, n_steps>0, and peak_vram_gib=None.
---
Duplicate comments:
In `@tests/ml/test_stage2_train_smoke.py`:
- Around line 409-418: Update the fixture values for the 283-batch, 2-step
accumulation case so trailing partial groups are counted: change
n_optimizer_steps, expected_n_optimizer_steps, and
total_scheduled_optimizer_steps from 1415 to 1420. Keep
optimizer_steps_per_epoch_per_rank at 142 and ensure the related expected values
remain internally consistent with the loop’s flush behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 0cdc8f19-3205-45c3-8b1e-5b9ca916e258
⛔ Files ignored due to path filters (1)
analyses/phase3_log.qmdis excluded by!**/*.qmd
📒 Files selected for processing (8)
conf/train/stage2.yamlscripts/backfill_stage2_provenance.pyslurm/p3/stage2_lora_finetune.sbatchsrc/tbox_finder/stage2/sizing.pysrc/tbox_finder/stage2/train.pytests/ml/test_stage2_sizing.pytests/ml/test_stage2_train_smoke.pytests/unit/test_backfill_provenance.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tests/unit/test_backfill_provenance.py
- scripts/backfill_stage2_provenance.py
- slurm/p3/stage2_lora_finetune.sbatch
- src/tbox_finder/stage2/train.py
Comment on lines
+261
to
+263
| "measured_the_requested_batch": all( | ||
| m.get("measured_batch_size") == m.get("batch_size") for m in with_numbers | ||
| ), |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Require VRAM evidence for every completed measurement.
with_numbers excludes non-OOM records that lack peak_vram_gib. Such records still satisfy swept when they have steps, but both enough_steps_for_optimizer_state and measured_the_requested_batch can pass vacuously.
Build these clauses from every non-OOM record. Also make swept reject a non-OOM record without a recorded peak. Add a regression case with oom=False, n_steps>0, and peak_vram_gib=None.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/tbox_finder/stage2/sizing.py` around lines 261 - 263, Update the sizing
completion logic around the with_numbers-derived clauses to evaluate every
non-OOM record, including those with peak_vram_gib=None, so
enough_steps_for_optimizer_state and measured_the_requested_batch cannot pass
vacuously. Make swept require a recorded peak_vram_gib for each non-OOM record,
and add a regression case covering oom=False, n_steps>0, and peak_vram_gib=None.
…corrupting gradients Three findings from the manual review (auto-review is disabled on this repo, so the manual comment is the live path). All real. (High) The producer schema was fixed in r2 but the ARTIFACTS stayed misleading: a reader saw best_val_total: 0.00037977 beside saved_from_epoch: 9 and had to infer they describe different weights. Four of six arms have a best epoch before 9; two have best == final, so an annotation claiming a gap everywhere would itself be wrong. Each report now carries a `legacy` block with saved_val_total derived from that file's OWN val.history, an explicit statement that best_val_* describe discarded weights, and the clauses it fails under schema 2. ANNOTATED, NOT REWRITTEN — 84 insertions, 0 deletions, no measured value touched. (Medium) grad_clip == 0 was silently corrupting gradients. I wrote clip_grad_norm_(params, inf) believing it a no-op; PyTorch scales by max_norm/(total_norm+1e-6), which is nan when both are infinite, and multiplies it through. Verified on the pinned torch 2.7.1: [inf,1,2] came back [nan,nan,nan] while the returned norm still read inf — detection signal fine, optimiser input not. gradient_total_norm now uses get_total_norm when clipping is off. (Medium) TWO REGRESSION TESTS WRITTEN FOR THE PREVIOUS ROUND WERE STILL VACUOUS. The step-count test reimplemented the loop instead of running it — deleting scheduler.step() from the flush left it green. The NaN test edited a report fixture — deleting the counter increment left it green. Both are now PRODUCER tests calling train_stage2 on CPU with odd micro-batches and accumulation 2, spying on the real optimizer/scheduler steps and injecting a real inf into a live gradient. Both of CodeRabbit's counterexamples now fail, verified by source mutation. Seven vacuous tests across this step; the pattern is that I keep testing the predicate and calling it a test of the producer. Three assertions I wrote this round were also wrong, each caught by running them: that the forward loss stays finite under a poisoned gradient (stepping on inf poisons the weights); that n_nonfinite_grad_steps >= n_nonfinite_steps (different denominators — micro-batches vs optimiser steps); and a sabotage that removed one of the counter's two increment sites. Confirmed by the reviewer, not assumed: backward() returns only after every DDP bucket all-reduce completes, including with find_unused_parameters=True, so the norm is a valid post-reduction signal. Validation: trainer 121, sizing 20, overrides 22, backfill 4; ruff + black clean. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
…erting it r5 via the CLI, scoped with --base-commit c6fc074 to the two unreviewed commits, and via the CLI rather than a PR comment so it consumed no GitHub review quota. ONE finding, minor. Series across the step: 11 -> 12 -> 3 -> 3 -> 1. The finding is a self-inflicted inconsistency: the legacy annotation hardcoded best_val_total_is_NOT_the_saved_weights: True on all six arms, while aux1.0_lr1e-4 and aux1.0_lr3e-4 have best epoch == 9 == saved. The test written in the SAME commit asserts that four of six differ — so the annotation contradicted its own test file. Now derived from the epoch comparison, with two note wordings: one for arms where best_val_total describes discarded weights, one for arms where the two agree and neither is misleading (and why that does not make the distinction moot). Still annotation-only: 84 insertions, 0 deletions against the pre-annotation commit. The test asserts the flag equals the epoch comparison and that the note matches its case. Verified non-tautological by sabotage: restoring the hardcoded True turns it red. Four consecutive rounds found real defects in the previous round's fixes, so none was a safe stopping point. This one found a documentation-consistency nit in an annotation block — no behaviour, no gate, no artifact. Validation: trainer 121, sizing 20, overrides 22, backfill 4; ruff + black clean. Claude-Session: https://claude.ai/code/session_01N8NJ1ksgJeJzV5Q1JHGkt1
This was referenced Aug 4, 2026
P3-08: the with/without-aux ablation on the calibrated binary head + the Stage-2 score producer
#101
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements imp.md step P3-06 and lands the run: six Stage-2 LoRA checkpoints from job 1064.
What this adds
src/tbox_finder/stage2/train.py— the Hydra/torchrun entrypoint: eligibility, dataset + collator, DDP loop, self-validating report, checkpoint save.src/tbox_finder/train/ddp.py— rank helpers, thePYTHONHASHSEEDverifier andShardedSampler, promoted out oftrain_stage1.py(which now imports them) because the two trainers run in different conda envs (ADR-0002 A4) and a forked copy would fix one while shipping the bug in the other.src/tbox_finder/stage2/sizing.py+slurm/p3/stage2_sizing_smoke.sbatch— the sizing measurement imp.md asked for.conf/train/stage2.yaml,conf/optim/stage2.yaml,slurm/p3/stage2_lora_finetune.sbatch.scripts/backfill_stage2_provenance.py— one-shot repair for the sidecars job 1064 could not write.tests/ml/test_stage2_train_smoke.py,tests/ml/test_stage2_sizing.py,tests/unit/test_backfill_provenance.py.The §7 decision this step owned
imp.mdhands the Stage-2 train-eligibility policy here. Settled by user sign-off (2026-08-03) and encoded inrow_eligibility:D5 because PRD §12 grades leave-one-order-out on the two-stage system; the scheme-A intersection so P3-07 fits its temperature on the disjoint
calibsplit and P3-10 grades GATE-2 on this checkpoint rather than an ADR-0004 A6 eval twin. Measured: 9,059 train (5,199 corpus + 3,860 parentless decoy; 4,795 pos / 4,264 neg), 2,070 val, row-id overlap 0, partition exact.Result (job 1064)
Six points, six passing gates, six 48 MB checkpoints, zero real failures.
dvc add→stage2_rinalmo.dvc(313,437,285 B, 30 files);dvc pushdeferred to the phase-exit gate.⚠ These totals are not comparable across arms. At
aux_weight=0the objective is the binary term alone; at 0.5 it is that plus five weighted aux terms. Arm comparison is P3-08's job, on the calibrated binary head with matched metrics.Reviewer: the things I'd most want a second pair of eyes on
derive_clausesintrain.py— the gate is scoped to rank 0 because the evidence (checkpoint, git snapshot) is rank 0's. Job 1053 died at the finish line when every rank judged it. Is the scoping right, and does anything else in that function read rank-local state?row_eligibilityis fail-closed and refuses by name; the parentless-decoy admit branch raises if a row with a real fold assignment reaches it. Is there a row shape that slips through?backgroundper the Stage-1 convention (data/negatives.py), not ignored. Fail-closed tosource == decoy AND is_tbox is False.scripts/backfill_stage2_provenance.py— deliberately forkscheckpoint_output_filesbecause it must run at the run's commit, which predates the shared helper.tests/unit/test_backfill_provenance.pypins the two together. Is that fork acceptable, and is the git_sha guard sound?.OKmarkers were never written (the sidecar crash preceded them), so §9.3 verification rests on reports + checkpoints +.errscan; and the provenance sidecars are reconstructed, each labelledreconstructed: truewith its limitation.Not in scope
dvc push(phase-exit gate), and hand-wiringtorch.utils.checkpoint— gradient checkpointing is a measured no-op on this backbone (saving ratio 0.9986;modeling_rinalmo.pyadvertises support at :82, stores the flag at :118, never calls_gradient_checkpointing_funcat :704), so it shipsfalsewith a drift-guard test.Summary by CodeRabbit
New Features
Documentation
Tests