Skip to content

Make config chunk_size maximum rather than minimum #227

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
christinaflo merged 11 commits into from
Jul 1, 2026
Merged

Make config chunk_size maximum rather than minimum #227

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
d10b027
Restructure chunk size settings
GMNGeoffrey May 20, 2026
52236b7
Add tests for chunk size tuner
GMNGeoffrey May 20, 2026
476eafc
Add e2e test of chunk_size=1
GMNGeoffrey May 21, 2026
d89aeae
Fix chunk size // 4 leading to attention chunk 0
GMNGeoffrey May 21, 2026
8bbcb5e
Set chunk size to 1024 in model presets
GMNGeoffrey May 28, 2026
bf1f80e
Don't set chunk size for training
GMNGeoffrey May 29, 2026
de89051
Merge branch 'main' into chunk-config
GMNGeoffrey Jun 4, 2026
40ae604
Rename tests for consistency
GMNGeoffrey Jun 4, 2026
8880531
Remove min_chunk_size arg from unit tests
GMNGeoffrey Jun 23, 2026
a287563
Merge branch 'main' into chunk-config
jandom Jun 30, 2026
028a48b
Update expected chunk config value
GMNGeoffrey Jun 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions examples/example_runner_yamls/cuequivariance.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
model_update:
presets:
presets:
- predict
- low_mem # to use low memory settings
custom:
settings:
memory:
eval:
use_cueq_triangle_kernels: true
use_deepspeed_evo_attention: true
use_deepspeed_evo_attention: true
chunk_size: 1024
Comment thread
GMNGeoffrey marked this conversation as resolved.
Outdated
1 change: 1 addition & 0 deletions examples/example_runner_yamls/triton.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,4 @@ model_update:
use_triton_triangle_kernels: true
use_deepspeed_evo_attention: false
use_cueq_triangle_kernels: false
chunk_size: 1024
16 changes: 4 additions & 12 deletions openfold3/core/model/latent/base_stacks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,7 @@
from torch import nn

from openfold3.core.utils.checkpointing import checkpoint_blocks
from openfold3.core.utils.chunk_utils import (
DEFAULT_MAX_CHUNK_SIZE,
FLASH_MAX_CHUNK_SIZE,
ChunkSizeTuner,
)
from openfold3.core.utils.chunk_utils import ChunkSizeTuner


