Unified memory for extra large systems - rematerialization problems

[GorgetteGorg]

Hi,

I'm trying to compute extra large systems (more than 5.120 residues) and as advised in #31 I tried to use unified memory.

I get this kind of error during the inference (details in out.txt) :

2025-05-27 11:26:29.427895: W external/xla/xla/service/hlo_rematerialization.cc:3005] Can't reduce memory use below -20.61GiB (-22125470549 bytes) by rematerialization; only reduced to 28.35GiB (30443270400 bytes), down from 28.35GiB (30443270400 bytes) originally

I have tried disabling rematerialization, and am waiting for the job to finish computing. This is the specifics of the node where this job is running, this is one of the usual nodes I use (sorry for bad formating):

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 560.35.03              Driver Version: 560.35.03      CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA H100 PCIe               On  |   00000000:17:00.0 Off |                    0 |
| N/A   68C    P0            309W /  310W |   60332MiB /  81559MiB |    100%      Default |
|                                         |                        |             Disabled |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A    533459      C   /alphafold3_venv/bin/python                 60320MiB |
+-----------------------------------------------------------------------------------------+

I am also running tests with A100 at the moment to see if it's a compatibility issue as mentioned in #341

my flags at the moment for my usual jobs are :

singularity exec \
     --nv \
     --bind $AF3_CODE_DIR:$AF3_CODE_DIR \
     --env XLA_FLAGS="--xla_disable_hlo_passes=custom-kernel-fusion-rewriter" \
     --env XLA_PYTHON_CLIENT_PREALLOCATE="false" \
     --env TF_FORCE_UNIFIED_MEMORY="true" \
     --env XLA_CLIENT_MEM_FRACTION="3.2" \
     --env JAX_PLATFORMS="cuda" \
     --bind $AF3_INPUT_DIR:/root/af_input \
     --bind $AF3_OUTPUT_DIR:/root/af_output \
     --bind $AF3_MODEL_PARAMETERS_DIR:/root/models \
     --bind $AF3_DATABASES_DIR:/root/public_databases \
     $AF3_IMAGE \
     python $AF3_CODE_DIR/run_alphafold.py \
     --json_path=/root/af_input/Nitrosopulimus_6mer.json \
     --model_dir=/root/models \
     --db_dir=/root/public_databases \
     --output_dir=/root/af_output \
     --flash_attention_implementation=xla

I am trying to debug as recommended in here: #341 (comment)

However any additional insight would be appreciated, as I am very much a beginner in comp chem and way out of my depth.

Thanks in advance!

[Augustin-Zidek]

Hello,

A few more clarifying questions:

• How much RAM (non-GPU) do you have on the host system?
• How many tokens does your input have?
• Since you are running on A100/H100, why are you setting the --xla_disable_hlo_passes=custom-kernel-fusion-rewriter flag? You should instead set --xla_gpu_enable_triton_gemm=false, see

  alphafold3/docker/Dockerfile

  Line 61 in 17afe15

  # To work around a known XLA issue causing the compilation time to greatly

[GorgetteGorg]

Hi,

Thanks for your answer.

• I ask for 380G in my slurm file. I have tried asking for more but it seems the nodes that have more are not available at the moment.

• the number of tokens ranges from 2,000 to 15,000. I am able to compute normally up to 5,120 tokens.

• IT provided template submission files for AF3, and when I asked about unified memory, they advised me to add this flag (--xla_disable_hlo_passes=custom-kernel-fusion-rewriter) as well as the 3 other ones. At this point I was still getting started with everything so I did not think to check. However I have launched one with --xla_gpu_enable_triton_gemm=false and it failed as well. I will try again (maybe with a smaller system, but still above 5,120 tokens) and let you know.

Thanks again!

[Augustin-Zidek]

Thanks for the additional information.

I ask for 380G in my slurm file.

Right, 380 GB should be sufficient.

The number of tokens ranges from 2,000 to 15,000. I am able to compute normally up to 5,120 tokens.

Useful to know you can run up to 5,120.

IT provided template submission files for AF3

Yeah, this is more about efficiency, so should not matter, but I think you are throwing some performance away by disabling that XLO pass and you should run with the recommended flags for optimal performance.

A few more debugging tips:

1. What is your host operating system? Is it Linux?
2. Could you run with without --flash_attention_implementation=xla? Again, you are leaving some performance and memory optimisations on the table.
3. Could you try updating past 6472373? I updated the CUDA used in the image from 12.6.0 -> 12.6.3 which might fix the issue you are seeing. Alternatively, you could go even further and try using nvidia/cuda:12.9.0-base-ubuntu24.04 instead of nvidia/cuda:12.6.3-base-ubuntu24.04 (look for it at the top of the Dockerfile). It might fix the issue, especially if you are seeing this only on the H100 but not on the A100.

