Repository navigation
Make --mix_dtype work in SNP/ISM scoring; add --compile - #3
Merged
Merged
Conversation
A pandas Index output_slice (e.g. targets_df.index from hound_snp) made torch.compile skip SeqNNMod.forward entirely (graph break inside the head loop). As a non-persistent long-tensor buffer it traces into a single graph, follows .to(device), and leaves the state_dict unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- hound_snp never set SeqNN.mix_dtype, so SeqNN.__call__'s inner float32 autocast disabled the requested bf16/fp16: -m was a no-op. - SeqNN.__call__ now returns float32 predictions; reduced-precision outputs crashed scoring (numpy has no bfloat16) and summed imprecisely. - hound_snp_folds gains -m and passes it through. - -m is restricted to float32/bfloat16/float16 in all scripts. - snps.parse_mix_dtype warns that reduced precision can dominate variant effect scores (small alt - ref differences). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
hound_snp, hound_ism_snp, hound_ism_bed and their _folds wrappers gain --compile, which calls SeqNN.compile() after restore. Default mode only: reduce-overhead's CUDA graphs overwrite outputs held across calls (ref predictions). SeqNN.compile only queries CUDA capability on CUDA devices. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davek44
force-pushed
the
seqnn-output-slice-tensor
branch
from
October 5, 2026 01:30
2c0d58f to
99a7f9f
Compare
davek44
changed the title
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
1 participant
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.
Summary
-m/--mix_dtypenow takes effect inhound_snp. It never setSeqNN.mix_dtype, soSeqNN.__call__'s own float32 autocast turned the requested bf16/fp16 back off (output bitwise identical to fp32).SeqNN.__call__returns float32 predictions whatever the autocast dtype. Reduced-precision outputs crashed scoring on CPU (numpy has no bfloat16;hound_ism_snp -m bfloat16always failed this way) and would sum imprecisely.snps.parse_mix_dtype, shared byhound_snp/hound_ism_snp/hound_ism_bed). On Cerberus, bf16 variant scores correlate with fp32 at only r ≈ 0.57–0.78 (random-sequence test), because rounding noise swamps small alt − ref differences. A GTEx eQTL comparison is pending.-mrestricted tofloat32|bfloat16|float16;hound_snp_foldsgains-mand passes it through.--compileforhound_snp,hound_ism_snp,hound_ism_bedand their_foldswrappers:SeqNN.compile()after restore, default mode. 1.33× steady-state per SNP for Cerberus at fp32 on an RTX 4090 (scores r ≥ 0.990 vs eager).reduce-overheadgives no further gain and breakshound_snp(CUDA graphs overwrite held ref outputs), so it isn't offered.SeqNNMod.output_slicestored as an index-tensor buffer. A pandasIndexmade dynamo skipSeqNNMod.forward; it now traces as one graph. Non-persistent, so checkpoints are unchanged.Companion: calico/westminster
966ff91adds-m/--compiletowestminster_eqtl_folds.Test plan
test_snp,test_ism_snp,test_ism_bed,test_snp_folds,test_modelon CPU (48 passed, 1 GPU-only skip)test_snp -k "compile or mix_dtype"(5 passed), including compiled vs eager agreement-m bfloat16changes scores and warns; folds pass-m/--compile; ISM runs in bf16🤖 Generated with Claude Code