float16

[AlexanderMath]

Question 1. Did anyone get Orb running in float16 or bfloat16?

Question 2. If I take float32 weights and train a bit in float16, what would convince us model didn't break? (e.g., would low validation MAE on MPTraj be sufficient?)

Apologies if I'm missing anything.

(update) Sorry for wall of text below, added details to make others didn't have to reproduce. If useful happy to write PR (with non-hacky fix for brute_force_kNN).

TLDR: bf16 is 2-3x faster on rtx4090. To get xs model working I had to hack brute_force_kNN. This gave 6.5steps/s = 0.28ns/day for a 36k atom system. I expect around 0.6-0.8ns/day on h100. Big thing left is to evaluate performance drop, e.g., reproduce Figure 2 in bf16.

[AlexanderMath]

notes on timing orb+rtx4090+crispr (36k atoms incl H+water)

• orb_d3_v2 got 1it/s traj 100% nvidia-smi util [1]
• orb_d3_v2 got 2.27it/s with bf16+torch.compile traj ~80% util [2]
• orb_d3_xs_v2 got 4.1it/s with bf16+torch.compile traj ~50% util [3]

surprises:

• bf16 and torch.compile only improved when combined
• can get >95% bf16 improvement only from self._edge_mlp (may be specific to input)
• orb_d3_xs_v2 only reduce message_passing=15->5 but drop GPU util 100%->50% as nvidia-smi
• fixing orb_d3_xs_v2 GPU util could give 8it/s or 0.34ns/day for CRISPR on cheap $0.2/h rtx4090 (or $14/ns).

todo
[ ] eval on PDB without missing residues. compute MAE(force_orb_d3_v2(frame_i), force_orb_d3_xs_v2_bf16(frame_i)) for frames from orb_d3_v2 trajectory

[1]

[2]

[3]

[AlexanderMath]

notes

• low util for orb_d3_xs_v2 was caused by brute_force_kNN; quick hack got me to 6.7it/s and 80% util for non-pbc

[AlexanderMath]

the hack.

step = 0
top1024 = None
def brute_force_knn(positions: torch.Tensor, k: int) -> Tuple[torch.Tensor, torch.Tensor]:
    global step, top1024
    if step % 20 == 0: 
        dist = torch.cdist(positions, positions)
        vals, indxs = torch.topk(dist, k, largest=False, sorted=True)
        _ , top1024 = torch.topk(dist, 1024, largest=False, sorted=True)
    else: 
        new_dists = torch.linalg.norm(positions.view(-1, 1, 3) - positions[top1024] , dim=2)
        vals, new_indxs = torch.topk(new_dists, k, largest=False, sorted=True)
        indxs = top1024.gather(1, new_indxs)
    step += 1
    return vals, indxs 

warning: In double precision the hack passes th.allclose( old_algo, new_algo ) but it breaks for

• f32 bcs cdist disagrees with th.linalg.norm(a.view(-1,1,3)-b.view(1,-1,3)
• atoms moving from outside top 1024 to within top k; in my case this is ~5A within 10fs and guesstimate atoms move around <0.5A.

notes

• torch.cdist takes 80ms ~40% (GNN takes 125ms)
• cdist in bf16 takes 40ms
• above hack is 14ms so still ~12%.

[AlexanderMath]

At 6.5step/s = 560k steps/day. Not sure there's anything left to do wrt model layer inferences. Main thing left is 20ms (~15%) of GPU not doing anything after forward pass (presumably ASE doing it's thing on CPU) -- might take a look, but think this is sufficiently fast if i switch to h100.

Next
[ ] try timing H100; best case we get 3x speedup and reach 1.5m steps/day.
[ ] redo figure 2 in bf16
[ ] train a bit in bf16 and see how much perf we can recover for figure 2 case.

[benrhodes26]

Hi Alexander,

Sorry for the delay in responding; as you can see, we've been busy working on major updates and have just released a new version (0.5.0) with a suite of new models (orb-v3-*).

In particular, orb-v3-direct-20-omat is the most similar to orb-v2, but substantially faster, and more performant in our benchmarking.

Speed-wise, the main changes include: compilation, shallower but wider models (5 layers instead of 15, 1024 width in MLPs rather than 512) and better graph-construction methods (we use cuML KNN routines for large systems).

I would be curious to hear how these changes impact your use-case.

We would also be interested in a PR contribution from you for bf16 --- but please benchmark it against our default precision=float32-high, which calls torch.set_float32_matmul_precision("high").

Your hack is also well-motivated, and we're interested in implementing something like that. However, we may take a bit more time to evaluate it in detail, as we don't want to risk implementing something that is flaky and confusing for users.

[AlexanderMath]

Christmas arrived early -- amazing results on accuracy/speed trade-off 🎉

I would be curious to hear how these changes impact your use-case.

# v2
orb_d3_v2 # 1it/s
orb_d3_v2  # 2.27it/s with bf16+torch.compile
orb_d3_xs_v2 # 6.5it/s bf16+compile+kNN hack

