diff --git a/src/baskerville/scripts/hound_ism_bed.py b/src/baskerville/scripts/hound_ism_bed.py index 9e3db62..3116c42 100755 --- a/src/baskerville/scripts/hound_ism_bed.py +++ b/src/baskerville/scripts/hound_ism_bed.py @@ -88,8 +88,15 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + parser.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) parser.add_argument( "-o", "--out_dir", @@ -151,12 +158,7 @@ def main(): args = parser.parse_args() # parse options - if args.mix_dtype == "float16": - args.mix_dtype = torch.float16 - elif args.mix_dtype == "bfloat16": - args.mix_dtype = torch.bfloat16 - elif args.mix_dtype == "float32": - args.mix_dtype = torch.float32 + args.mix_dtype = snps.parse_mix_dtype(args.mix_dtype) args.shifts = [int(shift) for shift in args.shifts.split(",")] args.snp_stats = [snp_stat for snp_stat in args.snp_stats.split(",")] @@ -219,6 +221,8 @@ def main(): seqnn_model.ensemble_shifts = args.shifts seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() + if args.compile: + seqnn_model.compile() output_length = seqnn_model.output_length() output_stride = seqnn_model.output_stride() diff --git a/src/baskerville/scripts/hound_ism_bed_folds.py b/src/baskerville/scripts/hound_ism_bed_folds.py index 19b20a3..43caab0 100755 --- a/src/baskerville/scripts/hound_ism_bed_folds.py +++ b/src/baskerville/scripts/hound_ism_bed_folds.py @@ -102,8 +102,15 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + ism_group.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) ism_group.add_argument( "-o", dest="out_dir", @@ -540,6 +547,9 @@ def build_ism_bed_cmd(args, model_file, out_dir): if args.mix_dtype != "float32": cmd_parts.extend(["-m", args.mix_dtype]) + if args.compile: + cmd_parts.append("--compile") + cmd_parts.extend(["-o", out_dir]) if args.rc: diff --git a/src/baskerville/scripts/hound_ism_snp.py b/src/baskerville/scripts/hound_ism_snp.py index f8b1982..08d49c3 100755 --- a/src/baskerville/scripts/hound_ism_snp.py +++ b/src/baskerville/scripts/hound_ism_snp.py @@ -87,8 +87,15 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + parser.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) parser.add_argument( "-o", dest="out_dir", @@ -138,12 +145,7 @@ def main(): args = parser.parse_args() # parse options - if args.mix_dtype == "float16": - args.mix_dtype = torch.float16 - elif args.mix_dtype == "bfloat16": - args.mix_dtype = torch.bfloat16 - elif args.mix_dtype == "float32": - args.mix_dtype = torch.float32 + args.mix_dtype = snps.parse_mix_dtype(args.mix_dtype) args.shifts = [int(shift) for shift in args.shifts.split(",")] args.snp_stats = args.snp_stats.split(",") @@ -203,6 +205,8 @@ def main(): seqnn_model.ensemble_shifts = args.shifts seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() + if args.compile: + seqnn_model.compile() output_length = seqnn_model.output_length() output_stride = seqnn_model.output_stride() diff --git a/src/baskerville/scripts/hound_ism_snp_folds.py b/src/baskerville/scripts/hound_ism_snp_folds.py index bbd058d..972a2ae 100755 --- a/src/baskerville/scripts/hound_ism_snp_folds.py +++ b/src/baskerville/scripts/hound_ism_snp_folds.py @@ -99,8 +99,15 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + ism_group.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) ism_group.add_argument( "-o", dest="out_dir", @@ -538,6 +545,9 @@ def build_ism_snp_cmd(args, model_file, out_dir): if args.mix_dtype != "float32": cmd_parts.extend(["-m", args.mix_dtype]) + if args.compile: + cmd_parts.append("--compile") + cmd_parts.extend(["-o", out_dir]) if args.rc: diff --git a/src/baskerville/scripts/hound_snp.py b/src/baskerville/scripts/hound_snp.py index fa6bf8f..907c450 100755 --- a/src/baskerville/scripts/hound_snp.py +++ b/src/baskerville/scripts/hound_snp.py @@ -17,9 +17,8 @@ import os import tempfile import shutil -import torch -from baskerville.snps import score_snps, score_gene_snps +from baskerville.snps import parse_mix_dtype, score_snps, score_gene_snps from baskerville.helpers.gcs_utils import ( upload_folder_gcs, download_rename_inputs, @@ -104,8 +103,15 @@ def main(): "--mix_dtype", dest="mix_dtype", default="float32", + choices=["float32", "bfloat16", "float16"], help="Mixed precision dtype", ) + parser.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) parser.add_argument( "-n", "--norm", @@ -207,12 +213,7 @@ def main(): args.targets_file = download_rename_inputs(args.targets_file, temp_dir) # parse options - if args.mix_dtype == "float16": - args.mix_dtype = torch.float16 - elif args.mix_dtype == "bfloat16": - args.mix_dtype = torch.bfloat16 - elif args.mix_dtype == "float32": - args.mix_dtype = torch.float32 + args.mix_dtype = parse_mix_dtype(args.mix_dtype) args.shifts = [int(shift) for shift in args.shifts.split(",")] args.snp_stats = args.snp_stats.split(",") diff --git a/src/baskerville/scripts/hound_snp_folds.py b/src/baskerville/scripts/hound_snp_folds.py index e8a4039..80dcfca 100755 --- a/src/baskerville/scripts/hound_snp_folds.py +++ b/src/baskerville/scripts/hound_snp_folds.py @@ -108,6 +108,20 @@ def main(): type=int, help="Local window size in bp for targets with window=local", ) + snp_group.add_argument( + "-m", + "--mix_dtype", + dest="mix_dtype", + default="float32", + choices=["float32", "bfloat16", "float16"], + help="Mixed precision dtype", + ) + snp_group.add_argument( + "--compile", + default=False, + action="store_true", + help="Compile the model with torch.compile", + ) snp_group.add_argument( "-n", "--norm", @@ -650,6 +664,12 @@ def build_snp_cmd(args, model_file, start_i, end_i, out_dir=None, fold_cross=Non if hasattr(args, "local_window") and args.local_window != 2048: cmd_parts.extend(["--local_window", str(args.local_window)]) + if hasattr(args, "mix_dtype") and args.mix_dtype != "float32": + cmd_parts.extend(["-m", args.mix_dtype]) + + if hasattr(args, "compile") and args.compile: + cmd_parts.append("--compile") + # Add normalization file if specified if hasattr(args, "norm_subdir") and args.norm_subdir and fold_cross is not None: norm_file = f"{args.models_dir}/{fold_cross}/{args.norm_subdir}/scores.h5" diff --git a/src/baskerville/seqnn.py b/src/baskerville/seqnn.py index f151e7f..bbbc194 100644 --- a/src/baskerville/seqnn.py +++ b/src/baskerville/seqnn.py @@ -15,6 +15,18 @@ from baskerville.types import BatchData, ModelOutput +def _float_output(out: ModelOutput, keep_gradients: bool) -> ModelOutput: + """Cast predictions to float32 (autocast may leave them reduced), detaching unless kept.""" + + def cast(x): + if x is None: + return None + x = x.float() + return x if keep_gradients else x.detach() + + return ModelOutput(coverage=cast(out.coverage), gene=cast(out.gene)) + + class SeqNN: """Sequence neural network model. @@ -142,12 +154,7 @@ def __call__( gene_out_mask=gene_out_mask, gene_presence=gene_presence, ) - if not keep_gradients: - out = ModelOutput( - coverage=out.coverage.detach() if out.has_coverage else None, - gene=out.gene.detach() if out.has_gene else None, - ) - outputs.append(out) + outputs.append(_float_output(out, keep_gradients)) # reverse complement if self.ensemble_rc: @@ -166,22 +173,17 @@ def __call__( gene_out_mask=rc_gene_out_mask, gene_presence=gene_presence, ) + out_rc = _float_output(out_rc, keep_gradients) # Process coverage: flip and apply strand pairing coverage_rc = None if out_rc.has_coverage: coverage_rc = torch.flip(out_rc.coverage, [2]) - if not keep_gradients: - coverage_rc = coverage_rc.detach() if strand_pair_hi is not None: coverage_rc = coverage_rc[:, strand_pair_hi, :] # Gene predictions: slices were transformed above, no post-hoc flipping needed - gene_rc = None - if out_rc.has_gene: - gene_rc = out_rc.gene - if not keep_gradients: - gene_rc = gene_rc.detach() + gene_rc = out_rc.gene outputs.append(ModelOutput(coverage=coverage_rc, gene=gene_rc)) @@ -200,10 +202,12 @@ def __call__( return ModelOutput(coverage=coverage_avg, gene=gene_avg) def compile(self): - """Compile model for faster inference. - (Not working well with augmentation ensembling.) + """Compile model for faster inference; call after restore. + + Default mode: reduce-overhead's CUDA graphs overwrite outputs that + callers hold across calls (e.g. ref predictions). """ - if torch.cuda.get_device_capability()[0] < 7: + if self.device == "cuda" and torch.cuda.get_device_capability()[0] < 7: print("Warning: CUDA device capability < 7.0, skipping compilation.") else: self.model = torch.compile(self.model) @@ -805,7 +809,10 @@ def __init__( self.heads_gene_def = heads_gene_def self.output_stride = 1 self.output_crop_bp = 0 - self.output_slice = output_slice + # index tensor (not e.g. a pandas Index) so torch.compile can trace the slice + if output_slice is not None: + output_slice = torch.as_tensor(np.asarray(output_slice), dtype=torch.long) + self.register_buffer("output_slice", output_slice, persistent=False) global_vars = global_vars or {} self.global_vars = global_vars self.seq_length = seq_length diff --git a/src/baskerville/snps.py b/src/baskerville/snps.py index a96da3a..635f021 100644 --- a/src/baskerville/snps.py +++ b/src/baskerville/snps.py @@ -1,6 +1,8 @@ import concurrent import json import os +import sys + import h5py import numpy as np import pandas as pd @@ -17,6 +19,24 @@ GENE_ASSAYS = {"RNA", "RNA3", "CAGE"} +MIX_DTYPES = { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, +} + + +def parse_mix_dtype(name: str) -> torch.dtype: + """Map a --mix_dtype name to a torch dtype, warning on reduced precision.""" + if name != "float32": + print( + f"WARNING: --mix_dtype {name} adds rounding noise that can dominate " + "variant effect scores (small alt - ref differences). Use float32 " + "unless validated for your model and stats.", + file=sys.stderr, + ) + return MIX_DTYPES[name] + def gene_track_mask(targets_df) -> np.ndarray: """Boolean mask over targets selecting gene-appropriate tracks. @@ -164,8 +184,11 @@ def score_snps(args): seqnn_model = seqnn.SeqNN(params_model, output_slice=targets_df.index) seqnn_model.restore(args.model_file) seqnn_model.ensemble_rc = args.rc + seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() seqnn_model.model_di = params.get("train", {}).get("model_di", None) + if args.compile: + seqnn_model.compile() # shift outside seqnn num_shifts = len(args.shifts) @@ -608,8 +631,11 @@ def score_gene_snps(args): seqnn_model = seqnn.SeqNN(params_model, output_slice=targets_df.index) seqnn_model.restore(args.model_file) seqnn_model.ensemble_rc = args.rc + seqnn_model.mix_dtype = args.mix_dtype seqnn_model.model.eval() seqnn_model.model_di = params.get("train", {}).get("model_di", None) + if args.compile: + seqnn_model.compile() # shift outside seqnn num_shifts = len(args.shifts) diff --git a/src/tests/test_ism_snp.py b/src/tests/test_ism_snp.py index de318c1..17654dc 100644 --- a/src/tests/test_ism_snp.py +++ b/src/tests/test_ism_snp.py @@ -299,3 +299,32 @@ def test_ism_snp_gene_track_selection(static_model_dir, stat): continue out_df = pd.read_csv(f"{temp_dir}/{name}", sep="\t", index_col=0) assert list(out_df.identifier) == expected_targets + + +def test_ism_snp_bfloat16(static_model_dir): + """-m bfloat16 runs (predictions cast to float32 for scoring) and warns.""" + test_dir = str(pathlib.Path(__file__).parent) + with tempfile.TemporaryDirectory() as temp_dir: + cmd = [ + "python", + "-m", + "baskerville.scripts.hound_ism_snp", + "-f", + f"{test_dir}/data/sc3.fa.gz", + "-t", + f"{test_dir}/data/targets_sc3_me.txt", + "-o", + temp_dir, + "-l", + "4", + "--stats", + "logSUM", + "-m", + "bfloat16", + f"{test_dir}/data/params_sc3.json", + f"{static_model_dir}/sc3_model.pth", + f"{test_dir}/data/sc3_snps.vcf", + ] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, f"hound_ism_snp failed: {result.stderr}" + assert "WARNING: --mix_dtype bfloat16" in result.stderr diff --git a/src/tests/test_snp.py b/src/tests/test_snp.py index 2dd5b95..f54c303 100644 --- a/src/tests/test_snp.py +++ b/src/tests/test_snp.py @@ -5,11 +5,15 @@ import pytest import subprocess +import argparse + import h5py import numpy as np import pandas as pd +import torch -from baskerville.snps import partition_snp_stats +from baskerville.scripts.hound_snp_folds import build_snp_cmd +from baskerville.snps import parse_mix_dtype, partition_snp_stats @pytest.fixture @@ -714,3 +718,90 @@ def test_partition_snp_stats(stats, expected): original = stats.copy() assert partition_snp_stats(stats) == expected assert stats == original + + +def test_parse_mix_dtype_warns_below_float32(capsys): + assert parse_mix_dtype("float32") is torch.float32 + assert "WARNING" not in capsys.readouterr().err + assert parse_mix_dtype("bfloat16") is torch.bfloat16 + assert "WARNING" in capsys.readouterr().err + + +def test_snp_folds_passes_mix_dtype(): + args = argparse.Namespace( + mix_dtype="bfloat16", params_file="p.json", vcf_file="v.vcf" + ) + assert "-m bfloat16" in build_snp_cmd(args, "m.pth", 0, 1) + args.mix_dtype = "float32" + assert "-m" not in build_snp_cmd(args, "m.pth", 0, 1).split() + + +def test_snp_folds_passes_compile(): + args = argparse.Namespace(compile=True, params_file="p.json", vcf_file="v.vcf") + assert "--compile" in build_snp_cmd(args, "m.pth", 0, 1).split() + + +def test_snp_mix_dtype_takes_effect(static_model_dir, test_vcf_file): + """-m bfloat16 must change the model's precision (it was once a silent no-op).""" + test_dir = str(pathlib.Path(__file__).parent) + scores = {} + with tempfile.TemporaryDirectory() as temp_dir: + for mix_dtype in ["float32", "bfloat16"]: + out_dir = f"{temp_dir}/{mix_dtype}" + cmd = [ + "python", + "-m", + "baskerville.scripts.hound_snp", + "-f", + f"{test_dir}/data/sc3.fa.gz", + "-t", + f"{test_dir}/data/targets_sc3_me.txt", + "-o", + out_dir, + "--stats", + "logSUM", + "-m", + mix_dtype, + f"{test_dir}/data/params_sc3.json", + f"{static_model_dir}/sc3_model.pth", + test_vcf_file, + ] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, f"hound_snp failed: {result.stderr}" + assert ("WARNING: --mix_dtype" in result.stderr) == (mix_dtype != "float32") + with h5py.File(f"{out_dir}/scores.h5", "r") as h5: + scores[mix_dtype] = h5["cov/logSUM"][:] + assert not np.array_equal(scores["float32"], scores["bfloat16"]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="compile targets GPU") +def test_snp_compile_matches_eager(static_model_dir, test_vcf_file): + """--compile runs and reproduces eager scores up to fusion reassociation.""" + test_dir = str(pathlib.Path(__file__).parent) + scores = {} + with tempfile.TemporaryDirectory() as temp_dir: + for compile_flag in [[], ["--compile"]]: + out_dir = f"{temp_dir}/{len(compile_flag)}" + cmd = [ + "python", + "-m", + "baskerville.scripts.hound_snp", + "-f", + f"{test_dir}/data/sc3.fa.gz", + "-t", + f"{test_dir}/data/targets_sc3_me.txt", + "-o", + out_dir, + "--stats", + "logSUM", + "--rc", + *compile_flag, + f"{test_dir}/data/params_sc3.json", + f"{static_model_dir}/sc3_model.pth", + test_vcf_file, + ] + result = subprocess.run(cmd, capture_output=True, text=True) + assert result.returncode == 0, f"hound_snp failed: {result.stderr}" + with h5py.File(f"{out_dir}/scores.h5", "r") as h5: + scores[len(compile_flag)] = h5["cov/logSUM"][:].astype("float32") + np.testing.assert_allclose(scores[1], scores[0], rtol=1e-2, atol=1e-2)