From 176b6cc2052f9b288873953ccb59e2c6c18efb5e Mon Sep 17 00:00:00 2001 From: Zhenshan Xie Date: Mon, 9 Nov 2026 15:41:39 -0800 Subject: [PATCH] docs(openfold3): record why FP8 is not shipped Add a README for the openfold3 family with the FP8 self-quantization study on TensorRT 11.4 (B200, B300) and 11.1: recipe (static per-tensor FP8 W8A8 on the Pairformer pair-stack and diffusion token block projections, plus a token-only variant), measurement method, GPU time split, speed, plan size and structure quality against the fp16 engine and the BF16 eager reference. GPU time per prediction is only 1.08x to 1.09x faster and process wall time is not faster; the plans are 18% to 20% smaller, but lDDT against the fp16 engine drops to 0.79 to 0.96 and two FP8 builds of the same recipe gave different structures. FP8 is not shipped. NVFP4 was not tried (no part is bound by weight reads). No problem on main was found. Documentation only; no code change. Signed-off-by: Zhenshan Xie --- families/openfold3/README.md | 172 +++++++++++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) create mode 100644 families/openfold3/README.md diff --git a/families/openfold3/README.md b/families/openfold3/README.md new file mode 100644 index 0000000000..2188e2ac49 --- /dev/null +++ b/families/openfold3/README.md @@ -0,0 +1,172 @@ +# OpenFold3 + +This family builds OpenFold3 v0.5.0 (`of3-ob-2025-06-30-174k` checkpoint) for +protein structure prediction as a bundle of TensorRT plans that the native C++ +runtime runs on one CUDA stream: input embedding, the 48 Pairformer blocks and +the trunk cycle (3 recycles), diffusion conditioning, the atom transformers, the +24 diffusion token blocks with the 200-step sampler, and the confidence heads. +Bundles are exact-request: one protein chain, query-only MSA, one diffusion +sample, batch-1 plans. The default precision is mixed fp16 (projections fp16; +norms, attention, triangle contractions and the diffusion state fp32); mixed +bf16 is the alternative. See `examples/models/openfold3/README.md` for the +build and the qualified ubiquitin profile. The family has no quantized build. + +## FP8 quantization study (not shipped) + +FP8 self-quantization of the Pairformer pair stack and the diffusion token +blocks was built and measured on TensorRT 11.4 (and 11.1) for the qualified +ubiquitin request (76 tokens, 601 atoms). It is not shipped: the GPU time per +prediction is only 1.08x to 1.09x faster, the process wall time is not faster, +and the predicted structure moves well beyond the fp16 rebuild floor (lDDT 0.79 +to 0.96 against the fp16 engine). Two FP8 builds of the same recipe gave +different structures. This section records the recipe, the method and the +numbers, so that the study does not have to be repeated. The FP8 code was not +merged; this family has no FP8 build. + +### Summary and decision + +| Item | fp16 (default) | FP8 | FP8 vs fp16 | +|---|---|---|---| +| GPU time per prediction, B200 / B300 (TensorRT 11.4) | 575 / 528 ms | 525 to 531 / 484 to 489 ms | 1.08x to 1.09x | +| Process wall time, B200 / B300 | 2.6 to 3.4 / 2.5 to 2.7 s | 2.7 to 3.6 / 2.5 to 3.0 s | not faster (run-to-run spread 0.2 to 1.2 s) | +| Bundle plan size | 907 to 908 MiB | 724 to 741 MiB | -18% to -20% | +| lDDT vs fp16 engine (fp16 rebuild 1.000, bf16 0.998 to 1.000) | 1.000 | 0.785 to 0.962 | far below the floor | +| pTM | 0.852 to 0.853 | 0.798 to 0.837 | -0.015 to -0.055 | +| Non-finite outputs | none | none | none | + +Decision: do not ship FP8 for OpenFold3. + +- Speed: the diffusion token blocks are about 64% of the GPU time and their FP8 + plans are 1.13x to 1.15x faster, which gives 8% to 9% of GPU time. The + process wall time (2.5 to 3.5 s) is mostly bundle load and setup, so the gain + is not visible there. +- FP8 in the Pairformer pair stack is slower than fp16 (K = 128 GEMMs at 76 + tokens). +- Quality: FP8 moves the structure far beyond the floor, and two FP8 builds of + the same recipe on B300 landed in different structures (lDDT vs fp16 0.954 + and 0.785), while fp16 rebuilds stay at lDDT 1.000. +- bf16 (already supported) is 1.02x to 1.04x faster than fp16 at floor quality. + +### Recipe that was measured + +No FP8, NVFP4 or ModelOpt checkpoint of OpenFold3 exists on the Hub, and no +official quantization config exists for it. The recipe is FP8 W8A8 with static +per-tensor scales. Each quantized projection is built as +`input -> Quantize -> Dequantize(fp16) -> MatMul`, with the weight +`FP8 constant -> Dequantize(fp16)`. + +| Choice | Decision | +|---|---| +| Quantized, Pairformer (48 blocks) | triangle multiplication `linear_a_p`, `linear_b_p`, `linear_a_g`, `linear_b_g`, `linear_g`, `linear_z`; triangle attention `mha.linear_q/k/v/g/o`; pair transition `swiglu.linear_a/b`, `linear_out` | +| Quantized, diffusion token blocks (24) | attention `mha.linear_q/k/v/g/o`; adaLN `linear_g`, `linear_s`, `linear_ada_out`; conditioned transition `swiglu.linear_a/b`, `linear_out`, `linear_g` and its adaLN | +| Kept fp16 / fp32 | triangle-attention bias `linear_z` (4 heads), token pair-bias `linear_z`, sequence attention, single transition, norms, attention cores, triangle contractions, atom transformers, confidence heads | +| Weight scale | per tensor, amax / 448, at build time | +| Activation scale | static, per tensor, max calibration (max abs input / 448) with forward pre-hooks on the eager OpenFold3 v0.5.0 model in BF16 autocast, 4 single-sequence proteins (thioredoxin, CI2, engrailed homeodomain, calmodulin N-lobe), disjoint from ubiquitin | +| Outlier rule | a scale above 20x the median of its kind stays unquantized; measured: no outlier | +| Dequantize type | fp16; no non-finite value | + +Two variants were measured: "pair + token" (all of the above, 1536 +projections) and "token only" (diffusion token blocks only). + +### Method + +- GPUs: B200 and B300, one GPU per job; TensorRT 11.4.0.106 (target) and + 11.1.0.106 (B200 only); torch 2.14.0 (CUDA 13). +- Request: the qualified ubiquitin request (3 recycles, 200 diffusion steps, + 1 sample, seed 42), the same pinned random samples in every engine. +- GPU time per prediction: median time of each plan times its call count, + summed. Wall time: median of 3 or 5 interleaved runs of the native + qualification binary (whole process, including bundle load). +- Quality: lDDT, RMSD and pLDDT error against the fp16 engine (floor: an fp16 + rebuild and the bf16 engine), and lDDT against the BF16 eager OpenFold3 + reference from `examples/models/openfold3/generate_reference.py`. The fp16 + engine itself reaches lDDT 0.783 (CA RMSD 2.58 A) against that reference. +- The token-only rows come from a second job with a fresh fp16 build, so they + are compared to that build. + +### Results + +GPU time split per prediction, fp16, TensorRT 11.4: + +| Part | Calls | B200 ms | B300 ms | +|---|---|---|---| +| Diffusion token transformer (4 plans x 6 blocks) | 200 x 4 | 367 (64%) | 330 (63%) | +| Diffusion score input and output | 200 x 2 | 150 | 140 | +| Pairformer (8 plans x 6 blocks) and trunk cycle | 4 x 9 | 56 | 55 | +| Embedder, conditioning, confidence | - | 2 | 2 | +| Total | - | 575 | 528 | + +Speed and size, TensorRT 11.4 (x is fp16 GPU time over variant GPU time): + +| GPU | Variant | Pairformer plan ms | Token plan ms | GPU time ms | x | Wall s | Plans MiB | +|---|---|---|---|---|---|---|---| +| B200 | fp16 (default) | 1.568 | 0.459 | 575 | 1.00 | 3.43 | 908 | +| B200 | bf16 | 1.393 | 0.436 | 553 | 1.04 | 2.76 | 911 | +| B200 | FP8 pair + token | 1.760 | 0.400 | 531 | 1.08 | 3.61 | 740 | +| B200 | fp16 rebuild (second job) | 1.544 | 0.458 | 574 | 1.00 | 2.63 | - | +| B200 | FP8 token only (second job) | 1.536 | 0.399 | 525 | 1.09 | 2.67 | 726 | +| B300 | fp16 (default) | 1.569 | 0.413 | 528 | 1.00 | 2.65 | 907 | +| B300 | bf16 | 1.345 | 0.409 | 519 | 1.02 | 2.59 | 911 | +| B300 | FP8 pair + token | 1.635 | 0.361 | 484 | 1.09 | 2.98 | 741 | +| B300 | fp16 rebuild (second job) | 1.605 | 0.410 | 526 | 1.00 | 2.55 | - | +| B300 | FP8 token only (second job) | 1.607 | 0.363 | 489 | 1.08 | 2.47 | 724 | + +Plan times are for one plan of 6 blocks. TensorRT 11.1 (B200): fp16 and FP8 +pair + token bundles build (913 and 749 MiB) and run; the Pairformer plan goes +from 1.502 to 1.697 ms with FP8 and the token plan stays at 0.422 ms, so there +is no gain on 11.1. + +Quality, ubiquitin: + +| GPU | Variant | lDDT vs fp16 | RMSD vs fp16 A | pLDDT MAE vs fp16 | pTM | lDDT vs eager ref | +|---|---|---|---|---|---|---| +| B200 | fp16 rebuild | 1.000 | 0.03 | 0.21 | 0.853 | 0.783 | +| B200 | bf16 | 1.000 | 0.04 | 0.21 | 0.854 | 0.783 | +| B200 | FP8 pair + token | 0.944 | 0.80 | 2.84 | 0.837 | 0.777 | +| B200 | FP8 token only | 0.961 | 0.68 | 2.32 | 0.825 | 0.791 | +| B300 | fp16 rebuild | 1.000 | 0.04 | 0.23 | 0.852 | 0.783 | +| B300 | bf16 | 0.998 | 0.20 | 0.75 | 0.852 | 0.782 | +| B300 | FP8 pair + token, build 1 | 0.954 | 0.70 | 6.08 | 0.806 | 0.780 | +| B300 | FP8 pair + token, build 2 | 0.785 | 2.41 | 6.44 | 0.804 | 0.880 | +| B300 | FP8 token only | 0.962 | 0.65 | 1.74 | 0.832 | 0.787 | +| B200, TensorRT 11.1 | FP8 pair + token | 0.913 | 1.66 | 6.65 | 0.798 | 0.781 | + +All outputs are finite for every variant. With FP8, mean pLDDT drops by 1.7 to +6.3 points. A scan over many inputs was not run: bundles are exact-request +(one 15 to 25 min build per input). + +### Why there is no gain + +- Only the diffusion token blocks gain from FP8 (1.13x to 1.15x per plan, fp16 + projections at 76 rows). They are about 64% of the GPU time, so the best case + is 8% to 9% of GPU time. +- The pair-stack GEMMs are small (K = 128 at 76 tokens) and do not gain from + FP8: the Pairformer plans are 4% to 13% slower and slightly larger (59.3 vs + 57.3 MiB). +- The process wall time is dominated by bundle load and setup, and its spread + between runs is larger than any FP8 effect. +- The diffusion rollout is sensitive: small changes in the token blocks move the + sample to a different nearby structure, so FP8 error shows up as large lDDT + changes, and two FP8 builds of the same recipe can disagree. + +### NVFP4 (not tried) + +No part is bound by weight reads (76-row GEMMs, small K in the pair stack), so +NVFP4 has no speed to gain, and FP8 already fails the quality bar. + +### What was not tried + +| Item | Note | +|---|---| +| FP8 attention cores, triangle contractions, atom transformers | not measured | +| Percentile or entropy calibration, per-channel weight scales | not measured | +| Longer sequences (up to 128 tokens), other proteins, several diffusion samples | not measured; one 15 to 25 min build per request | +| H100 | not measured | + +### When to revisit + +- Longer sequences or batched diffusion samples, where the token and pair GEMMs + have more rows and FP8 GEMMs can pay for the Quantize cost; the family builds + batch-1 exact-request bundles today. +- A recipe that keeps the diffusion rollout stable (lDDT vs fp16 near the fp16 + rebuild floor across several FP8 builds).