# v3
orb-v3-direct-20 # 2.25it/s 
orb-v3-direct-20 # 3.7it/s bf16(gnn_stack) + kNN with cuml=12.2

Note: gnn_stack improved from 0.382m to 0.229ms when switching f32 to bf16. I was expecting 3x reduction ~126ms, will debug w/ nsys profile when i get back to desktop (nsys-ui doesn't work over ssh...).

Note: rtx4090 is 6x cheaper than a100 (~$0.2/h vs $1.2/h) and has 1.05x more TFLOPs (330 tflops [1] vs 312tflops [2]).

We would also be interested in a PR contribution from you for bf16 --- but please benchmark it against our default precision=float32-high, which calls torch.set_float32_matmul_precision("high").

Got it. LMK if you have preference for test-set.

Your hack is also well-motivated, and we're interested in implementing something like that. However, we may take a bit more time to evaluate it in detail, as we don't want to risk implementing something that is flaky and confusing for users.

Yeh agree, not sure this would ever be good default behaviour.

[1] https://images.nvidia.com/aem-dam/Solutions/geforce/ada/nvidia-ada-gpu-architecture.pdf page 30
[2] https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a100/pdf/nvidia-a100-datasheet-us-nvidia-1758950-r4-web.pdf

[benrhodes26]

Got it. LMK if you have preference for test-set.

We just added internal/speed.py (v0.5.1) which can be run like so:

python internal/speed.py --method orb-forward --extra_kwargs "name=orb_v3_direct_20_omat precision=float32-high compile=true"

This should report speed and memory use across a range of system-sizes (up to 100k atoms).

[beckernick]

Hi! Came across this issue due to the cuML reference. I work on accelerated data science at NVIDIA.

I'm not familiar with this project, but it looks very cool.

I wanted to mention that we've done a lot of work optimizing approximate nearest neighbors algorithms recently. If exact KNN isn't critical, these can provide a major speedup.

We recently migrated cuML's GPU-accelerated UMAP implementation from exact KNN to ANN (using the nn-descent algorithm, which is what the CPU version of UMAP uses). This enabled huge performance and scalability benefits, without materially affecting UMAP quality.

This is currently available in C++ and Python via cuVS, but is something we could explore making available e.g., via cuML's NearestNeighbors class, too.

If the KNN is a non-trivial amount of time, could be worth exploring. Happy to chat further, if interested.

[benrhodes26]

Very cool @beckernick! Before using cuML, nearest-neighbors was a significant bottleneck for us.

It is now a 'moderate' bottleneck i.e. if knn was free, our total throughput would be between 1.2 to 2x for the types of data we work with right now.

Up to 2x would be great, but using approximate instead of exact KNN could actually be a problem for our use-case. The models we train are fairly sensitive to neighbor lists.

I could definitely see a benefit in the future however. If our atomic systems were substantially larger, or our models forward pass got even cheaper, then KNN would again become a major bottleneck and we be very interested in speeding it up.

[AlexanderMath]

Update: Looks like I got ~10x cost reduction by switching from H100 (~500 TF32 @ $3/hr → $72/ns) to RTX 4090 (330 TF16 @ $0.2/hr → ~$5/ns).

Hard part: pytorch by defaults accumulates A.to(th.float16)@B.to(th.float16) in float32. Fixed by rewriting edge_mlp matmul to Triton (also fused in Silu).

Surprise: MAE(orb(crispr), orb_f16(crispr)) ~ 0.007 seems ok, will test more soon, and finetune if find issue. i just keep featurizer/decoder in f32 and appears ok.

Not ready for PR, will require much more testing.

[thomasloux]

Did you have the chance to look back at it? Was trying to see if float16 could be any faster. In my test it's hardly 10% faster than float32-high. I removed the graph construction in this benchmark.

[AlexanderMath]

Yes. I got ~3x speedup for the orbmol model when switching tf32->fp16/fp16 on rtx4090.

Happy to vchat and discuss further.

It's quite involved, naively, you don't get anything because pytorch defaults to accumulate fp16 in fp32 which on ?most? GPUs gives less speedups.

[benrhodes26]

@thomasloux what GPU did you experiment with? I think there is likely a significant difference between an rtx4090 and an A100/H100 (which I think receive greater acceleration from fp32 -> tf32, but perhaps less from tf32 -> fp16)

[thomasloux]

Indeed I use a H100. It's also possible that the float32 high directly uses bfloat16 on this GPU, as its described on the PyTorch page. Also, I tried to run with float32-medium precision, and there is no difference with float32-high. I don't really know why but I'm not an expert in this low level optimisation

[AlexanderMath]

@benrhodes26 I've used orbmol to simulate waters phase transition with/without ice-blockers.

I works even with the edge_mlp in fp8.

Does anyone have an intuition for the features learnt by the edge_mlp (e.g. mechanistic interpretability) and how these features don't get destroyed by fp8?

obs: this is great for 4090 80tflops->660tflops, but for h200 i'd wager edge_list creation dominates to fp8 may not help.