[GorgetteGorg]

Hi,
Thanks for the reply !

• I have launched H100 jobs asking for 500G yesterday and they have failed due to Resources Exhausted as well so memory is indeed not the issue here.

• Just to be precise, when I said I could run up to 5,120 : I have not tested if it is actually an exact limit, but jobs with ~4600 tokens manage to compute, and jobs with 6,000 don't. I am running a test on this.

• I have launched jobs without the --flash_attention_implementation=xla flag, I will post here with an update when I have results.

• Yes the OS is Linux (CentOS Linux, if that matters ?)

• I will contact IT about your suggestion.

• Some additional info : I have tried to run the jobs with the --xla_gpu_enable_triton_gemm=false flag and I end up running into the same issues as in 5746 tokens job: RESOURCE_EXHAUSTED (H100 80G, despite unified memory, jax flag and increase bucket size) #341 (jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the function failed to launch on the GPU). I am still waiting for an A100 job (with the --xla_gpu_enable_triton_gemm=false flag, 320GB of memory, and 6,155 tokens), I will report back when it is finished.

Thanks again for the advice, and hopefully this will be resolved soon !
Best,

[GorgetteGorg]

Hi !

I thought I'd update :

• I have been able to run jobs for systems with 6,155 and 6,543 tokens on 40GB A100 using unified memory, they took like 4 and 5 days respectively.
• A system at 7,386 tokens fails
• We have H100 gpus with cuda 12.9 so I've been trying to use those for the 7,386 tokens systems, but there is an issue which had nothing to do with alphafold in itself. I will update if cuda 12.9 ends up solving things.

I just have a question for @Augustin-Zidek : is there a specific reason why the flag for unified memory is 3.2 ? From what I understand, this just dictates how much memory one is allowed to spill on the cpu relatively to the memory of the gpu, but if I were to set it at like 4.5 or higher (provided my cpu memory is large enough) could that impact the quality of the results / would that setup have to be re-tested against benchmark examples ?

Thanks in advance, and thanks for your help in the past : )

[Augustin-Zidek]

I just have a question for @Augustin-Zidek : is there a specific reason why the flag for unified memory is 3.2?

Not really, this is what empirically worked well and what we tested with. So if your machine has a lot of RAM, feel free to try a higher number to be able to fold larger proteins. It will be slow though. Also note that you would be running AlphaFold 3 well outside the rigorously tested regime, so you should be cautious about the results and be prepared to hit unusual behaviors.

Alternatively, you could consider chopping the sequence into multiple overlapping chunks (e.g. 1024 residues each, 256 residue overlaps), folding each separately, then stitching them back together -- but quality might suffer (understandably).

[GorgetteGorg]

Hi @Augustin-Zidek ,

I just wanted to let you know that I have been able to launch a job with 7,386 tokens on a GH200 with 12.9 cuda version (with the update you suggested a few weeks ago), with unified memory (3.2 as advised). It finished in 6 days.

My guess is that since the GH200 has a huge memory and the 7,386 tokens job finished, it can probably handle more. The limiting factor would be time here, and I don't have access to it anymore to run larger tests (nor would I want to if it blocks other users for that long).

I would also think that a smaller job with unified memory could finish on a H100 with cuda 12.9 instead of just failing from the start. However I am not a computer science person by any means and the cluster I'm using doesn't have that version on the H100, so that's just a guess.

Since I don't have access to the GH200 anymore I can't share the specs from nvidia-smi for people who might be interested, but at least for the record here are my flags as discussed above :

# Run AlphaFold3 container

singularity exec \
            --nv \
            --bind $AF3_CODE_DIR:$AF3_CODE_DIR \
            --env XLA_FLAGS="--xla_gpu_enable_triton_gemm=false" \
            --env XLA_PYTHON_CLIENT_PREALLOCATE="false" \
            --env TF_FORCE_UNIFIED_MEMORY="true" \
            --env XLA_CLIENT_MEM_FRACTION="3.2" \
            --env JAX_PLATFORMS="cuda" \
            --bind $LOCAL_INPUT:/root/af_input \
            --bind $LOCAL_OUTPUT:/root/af_output \
            --bind $AF3_MODEL_PARAMETERS_DIR:/root/models \
            --bind $AF3_DATABASES_DIR:/root/public_databases \
            $AF3_IMAGE \
            python $AF3_CODE_DIR/run_alphafold.py \
            --json_path=/root/af_input/Saccharolobus_6mer.json \
            --model_dir=/root/models \
            --db_dir=/root/public_databases \
            --output_dir=/root/af_output

