Flags enabling unified memory seemingly no longer working on most recent version of AF3

[ryanemenecker]

Specific objective: Running AF3 on large complexes (>5,120 tokens).
Specific issue: Flags previously used to enable unified memory do not seem to work on most recent version of AF3.
Have I ever gotten AF3 to work? Yes! I've gotten a previous version of AF3 before the tokamax changes on Novermber 29, 2025 to run no problem.
Details: I tried building a Docker container with the most recent AF3 modifications (pulled from most recent update 5 days ago). The new version will not run large complexes. I have pulled the most recent version and rebuilt my Docker container to make sure that bugs fixed recently are not causing the issue. I suspect that the flags previously recommended are not working for some reason. In my docker container I set
ENV XLA_PYTHON_CLIENT_PREALLOCATE=false
ENV TF_FORCE_UNIFIED_MEMORY=true
ENV XLA_CLIENT_MEM_FRACTION=3.2
as is recommended in the documentation. I am still getting OOM issues. These flags were sufficient to get the previous version to work when predicting large complexes without issues.

Hardware specifications
GPU: NVIDIA RTX PRO 6000 Blackwell Workstation Edition (96GB)
CUDA Driver Version: 580.105.08
Local CUDA Version: 13.0
Note: These previous specifications were not a problem when running the older version of AF3.

[Augustin-Zidek]

Hi, thanks a lot for reporting this. I've reported this to the Tokamax team, I will get back to you.

[ryanemenecker]

Thanks for the quick response! Let me know if there is any other information that I can provide that will help.

[Augustin-Zidek]

My current suspicion is that the TensorFlow import coming via Tokamax might be causing the issue. We have a PR to remove it, it should land in Tokamax soon. Once Tokamax tags 0.0.10 that includes this change, I will update AF3 and I hope that will fix it. I will let you know once that is done.

[Augustin-Zidek]

Actually, could you try running the offending input with:

• The latest AF3 version
• The ENV XLA_PYTHON_CLIENT_PREALLOCATE=false, ENV TF_FORCE_UNIFIED_MEMORY=true, ENV XLA_CLIENT_MEM_FRACTION=3.2
• The flag --flash_attention_implementation=xla

[ryanemenecker]

I gave that a run and ran into the same issue. For these builds I always have the following flags:
ENV XLA_PYTHON_CLIENT_PREALLOCATE=false
ENV TF_FORCE_UNIFIED_MEMORY=true
ENV XLA_CLIENT_MEM_FRACTION=3.2
I've tried a few flags in addition to the three above as follows (none of which have worked):

1. (your suggestion)ENV flash_attention_implementation=xla
2. ENV XLA_PYTHON_CLIENT_ALLOCATOR=platform
3. ENV XLA_PYTHON_CLIENT_ALLOCATOR=platform, ENV xla_disable_hlo_passes=custom-kernel-fusion-rewriter , ENV xla_gpu_strict_conv_algorithm_picker=false
   I've also tried these flags after changing the model_config.py pair_transition_shard_spec to

( (2048, None), (3072, 1024), (None, 512), )
as was recommended to reduce memory usage in the performance documentation (though this seemed overkill given I've got 96GB vRAM).
Thanks again for the response!

[ryanemenecker]

In case this helps, I only have this issue when the input is sufficiently large. If everything fits onto my GPU (no unified memory necessary) the most recent version of AlphaFold3 works fine with the exception that I had to add a line to my Docker file to reinstall jax[cuda12] after installing dependencies because it kept installing the CPU-only version of jax.

[Augustin-Zidek]

Tokamax removed the TF dependency in openxla/tokamax@ce4567e. Once they tag a new version, I will update AlphaFold 3 to use it. If you don't want to wait, you can change AF3 pyproject.toml to depend on the exact Tokamax commit instead of tokamax==0.0.8.

[Augustin-Zidek]

Sorry this took a bit longer. Could you try updating past ebfe70a and trying whether that fixes it?