# TODO: Rename to CheckpointStack and generalize any kind of block (i.e. remove
Expand Down Expand Up @@ -131,9 +127,6 @@ def block_with_cache_clear(block, *args, **kwargs):
or use_triton_triangle_kernels
or use_deepspeed_evo_attention
)
max_chunk_size = (
FLASH_MAX_CHUNK_SIZE if use_flash_kernels else DEFAULT_MAX_CHUNK_SIZE
)
tuned_chunk_size = self.chunk_size_tuner.tune_chunk_size(
representative_fn=blocks[0],
# Tensors cloned to avoid getting written to in-place
Expand All @@ -143,19 +136,18 @@ def block_with_cache_clear(block, *args, **kwargs):
m.clone(),
z.clone(),
),
min_chunk_size=chunk_size,
max_chunk_size=max_chunk_size,
max_chunk_size=chunk_size,
)
attn_chunk = (
tuned_chunk_size if use_flash_kernels else (tuned_chunk_size // 4)
tuned_chunk_size if use_flash_kernels else max(1, tuned_chunk_size // 4)
)
blocks = [
partial(
b,
chunk_size=tuned_chunk_size,
# A temporary measure to address torch's occasional
# inability to allocate large tensors
_attn_chunk_size=max(chunk_size, attn_chunk),
_attn_chunk_size=attn_chunk,
Comment thread
christinaflo marked this conversation as resolved.
)
for b in blocks
]
Expand Down
16 changes: 4 additions & 12 deletions openfold3/core/model/latent/pairformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,7 @@
from openfold3.core.model.layers.attention_pair_bias import AttentionPairBias
from openfold3.core.model.layers.transition import SwiGLUTransition
from openfold3.core.utils.checkpointing import checkpoint_blocks
from openfold3.core.utils.chunk_utils import (
DEFAULT_MAX_CHUNK_SIZE,
FLASH_MAX_CHUNK_SIZE,
ChunkSizeTuner,
)
from openfold3.core.utils.chunk_utils import ChunkSizeTuner
from openfold3.core.utils.tensor_utils import add


Expand Down Expand Up @@ -376,29 +372,25 @@ def block_with_cache_clear(block, *args, **kwargs):
or use_triton_triangle_kernels
or use_deepspeed_evo_attention
)
max_chunk_size = (
FLASH_MAX_CHUNK_SIZE if use_flash_kernels else DEFAULT_MAX_CHUNK_SIZE
)
tuned_chunk_size = self.chunk_size_tuner.tune_chunk_size(
representative_fn=blocks[0],
# We don't want to write in-place during chunk tuning runs
args=(
s.clone(),
z.clone(),
),
min_chunk_size=chunk_size,
max_chunk_size=max_chunk_size,
max_chunk_size=chunk_size,
)
attn_chunk = (
tuned_chunk_size if use_flash_kernels else (tuned_chunk_size // 4)
tuned_chunk_size if use_flash_kernels else max(1, tuned_chunk_size // 4)
)
blocks = [
partial(
b,
chunk_size=tuned_chunk_size,
# A temporary measure to address torch's occasional
# inability to allocate large tensors
_attn_chunk_size=max(chunk_size, attn_chunk),
_attn_chunk_size=attn_chunk,
)
for b in blocks
]
Expand Down
16 changes: 4 additions & 12 deletions openfold3/core/model/latent/template_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,11 +35,7 @@
from openfold3.core.model.latent.base_blocks import PairBlock
from openfold3.core.model.primitives import LayerNorm, Linear
from openfold3.core.utils.checkpointing import checkpoint_blocks, checkpoint_section
from openfold3.core.utils.chunk_utils import (
DEFAULT_MAX_CHUNK_SIZE,
FLASH_MAX_CHUNK_SIZE,
ChunkSizeTuner,
)
from openfold3.core.utils.chunk_utils import ChunkSizeTuner
from openfold3.core.utils.tensor_utils import add


Expand Down Expand Up @@ -391,23 +387,19 @@ def _prep_blocks(
or use_triton_triangle_kernels
or use_deepspeed_evo_attention
)
max_chunk_size = (
FLASH_MAX_CHUNK_SIZE if use_flash_kernels else DEFAULT_MAX_CHUNK_SIZE
)
tuned_chunk_size = self.chunk_size_tuner.tune_chunk_size(
representative_fn=blocks[0],
args=(t.clone(),),
min_chunk_size=chunk_size,
max_chunk_size=max_chunk_size,
max_chunk_size=chunk_size,
)
attn_chunk = (
tuned_chunk_size if use_flash_kernels else (tuned_chunk_size // 4)
tuned_chunk_size if use_flash_kernels else max(1, tuned_chunk_size // 4)
)
blocks = [
partial(
b,
chunk_size=tuned_chunk_size,
_attn_chunk_size=max(chunk_size, attn_chunk),
_attn_chunk_size=attn_chunk,
)
for b in blocks
]
Expand Down
3 changes: 1 addition & 2 deletions openfold3/core/model/layers/diffusion_conditioning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,8 +196,7 @@ def _chunk_forward(
zij.clone(),
token_mask,
),
min_chunk_size=chunk_size,
max_chunk_size=2048,
max_chunk_size=chunk_size,
)

si, zij = self._forward(
Expand Down
19 changes: 5 additions & 14 deletions openfold3/core/utils/chunk_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,6 @@
tree_map,
)

DEFAULT_MAX_CHUNK_SIZE = 512
FLASH_MAX_CHUNK_SIZE = 1024


def _fetch_dims(tree):
shapes = []
Expand Down Expand Up @@ -354,15 +351,13 @@ def __init__(self):
self.cached_arg_data = None

@staticmethod
def _determine_favorable_chunk_size(fn, args, min_chunk_size, max_chunk_size):
def _determine_favorable_chunk_size(fn, args, max_chunk_size):
logging.info("Tuning chunk size...")

if min_chunk_size >= max_chunk_size:
return min_chunk_size

candidates = [2**l for l in range(int(math.log(max_chunk_size, 2)) + 1)]
candidates = [c for c in candidates if c > min_chunk_size]
candidates = [min_chunk_size] + candidates
# If someone passed a chunk size that wasn't a power of two, consider it as well
if candidates[-1] < max_chunk_size:
candidates.append(max_chunk_size)

def test_chunk_size(chunk_size):
try:
Expand Down Expand Up @@ -403,10 +398,7 @@ def tune_chunk_size(
self,
representative_fn: Callable,
args: tuple[Any],
min_chunk_size: int,
# Heuristically, runtimes for most of the modules in the network
# plateau earlier than this on all GPUs I've run the model on.
max_chunk_size=DEFAULT_MAX_CHUNK_SIZE,
max_chunk_size: int,
) -> int:
def remove_tensors(a):
return a.shape if type(a) is torch.Tensor else a
Expand All @@ -424,7 +416,6 @@ def remove_tensors(a):
self.cached_chunk_size = self._determine_favorable_chunk_size(
fn=representative_fn,
args=args,
min_chunk_size=min_chunk_size,
max_chunk_size=max_chunk_size,
)
self.cached_arg_data = arg_data
Expand Down
7 changes: 5 additions & 2 deletions openfold3/projects/of3_all_atom/config/model_setting_presets.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,11 @@ predict:
settings:
memory:
eval:
# This chunk size is the minimum value, it will be auto-tuned by default
chunk_size: 4
# When chunk-size tuning is enabled (the default), this is the maximum
# tried; the tuner chooses the largest power of 2 up to this (plus the
# value itself if not a power of 2) that doesn't OOM. When tuning is
# disabled, this value is used as-is.
chunk_size: 512
offload_inference:
confidence_heads: true
token_cutoff: 2800
Expand Down
28 changes: 26 additions & 2 deletions openfold3/tests/test_of3_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ def run_model(
reduce_model_size=True,
use_deepspeed_evo_attention=False,
use_triton_triangle_kernels=False,
chunk_size=None,
):
device = "cuda" if torch.cuda.is_available() else "cpu"

Expand All @@ -59,8 +60,12 @@ def run_model(
use_deepspeed_evo_attention
)

if use_triton_triangle_kernels:
config.settings.memory.eval.use_triton_triangle_kernels = True
config.settings.memory.eval.use_triton_triangle_kernels = (
use_triton_triangle_kernels
)
if chunk_size is not None:
config.settings.memory.eval.chunk_size = chunk_size
config.settings.memory.train.chunk_size = chunk_size

@christinaflo christinaflo May 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We never really run chunking with training because the activations stack up anyway in the backward pass so it doesnt save you much. I havent actually run this but i think it may fail some assert. Diffusion conditioning at least has a assert not self.training for its chunking function

@christinaflo christinaflo May 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah okay i see now you only have a test for eval mode anyway

@christinaflo christinaflo May 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we just delete this line anyway since it cant run

@GMNGeoffrey GMNGeoffrey May 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case, shouldn't we remove chunk_size from the train settings entirely? I set it here because if chunk size is overridden for the test it would be surprising if it weren't overridden for training as well IMO. I can delete it if you prefer though

@christinaflo christinaflo May 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had it there originally because I needed chunking enabled during validation for some large samples but not for training, so it'll pick what to use in the model like: chunk_size=mode_mem_settings.chunk_size depending on the stage it's in. I guess there's nothing stopping you from doing it during training, it's just not really worth it ever.

We could fix the assert in diffusion conditioning to match the other modules so it runs:

if chunk_size is not None and self.chunk_size_tuner is not None:
    assert not self.training

or change model.py to always set it to None for training and not reference the config, i thought it was easier to distinguish in the config though

@christinaflo christinaflo Jun 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh yeah thats a lot of memory. on 80gb gpus with no kernels I can see the chunking take effect with seq lengths > 1500 tokens. I have a really old and messy script for benchmarking inference speed + mem using random_of3_features. I can clean it up and share it, but I have a bit of a backlog this week so I can just send it for testing #213 so I don't block this PR. As you said, this is really just a config change. Btw realistic values are n_msa=16384 (this will get subsampled to 1024 each recycle) and n_templ=4.

@GMNGeoffrey GMNGeoffrey Jun 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perfect, thanks! I think that gives me enough to do testing without being limited by what I can find in the pdb. A standardized script in-repo would be nice, but you don't need to rush to clean up yours 🙂

So those don't vary much with input or scale with n_tok? I found that homoers used significantly less memory at very large n_tok I think due to msa reuse for the shared sequences, but I didn't dig in.

I can test for smaller memory caps and the chunk tuner behavior by limiting torch mem_fraction. The only problem is the combinatorial explosion of options. If 80gb is of particular interest (H100?) I can test it as well. The other option is to test fixed chunk sizes and just track memory. I've got a memory snapshotting callback (happy to clean up and upstream if it would be generally useful). The only issues there are that tuning itself can affect peak due to some clones and then diffusion conditioning getting chunk size 2048 is guarded by the tuner getting on (is the diff there worth it?)

@christinaflo christinaflo Jun 3, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's just the max allowable input n_msa and n_templ, there could be less but it'll get capped at those numbers. I'm not sure why large homomers use less memory either off the top of my head, I'd have to look into that as well, since msas should still be capped at 1024 per recycle I'd expect they would still reach that threshold.
Yeah 80 gb h100 is generally what we have access to so that is of particular interest. Even on mi300A I try to limit model gpu memory quite a bit anyway since our data loader processes take up a lot of memory, so I kind of treat it like an h100.
For the memory snapshotting, I have a callback for this also internally that will get merged in at some point, I assume they're probably the same.
For the diffusion conditioning chunk size, this PR changes it to the global max but I think it's fine, I highly doubt it's that much slower and I eyeballed that number to begin with :).

@GMNGeoffrey GMNGeoffrey Jun 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah I looked closer and found the issue. One of the homomers I had Claude dig up for me (6R7M-1) is a 40-chain homomer with an MSA depth of only 122. This was throwing things off and I thought it would be an issue with all homomers to a lesser extent, but it's really just this one that is weird. If n_msa gets subsampled to 1024 every recycle, does it actually make a difference if the random input has n_msa=1024 or n_msa=16384?

For the diffusion conditioning chunk size, this PR changes it to the global max but I think it's fine, I highly doubt it's that much slower and I eyeballed that number to begin with :).

Oh right, I forgot that I did that already :-D

@christinaflo christinaflo Jun 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No actually it doesn't make a difference, it's only if you want to exercise the subsampling logic which doesn't matter here.

config.architecture.loss_module.diffusion.chunk_size = 16

of3 = OpenFold3AllAtom(config).to(device=device, dtype=dtype)
Expand Down Expand Up @@ -153,6 +158,25 @@ def test_shape_small_fp32(self, model_phase):
use_deepspeed_evo_attention=False,
)

def test_shape_small_chunk_size_one(self):
Comment thread
christinaflo marked this conversation as resolved.
batch_size = consts.batch_size
n_token = 18
n_msa = 10
n_templ = 3

self.run_model(
batch_size=batch_size,
n_token=n_token,
n_msa=n_msa,
n_templ=n_templ,
dtype=torch.float32,
train=False,
reduce_model_size=True,
use_deepspeed_evo_attention=False,
use_triton_triangle_kernels=False,
chunk_size=1,
)

@compare_utils.skip_unless_triton_installed()
@compare_utils.skip_unless_cuda_available()
@pytest.mark.parametrize(
Expand Down
Loading
Loading