RF3 error ConvertTritonGPUToLLVM

[PSchebek]

Hi There
First of all thanks for your great work! I am getting an error when running structure prediction with rf3. I am using this command:
rf3 fold inputs=test_input/test_struct.pdb ckpt_path=./checkpoints/rf3_foundry_01_24_latest_remapped.ckpt out_dir=test_pdb

And this is my error:

Unsupported conversion from bf16 to f16 LLVM ERROR: Unsupported rounding mode for conversion. #blocked = #ttg.blocked<{sizePerThread = [4, 4], threadsPerWarp = [4, 8], warpsPerCTA = [4, 1], order = [1, 0]}> #blocked1 = #ttg.blocked<{sizePerThread = [1, 8], threadsPerWarp = [8, 4], warpsPerCTA = [4, 1], order = [1, 0]}> #blocked2 = #ttg.blocked<{sizePerThread = [8, 1], threadsPerWarp = [8, 4], warpsPerCTA = [1, 4], order = [0, 1]}> #blocked3 = #ttg.blocked<{sizePerThread = [8, 1], threadsPerWarp = [4, 8], warpsPerCTA = [1, 4], order = [0, 1]}> #shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> #shared1 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0, 1]}> #smem = #ttg.shared_memory module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:70", "ttg.threads-per-warp" = 32 : i32} { tt.func public @fused_sigmoid_gated_dual_gemm_forward_kernel(%arg0: !tt.ptr<bf16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<bf16> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<bf16> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32 {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32}, %arg6: !tt.ptr<bf16> {tt.divisibility = 16 : i32}) attributes {noinline = false} {
...
{-# external_resources: { mlir_reproducer: { pipeline: "builtin.module(triton-nvidia-mma-lowering, tritongpu-combine-tensor-select-and-if, tritongpu-allocate-warp-groups, convert-scf-to-cf, allocate-shared-memory, triton-tensor-memory-allocation, tritongpu-global-scratch-memory-allocation, convert-triton-gpu-to-llvm{compute-capability=70 ptx-version=84}, canonicalize{ max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true}, cse, convert-nv-gpu-to-llvm, convert-warp-specialize-to-llvm, canonicalize{ max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true}, cse, symbol-dce, enable-line-info)", disable_threading: false, verify_each: true } } #-} /work/ps40kyji-RFD3/foundry/.venv/lib/python3.12/site-packages/cuequivariance_ops/triton/gated_gemm_triton.py:19:0: error: Failures have been detected while processing an MLIR pass pipeline /work/ps40kyji-RFD3/foundry/.venv/lib/python3.12/site-packages/cuequivariance_ops/triton/gated_gemm_triton.py:19:0: note: Pipeline failed while executing [ConvertTritonGPUToLLVMon 'builtin.module' operation]: reproducer generated atstd::errs, please share the reproducer above with Triton project.Traceback (most recent call last):
What could cause this issue?
Thanks very much in advance!

[Ubiquinone-dot]

What accelerator are you using? Looks like the kernels might not be properly installed / supported

[PSchebek]

I tried running this on a Nvidia Tesla A30 and V100 with the same result. Colleagues of me used the same on the same HPC without issues, i am wondering if this is because i installed it recently (mid february 26) and their installation is slighty older.

[iiLiniesta]

I met the same problem as @PSchebek .
My versions: Python 3.12, rc-foundry 0.1.11.
I'm running rf3 using command rf3 predict inputs=xxx.json ckpt_path=xxx.ckpt outdir=xxx.
I suppose this is a length-related problem, according to my runs. I was trying to refold some sequences generated by RFd3.

First I successfully ran RF3 on some short (<50aa) sequences, as testing examples.

But the following input (50+51aa) would cause the error:
Unsupported conversion from bf16 to f16
LLVM ERROR: Unsupported rounding mode for conversion.

{
"name": "xyz50_51",
"components": [
{
"seq": "ALVAAAAAIAATAAALQAAALADAAASAQVIAQAAAARSAAADLADAIAA",
"chain_id": "A"
},
{
"seq": "DLAEALAAAESLAATAASVAAASSSAALQALAATIASLAAATAALLSAAEQ",
"chain_id": "B"
}
]
}

I'm not very sure, but for tasks of <=2 chains, no msa and only 20 classic amino acids, the limit seems to be: sum(lengths)<=100aa. I verified this pattern on lengths like 50+50, 86+14, 85+15, 100, 101, etc.

Another evidence is when I run using docs/releases/rf3/examples/7o1r_from_json.json, the same error occurred.
But if i shorten the sequence to 21aa, RF3 can give a result successfully; and on 22aa, it cannot. Like this:

[
{
"name": "7o1r_from_json",
"components": [
{
"seq": "MKSLSFSLALGFGSTLVYSAP",
"msa_path": "docs/releases/rf3/examples/msas/7o1r_A.a3m.gz",
"chain_id": "A"
},
{
"ccd_code": "MG"
},
{
"path": "docs/releases/rf3/examples/ligands/NAG.cif"
},
{
"path": "docs/releases/rf3/examples/ligands/HEM.sdf"
},
{
"smiles": "[nH]1cc[nH+]c1"
},
{
"ccd_code": "NAG",
"chain_id": "F"
}
]
}
]

I assume that this is because HEM.sdf has 75 atoms, and other components like "MG", "NAG" and the SMILES code might add up to a "sum(lengths)" value, so the lengths finally add up to roughly 100 when the protein sequences has 21aa.

I'm going to try some older versions of rc-foundry for now.
But on version 0.1.10 and 0.1.9, I met this: FileNotFoundError: Project root directory not found. Indicators: ['.project-root'], which is mentioned at #94

[PSchebek]

Thanks for your input. I tested it again using much shorter sequences and rf3 ran just fine. I only get this error when predicting structures < 100 aa

[iiLiniesta]

I thins this bug ultimately stems from Turing GPUs not supporting bfloat16. I was using RTX 2080 Ti.
See:
triton-lang/triton#6698
vllm-project/vllm#17639

Then I'm curious about why RF3 runs fine on my GPU, though with shorter sequences.
And since the conversion related with bfloat16 can successfully run on Turing GPU with shorter sequences, maybe (I hope) this bug could be fixed on longer sequences.