Thanks for all the help !

[anivar]

Hi @GorgetteGorg,

I've dealt with similar large-scale protein complex simulations hitting memory walls. The negative memory target (-20.61GiB) in your rematerialization error is particularly interesting - it suggests XLA's memory optimizer is getting confused about memory requirements.

Here's a comprehensive approach to handle extra-large systems (>5,120 residues) on your H100:

Understanding the Memory Challenge

Your error indicates XLA can't find a valid rematerialization strategy:

Can't reduce memory use below -20.61GiB (-22125470549 bytes)

This negative value typically means:

1. The memory estimator is overflow/underflowing due to the system size
2. The computation graph is too large for standard optimization passes
3. There's a conflict between unified memory and rematerialization

Solution Strategy

1. Optimized JAX/XLA Configuration

# Enhanced memory configuration for H100 (80GB)
export XLA_FLAGS="--xla_gpu_enable_command_buffer= \
  --xla_gpu_enable_triton_softmax_fusion=false \
  --xla_gpu_enable_custom_fusions=false \
  --xla_gpu_enable_while_loop_double_buffering=false \
  --xla_gpu_deterministic_ops=true \
  --xla_force_host_platform_device_count=1 \
  --xla_disable_hlo_passes=custom-kernel-fusion-rewriter,gpu-conv-rewriter"

# JAX-specific optimizations
export JAX_ENABLE_X64=0  # Use float32 to save memory
export JAX_DEFAULT_MATMUL_PRECISION="high"
export JAX_TRACEBACK_FILTERING=off  # Better debugging

# Unified memory with larger fraction
export TF_FORCE_UNIFIED_MEMORY=1
export XLA_PYTHON_CLIENT_PREALLOCATE=false
export XLA_CLIENT_MEM_FRACTION=4.0  # Allow oversubscription

2. Custom Memory Management Script

Create a wrapper script that handles memory more intelligently:

#!/usr/bin/env python3
# run_alphafold_large.py

import os
import gc
import jax
import jax.numpy as jnp
from jax.experimental import mesh_utils
from jax.experimental.shard_map import shard_map
import logging

def configure_memory_for_large_systems():
    """Configure JAX/XLA for extra-large protein systems"""
    
    # Force single-device mode to avoid memory fragmentation
    jax.config.update('jax_platforms', 'cuda')
    jax.config.update('jax_enable_x64', False)
    
    # Configure memory allocation
    os.environ['XLA_PYTHON_CLIENT_MEM_GROWTH'] = 'true'
    
    # Custom memory pool settings
    import tensorflow as tf
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
            # Set explicit memory limit (leave 5GB for system)
            tf.config.experimental.set_virtual_device_configuration(
                gpu,
                [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=75000)]
            )

def create_chunked_inference_pipeline(model_runner, chunk_size=2560):
    """
    Split large systems into overlapping chunks for inference
    """
    
    def chunked_inference(input_features, num_residues):
        if num_residues <= chunk_size:
            # Standard inference for smaller systems
            return model_runner.run_inference(input_features)
        
        # Calculate chunk parameters
        overlap = 512  # Overlap between chunks
        stride = chunk_size - overlap
        num_chunks = (num_residues - overlap) // stride + 1
        
        logging.info(f"Processing {num_residues} residues in {num_chunks} chunks")
        
        chunk_results = []
        for i in range(num_chunks):
            start_idx = i * stride
            end_idx = min(start_idx + chunk_size, num_residues)
            
            # Extract chunk features
            chunk_features = extract_chunk_features(
                input_features, start_idx, end_idx
            )
            
            # Run inference on chunk
            with jax.default_device(jax.devices('gpu')[0]):
                chunk_result = model_runner.run_inference(chunk_features)
                
            chunk_results.append(chunk_result)
            
            # Force memory cleanup between chunks
            jax.clear_caches()  # Clear compilation caches
            gc.collect()
            
        # Merge chunk results
        return merge_chunk_results(chunk_results, overlap)
    
    return chunked_inference

def optimize_model_for_large_systems(model):
    """
    Apply model-specific optimizations for memory efficiency
    """
    
    # Gradient checkpointing for transformer layers
    from jax.experimental import checkify
    
    def add_gradient_checkpointing(layer):
        return jax.checkpoint(layer, policy=jax.checkpoint_policies.dots_with_no_batch_dims)
    
    # Apply to attention layers
    for module_name in ['pair_attention', 'msa_attention', 'evoformer']:
        if hasattr(model, module_name):
            setattr(model, module_name, add_gradient_checkpointing(getattr(model, module_name)))
    
    return model

