HuggingFaceBio

/

Carbon-A-1.2B

Like 26

Follow

Hugging Face Biology Research 356

Hi,
Not a bug, but a suggestion.

Summary

annotation_pipeline/src/model_inference.py currently aborts on machines without a CUDA GPU:

available_gpus = torch.cuda.device_count()
if available_gpus <= 0:
    raise RuntimeError("No CUDA devices available.")

However, setup_model_for_gpu() already contains a CPU device fallback (gpu_id < 0 -> "cpu"), and every operator used by the model (SDPA attention, RMSNorm, GEMM) has a native CPU implementation. On recent PyTorch versions the fused flash_attention_for_cpu SDPA kernel also keeps 16,384-token windows within a few GB of RAM, so long-window inference works fine on CPU.

We therefore propose a minimal, opt-in CPU path: --gpu_count 0 runs inference on CPU in float32, instead of raising. This makes the pipeline usable on CPU-only HPC clusters, which are common in genomics groups.

Proposed change (unified diff)

--- a/annotation_pipeline/src/model_inference.py
+++ b/annotation_pipeline/src/model_inference.py
@@
-    device = torch.device(f"cuda:{gpu_id}" if gpu_id >= 0 else "cpu")
+    device = torch.device(
+        f"cuda:{gpu_id}" if (gpu_id >= 0 and torch.cuda.is_available()) else "cpu"
+    )
@@
-    available_gpus = torch.cuda.device_count()
-    if available_gpus <= 0:
-        raise RuntimeError("No CUDA devices available.")
-    if args.gpu_count == -1:
-        gpu_count = available_gpus
-    elif args.gpu_count <= 0:
-        raise ValueError(
-            f"--gpu_count must be -1 (all GPUs) or a positive integer, got {args.gpu_count}"
-        )
-    else:
-        gpu_count = min(args.gpu_count, available_gpus)
-    print(f"🎯 Using {gpu_count} GPU(s) out of {available_gpus} available")
+    available_gpus = torch.cuda.device_count()
+    if args.gpu_count == 0:
+        # Opt-in CPU fallback: one worker process, float32 weights.
+        gpu_count = 1
+        dtype_str = "float32"
+        print("INFO: --gpu_count 0: running model inference on CPU with float32 weights (slow)")
+    else:
+        if available_gpus <= 0:
+            raise RuntimeError("No CUDA devices available.")
+        if args.gpu_count == -1:
+            gpu_count = available_gpus
+        elif args.gpu_count < 0:
+            raise ValueError(
+                f"--gpu_count must be -1 (all GPUs) or a positive integer, got {args.gpu_count}"
+            )
+        else:
+            gpu_count = min(args.gpu_count, available_gpus)
+        print(f"Using {gpu_count} GPU(s) out of {available_gpus} available")

Notes for reviewers:

• Default behavior is unchanged; the CPU path is only entered when the user explicitly passes --gpu_count 0.
• dtype_str stays bfloat16 on GPUs. On CPU we deliberately use float32: many CPUs (e.g. Skylake-SP, Zen2) lack native bf16 instructions, where bf16 would be slower and less accurate.
• CPU workers should run with --attn-implementation sdpa (no flash-attn package needed). The fused flash_attention_for_cpu kernel is selected automatically by PyTorch; we measured no quadratic-memory blow-up at 16k tokens.
• The transition/confidence networks and the Viterbi decoding already work on CPU (fp32 for the two networks; the numba reference backend needs no C++ compiler), so no other file needs modification.

Usage

export OMP_NUM_THREADS=<physical_cores>
bash Carbon-A-1.2B/annotation_pipeline/run_annotation.sh \
    --input genome.fa --gpu-count 0 --attn-implementation sdpa --batch-size 4

Validation performed

• End-to-end correctness (CPU, fp32): a 500 kb apricot fragment yielded 34 genes, 33 with confidence ≥ 0.1 (median 0.859); a 5 Mb fragment yielded 508 genes, 485 ≥ 0.1 (median 0.778). GFF structure (exon phases, start/stop codons) and the codon-table check passed. No _WARNING was raised.
• Memory: ~6.6 GB RSS at batch 1, ~10 GB at batch 4 (fp32).
• Throughput (per 16,384-token window, fp32, batch 4, 10 threads):
  • Intel Skylake 2.4 GHz: ~110 s (attention 67% of runtime, MKL GEMM 27%)
  • AMD Rome 2.9 GHz (Zen2): ~175 s
• The resulting per-window times make whole-genome annotation practical on multi-node CPU clusters.

Optional documentation addition

The model card / pipeline README could mention the CPU mode explicitly, e.g.:

On machines without a CUDA GPU, pass --gpu-count 0 --attn-implementation sdpa. Inference then runs on CPU in float32 (slower but fully supported); post-processing already runs on CPU.

Best,
Kun