3. Modified Singularity Launch Script

#!/bin/bash
# launch_alphafold3_large.sh

# System-specific optimizations for H100
export CUDA_VISIBLE_DEVICES=0
export NVIDIA_TF32_OVERRIDE=1

# Memory pressure handling
echo 3 > /proc/sys/vm/drop_caches  # Clear system cache (requires sudo)

# Launch with optimized settings
singularity exec \
    --nv \
    --bind $AF3_CODE_DIR:$AF3_CODE_DIR \
    --env XLA_FLAGS="--xla_gpu_enable_command_buffer= \
        --xla_gpu_deterministic_ops=true \
        --xla_gpu_enable_triton_softmax_fusion=false \
        --xla_disable_hlo_passes=custom-kernel-fusion-rewriter,rematerialization" \
    --env XLA_PYTHON_CLIENT_PREALLOCATE="false" \
    --env TF_FORCE_UNIFIED_MEMORY="1" \
    --env XLA_CLIENT_MEM_FRACTION="4.2" \
    --env JAX_ENABLE_X64="0" \
    --env TF_GPU_ALLOCATOR="cuda_malloc_async" \
    --env TF_CUDA_MALLOC_ASYNC_SUPPORTED_PREALLOC="0" \
    --bind $AF3_INPUT_DIR:/root/af_input \
    --bind $AF3_OUTPUT_DIR:/root/af_output \
    --bind $AF3_MODEL_PARAMETERS_DIR:/root/models \
    --bind $AF3_DATABASES_DIR:/root/public_databases \
    $AF3_IMAGE \
    python /path/to/run_alphafold_large.py \
    --json_path=/root/af_input/Nitrosopulimus_6mer.json \
    --model_dir=/root/models \
    --db_dir=/root/public_databases \
    --output_dir=/root/af_output \
    --chunk_size=2560 \
    --enable_chunking=true

4. Debugging Memory Issues

Add this debugging code to understand memory usage:

def profile_memory_usage():
    """Profile GPU memory usage during inference"""
    import nvidia_ml_py as nvml
    
    nvml.nvmlInit()
    handle = nvml.nvmlDeviceGetHandleByIndex(0)
    
    def get_memory_info():
        info = nvml.nvmlDeviceGetMemoryInfo(handle)
        return {
            'total': info.total / 1e9,
            'used': info.used / 1e9,
            'free': info.free / 1e9
        }
    
    # Monitor memory during key operations
    memory_log = []
    
    # Add hooks to model
    def memory_logging_wrapper(func):
        def wrapped(*args, **kwargs):
            pre_memory = get_memory_info()
            result = func(*args, **kwargs)
            post_memory = get_memory_info()
            
            memory_log.append({
                'function': func.__name__,
                'pre_memory_gb': pre_memory['used'],
                'post_memory_gb': post_memory['used'],
                'delta_gb': post_memory['used'] - pre_memory['used']
            })
            
            return result
        return wrapped
    
    return memory_logging_wrapper, memory_log

5. Alternative: Multi-GPU Strategy

If single GPU isn't sufficient, here's a multi-GPU approach:

def setup_multi_gpu_inference():
    """Configure for multi-GPU inference on large systems"""
    
    # Get all available GPUs
    devices = jax.devices('gpu')
    n_devices = len(devices)
    
    # Create device mesh
    mesh = mesh_utils.create_device_mesh((n_devices,))
    
    # Shard large tensors across devices
    from jax.experimental import PartitionSpec as P
    
    def shard_input_features(features):
        # Shard MSA and pair features across devices
        sharding_spec = {
            'msa': P('devices', None, None),
            'pair': P('devices', None, None),
            'atom_positions': P(None, None, None)  # Keep on single device
        }
        
        return jax.tree_map(
            lambda x, spec: jax.device_put(x, sharding=spec),
            features,
            sharding_spec
        )
    
    return shard_input_features

Immediate Recommendations

1. Disable rematerialization completely (as you're trying):

   --xla_disable_hlo_passes=rematerialization

2. Use mixed precision to reduce memory:

   --enable_mixed_precision=true

3. Try the chunking approach for systems >5k residues

4. Monitor actual memory usage:

   watch -n 1 nvidia-smi

The key insight is that for extra-large systems, you often need to work around XLA's optimization assumptions rather than with them. The chunking approach has worked well for systems up to 10k residues in my experience.

Let me know if you need help implementing any of these solutions or if you're still hitting memory issues!