From d10b027de61855a0ed778d377f55e6f72ba39368 Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Wed, 20 May 2026 15:53:33 -0700 Subject: [PATCH 1/9] Restructure chunk size settings This converts the configured chunk size to the *maximum* rather than the *minimum*. I don't think there's a clear use case for this to be the minimum and users were confused. This also allows us to remove the hardcoded values for the max with default and flash-style kernels: these instead are just config settings. I removed the minimum chunk size entirely, so the tuner will search all powers of 2 up to the maximum. I'm struggling to think of a use case where someone would want to limit chunks size to a minimum of 4 (I think this number just came from the original suggested chunk size in the AlphaFold 3 paper). Note that this changes the hardcoded different chunk size for diffusion conditioning. I'm not sure how important this is to preserve and whether we saw actual latency improvements from raising the chunk size here vs just observing that it was possible to do so. If we want different chunk sizes for different modules, I think we should probably thread through a config setting per module because having a hardcoded value here is a bit surprising IMO (and note that it only gets applied when tuning is turned on). Probably the best analog in the config would be `offload_inference` as a per-module setting that isn't an init arg. If more things got added here, IDK whether we'd want to keep this structure or perhaps mirror the structure of architecture. --- .../example_runner_yamls/cuequivariance.yml | 5 +++-- examples/example_runner_yamls/triton.yml | 1 + openfold3/core/model/latent/base_stacks.py | 14 +++----------- openfold3/core/model/latent/pairformer.py | 14 +++----------- .../core/model/latent/template_module.py | 14 +++----------- .../model/layers/diffusion_conditioning.py | 3 +-- openfold3/core/utils/chunk_utils.py | 19 +++++-------------- .../config/model_setting_presets.yml | 7 +++++-- 8 files changed, 24 insertions(+), 53 deletions(-) diff --git a/examples/example_runner_yamls/cuequivariance.yml b/examples/example_runner_yamls/cuequivariance.yml index 49ded9071..a38fd417c 100644 --- a/examples/example_runner_yamls/cuequivariance.yml +++ b/examples/example_runner_yamls/cuequivariance.yml @@ -1,5 +1,5 @@ model_update: - presets: + presets: - predict - low_mem # to use low memory settings custom: @@ -7,4 +7,5 @@ model_update: memory: eval: use_cueq_triangle_kernels: true - use_deepspeed_evo_attention: true \ No newline at end of file + use_deepspeed_evo_attention: true + chunk_size: 1024 diff --git a/examples/example_runner_yamls/triton.yml b/examples/example_runner_yamls/triton.yml index b614acff7..10d0212b1 100644 --- a/examples/example_runner_yamls/triton.yml +++ b/examples/example_runner_yamls/triton.yml @@ -9,3 +9,4 @@ model_update: use_triton_triangle_kernels: true use_deepspeed_evo_attention: false use_cueq_triangle_kernels: false + chunk_size: 1024 diff --git a/openfold3/core/model/latent/base_stacks.py b/openfold3/core/model/latent/base_stacks.py index 22896844a..c00387384 100644 --- a/openfold3/core/model/latent/base_stacks.py +++ b/openfold3/core/model/latent/base_stacks.py @@ -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 @@ -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 @@ -143,8 +136,7 @@ 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) @@ -155,7 +147,7 @@ def block_with_cache_clear(block, *args, **kwargs): 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 ] diff --git a/openfold3/core/model/latent/pairformer.py b/openfold3/core/model/latent/pairformer.py index 8ef7b3540..ba3713818 100644 --- a/openfold3/core/model/latent/pairformer.py +++ b/openfold3/core/model/latent/pairformer.py @@ -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 @@ -376,9 +372,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], # We don't want to write in-place during chunk tuning runs @@ -386,8 +379,7 @@ def block_with_cache_clear(block, *args, **kwargs): 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) @@ -398,7 +390,7 @@ def block_with_cache_clear(block, *args, **kwargs): 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 ] diff --git a/openfold3/core/model/latent/template_module.py b/openfold3/core/model/latent/template_module.py index 1cf0bb035..790f47127 100644 --- a/openfold3/core/model/latent/template_module.py +++ b/openfold3/core/model/latent/template_module.py @@ -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 @@ -391,14 +387,10 @@ 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) @@ -407,7 +399,7 @@ def _prep_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 ] diff --git a/openfold3/core/model/layers/diffusion_conditioning.py b/openfold3/core/model/layers/diffusion_conditioning.py index 2059c62a0..7fd86fa5e 100644 --- a/openfold3/core/model/layers/diffusion_conditioning.py +++ b/openfold3/core/model/layers/diffusion_conditioning.py @@ -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( diff --git a/openfold3/core/utils/chunk_utils.py b/openfold3/core/utils/chunk_utils.py index 4c9b850ca..e609d67f1 100644 --- a/openfold3/core/utils/chunk_utils.py +++ b/openfold3/core/utils/chunk_utils.py @@ -25,9 +25,6 @@ tree_map, ) -DEFAULT_MAX_CHUNK_SIZE = 512 -FLASH_MAX_CHUNK_SIZE = 1024 - def _fetch_dims(tree): shapes = [] @@ -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: @@ -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 @@ -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 diff --git a/openfold3/projects/of3_all_atom/config/model_setting_presets.yml b/openfold3/projects/of3_all_atom/config/model_setting_presets.yml index 5a2b34e53..8249a4a1d 100644 --- a/openfold3/projects/of3_all_atom/config/model_setting_presets.yml +++ b/openfold3/projects/of3_all_atom/config/model_setting_presets.yml @@ -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 From 52236b73ef736301a7f043d361f3c5ca625cb6b8 Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Wed, 20 May 2026 16:46:17 -0700 Subject: [PATCH 2/9] Add tests for chunk size tuner --- openfold3/tests/utils/test_utils.py | 129 +++++++++++++++++++++++++++- 1 file changed, 128 insertions(+), 1 deletion(-) diff --git a/openfold3/tests/utils/test_utils.py b/openfold3/tests/utils/test_utils.py index 63a33b284..97a5395f3 100644 --- a/openfold3/tests/utils/test_utils.py +++ b/openfold3/tests/utils/test_utils.py @@ -18,7 +18,7 @@ import torch from openfold3.core.model.primitives import Linear -from openfold3.core.utils.chunk_utils import _chunk_slice, chunk_layer +from openfold3.core.utils.chunk_utils import ChunkSizeTuner, _chunk_slice, chunk_layer from openfold3.core.utils.rigid_utils import ( Rigid, Rotation, @@ -196,3 +196,130 @@ def test_chunk_slice_dict(self): chunked_flattened = x_flat[i:j] self.assertTrue(torch.all(chunked == chunked_flattened)) + + def test_chunk_size_tuner_picks_largest_viable(self): + # When the cutoff sits between two power-of-2 candidates, the tuner + # should pick the largest viable power of 2 at or below the cutoff. + cases = [ + # (max_viable, expected_chunk_size) + (1024, 1024), + (512, 512), + (511, 256), + (256, 256), + (255, 128), + (128, 128), + (4, 4), + (3, 2), + (1, 1), + ] + for max_viable, expected in cases: + with self.subTest(max_viable=max_viable): + + def fn(arg, chunk_size, _max=max_viable): + if chunk_size > _max: + raise RuntimeError("simulated OOM") + + result = ChunkSizeTuner._determine_favorable_chunk_size( + fn, args=(None,), max_chunk_size=1024 + ) + self.assertEqual(result, expected) + + def test_chunk_size_tuner_caps_at_max_chunk_size(self): + # max_chunk_size is the config-level ceiling: even when much larger + # values would fit, the tuner must not exceed it. + for max_chunk_size in (4, 16, 128, 512, 1024): + with self.subTest(max_chunk_size=max_chunk_size): + + def fn(arg, chunk_size): + return None # never raises -- any chunk_size "fits" + + result = ChunkSizeTuner._determine_favorable_chunk_size( + fn, args=(None,), max_chunk_size=max_chunk_size + ) + self.assertEqual(result, max_chunk_size) + + def test_chunk_size_tuner_caches_for_same_args(self): + # Repeated calls with identical arg shapes should be a cache hit: the + # fn must not be re-invoked after the initial tuning pass. + tuner = ChunkSizeTuner() + tested = [] + + def fn(t, chunk_size, tested=tested): + tested.append(chunk_size) + + args = (torch.zeros(2, 3, 4),) + first = tuner.tune_chunk_size( + representative_fn=fn, args=args, max_chunk_size=64 + ) + after_first = len(tested) + second = tuner.tune_chunk_size( + representative_fn=fn, args=args, max_chunk_size=64 + ) + + self.assertEqual(first, second) + self.assertGreater(after_first, 0) + self.assertEqual( + len(tested), + after_first, + f"fn was re-invoked on cache hit: {tested[after_first:]}", + ) + + def test_chunk_size_tuner_retunes_for_different_shape(self): + # Different arg shapes should invalidate the cache and trigger + # re-tuning. + tuner = ChunkSizeTuner() + tested = [] + + def fn(t, chunk_size, tested=tested): + tested.append(chunk_size) + if chunk_size > t.shape[-1]: + raise RuntimeError("simulated OOM") + + first = tuner.tune_chunk_size( + representative_fn=fn, + args=(torch.zeros(2, 3, 16),), + max_chunk_size=256, + ) + after_first = len(tested) + second = tuner.tune_chunk_size( + representative_fn=fn, + args=(torch.zeros(2, 3, 128),), + max_chunk_size=256, + ) + + self.assertNotEqual( + first, + second, + "Chunk size should have been re-tuned for new arg shape", + ) + self.assertGreater( + len(tested), + after_first, + "fn was not re-invoked on cache miss", + ) + + def test_chunk_size_tuner_non_power_of_two_max(self): + # When max_chunk_size isn't a power of 2, it should still be tried as + # a candidate (and returned when viable). + def fits_all(arg, chunk_size): + return None + + self.assertEqual( + ChunkSizeTuner._determine_favorable_chunk_size( + fits_all, args=(None,), max_chunk_size=500 + ), + 500, + ) + + # And when only powers of 2 below the max are viable, fall back to the + # largest such power of 2. + def fits_up_to_256(arg, chunk_size): + if chunk_size > 256: + raise RuntimeError("simulated OOM") + + self.assertEqual( + ChunkSizeTuner._determine_favorable_chunk_size( + fits_up_to_256, args=(None,), max_chunk_size=500 + ), + 256, + ) From 476eafc8681643e5be1f20a8227e518064a933d6 Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Wed, 20 May 2026 21:39:33 -0700 Subject: [PATCH 3/9] Add e2e test of chunk_size=1 This would have caught the issue with the chunk size for attention getting divided by 4. I think for only that bug, it's not really carrying its weight, but testing this boundary case still seems useful. --- openfold3/tests/test_of3_model.py | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/openfold3/tests/test_of3_model.py b/openfold3/tests/test_of3_model.py index 1dbf07a1e..788c488a9 100644 --- a/openfold3/tests/test_of3_model.py +++ b/openfold3/tests/test_of3_model.py @@ -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" @@ -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 config.architecture.loss_module.diffusion.chunk_size = 16 of3 = OpenFold3AllAtom(config).to(device=device, dtype=dtype) @@ -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): + 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( From d89aeae80e0211a342037786e266118082211ed7 Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Wed, 20 May 2026 21:40:41 -0700 Subject: [PATCH 4/9] Fix chunk size // 4 leading to attention chunk 0 --- openfold3/core/model/latent/base_stacks.py | 2 +- openfold3/core/model/latent/pairformer.py | 2 +- openfold3/core/model/latent/template_module.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/openfold3/core/model/latent/base_stacks.py b/openfold3/core/model/latent/base_stacks.py index c00387384..03af31c6d 100644 --- a/openfold3/core/model/latent/base_stacks.py +++ b/openfold3/core/model/latent/base_stacks.py @@ -139,7 +139,7 @@ def block_with_cache_clear(block, *args, **kwargs): 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( diff --git a/openfold3/core/model/latent/pairformer.py b/openfold3/core/model/latent/pairformer.py index ba3713818..7f058c8c2 100644 --- a/openfold3/core/model/latent/pairformer.py +++ b/openfold3/core/model/latent/pairformer.py @@ -382,7 +382,7 @@ def block_with_cache_clear(block, *args, **kwargs): 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( diff --git a/openfold3/core/model/latent/template_module.py b/openfold3/core/model/latent/template_module.py index 790f47127..9017d8a19 100644 --- a/openfold3/core/model/latent/template_module.py +++ b/openfold3/core/model/latent/template_module.py @@ -393,7 +393,7 @@ def _prep_blocks( 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( From 8bbcb5ec82a407c524ef1a02b75221d74f00be43 Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Thu, 28 May 2026 09:45:50 -0700 Subject: [PATCH 5/9] Set chunk size to 1024 in model presets --- examples/example_runner_yamls/cuequivariance.yml | 1 - examples/example_runner_yamls/triton.yml | 1 - .../projects/of3_all_atom/config/model_setting_presets.yml | 2 +- 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/example_runner_yamls/cuequivariance.yml b/examples/example_runner_yamls/cuequivariance.yml index a38fd417c..14539698f 100644 --- a/examples/example_runner_yamls/cuequivariance.yml +++ b/examples/example_runner_yamls/cuequivariance.yml @@ -8,4 +8,3 @@ model_update: eval: use_cueq_triangle_kernels: true use_deepspeed_evo_attention: true - chunk_size: 1024 diff --git a/examples/example_runner_yamls/triton.yml b/examples/example_runner_yamls/triton.yml index 10d0212b1..b614acff7 100644 --- a/examples/example_runner_yamls/triton.yml +++ b/examples/example_runner_yamls/triton.yml @@ -9,4 +9,3 @@ model_update: use_triton_triangle_kernels: true use_deepspeed_evo_attention: false use_cueq_triangle_kernels: false - chunk_size: 1024 diff --git a/openfold3/projects/of3_all_atom/config/model_setting_presets.yml b/openfold3/projects/of3_all_atom/config/model_setting_presets.yml index 8249a4a1d..8b4875d71 100644 --- a/openfold3/projects/of3_all_atom/config/model_setting_presets.yml +++ b/openfold3/projects/of3_all_atom/config/model_setting_presets.yml @@ -15,7 +15,7 @@ predict: # 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 + chunk_size: 1024 offload_inference: confidence_heads: true token_cutoff: 2800 From bf1f80e8f10f0e01234dbecfe9fc24f6853023e4 Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Fri, 29 May 2026 12:41:18 -0700 Subject: [PATCH 6/9] Don't set chunk size for training I added an assert to guard against this so there isn't some surprising behavior. --- openfold3/tests/test_of3_model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/openfold3/tests/test_of3_model.py b/openfold3/tests/test_of3_model.py index 788c488a9..948fd9a07 100644 --- a/openfold3/tests/test_of3_model.py +++ b/openfold3/tests/test_of3_model.py @@ -65,7 +65,8 @@ def run_model( ) if chunk_size is not None: config.settings.memory.eval.chunk_size = chunk_size - config.settings.memory.train.chunk_size = chunk_size + assert not train + config.architecture.loss_module.diffusion.chunk_size = 16 of3 = OpenFold3AllAtom(config).to(device=device, dtype=dtype) From 40ae604de90c8802801faae94a1c26d1edecef8d Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Thu, 4 Jun 2026 11:23:53 -0700 Subject: [PATCH 7/9] Rename tests for consistency --- openfold3/tests/utils/test_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/openfold3/tests/utils/test_utils.py b/openfold3/tests/utils/test_utils.py index ee648fb98..523a8cd01 100644 --- a/openfold3/tests/utils/test_utils.py +++ b/openfold3/tests/utils/test_utils.py @@ -323,7 +323,7 @@ def fn(arg, chunk_size): ) self.assertEqual(result, max_chunk_size) - def test_chunk_size_tuner_retunes_for_different_shape(self): + def test_chunk_size_tuner_retunes_different_shape(self): # Different arg shapes should invalidate the cache and trigger # re-tuning. tuner = ChunkSizeTuner() @@ -349,7 +349,7 @@ def fn(t, chunk_size): "Chunk size should have been re-tuned for new arg shape", ) - def test_chunk_size_tuner_handles_arg_rank_change(self): + def test_chunk_size_tuner_retunes_different_rank(self): tuner = ChunkSizeTuner() def fn(t, chunk_size): @@ -374,7 +374,7 @@ def fn(t, chunk_size): first, second, "Chunk size should have been re-tuned for new arg rank" ) - def test_chunk_size_tuner_handles_dtype_bytes_change(self): + def test_chunk_size_tuner_retunes_different_dtype_bytes(self): tuner = ChunkSizeTuner() def fn(t, chunk_size): @@ -399,7 +399,7 @@ def fn(t, chunk_size): first, second, "Chunk size should have been re-tuned for new dtype bytes" ) - def test_chunk_size_tuner_handles_arg_count_change(self): + def test_chunk_size_tuner_retunes_different_arg_count(self): tuner = ChunkSizeTuner() def fn(*args, chunk_size): From 8880531810b5ca2101c23f04e970805ed762e044 Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Tue, 23 Jun 2026 09:20:54 -0700 Subject: [PATCH 8/9] Remove min_chunk_size arg from unit tests Missed this when merging --- openfold3/tests/utils/test_utils.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/openfold3/tests/utils/test_utils.py b/openfold3/tests/utils/test_utils.py index 523a8cd01..3af3e6d2f 100644 --- a/openfold3/tests/utils/test_utils.py +++ b/openfold3/tests/utils/test_utils.py @@ -210,7 +210,6 @@ def fn(t, chunk_size): first = tuner.tune_chunk_size( representative_fn=spy_fn, args=(torch.randn(2, 3, 4, 5),), - min_chunk_size=4, max_chunk_size=256, ) @@ -218,7 +217,6 @@ def fn(t, chunk_size): second = tuner.tune_chunk_size( representative_fn=spy_fn, args=(torch.randn(2, 3, 4, 5),), - min_chunk_size=4, max_chunk_size=256, ) @@ -246,7 +244,7 @@ def fn(arg, chunk_size, _max=max_viable, tested=tested): raise RuntimeError("simulated OOM") ChunkSizeTuner._determine_favorable_chunk_size( - fn, args=(None,), min_chunk_size=4, max_chunk_size=1024 + fn, args=(None,), max_chunk_size=1024 ) self.assertEqual( @@ -360,13 +358,11 @@ def fn(t, chunk_size): first = tuner.tune_chunk_size( representative_fn=fn, args=(torch.zeros(2, 3, 4, 5),), - min_chunk_size=4, max_chunk_size=256, ) second = tuner.tune_chunk_size( representative_fn=fn, args=(torch.zeros(2, 3, 4, 5, 6),), - min_chunk_size=4, max_chunk_size=256, ) @@ -385,13 +381,11 @@ def fn(t, chunk_size): first = tuner.tune_chunk_size( representative_fn=fn, args=(torch.zeros(2, 3, 4, 5, dtype=torch.float32),), - min_chunk_size=4, max_chunk_size=256, ) second = tuner.tune_chunk_size( representative_fn=fn, args=(torch.zeros(2, 3, 4, 5, dtype=torch.bfloat16),), - min_chunk_size=4, max_chunk_size=256, ) @@ -410,13 +404,11 @@ def fn(*args, chunk_size): first = tuner.tune_chunk_size( representative_fn=fn, args=(1, 2, 3, 4, 5), - min_chunk_size=4, max_chunk_size=256, ) second = tuner.tune_chunk_size( representative_fn=fn, args=(1, 2, 3, 4, 5, 6), - min_chunk_size=4, max_chunk_size=256, ) From 028a48b79799c8b8380c68b03ce3e10e1390c1de Mon Sep 17 00:00:00 2001 From: Geoffrey Martin-Noble Date: Tue, 30 Jun 2026 07:30:58 -0700 Subject: [PATCH 9/9] Update expected chunk config value --- openfold3/tests/test_entry_points.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openfold3/tests/test_entry_points.py b/openfold3/tests/test_entry_points.py index d13cb02b3..a5de59242 100644 --- a/openfold3/tests/test_entry_points.py +++ b/openfold3/tests/test_entry_points.py @@ -314,7 +314,7 @@ def test_low_mem_model_config_preset(self, tmp_path, dummy_ckpt_file): assert model_cfg.architecture.msa.msa_module_embedder.subsample_all_msa # check low memory settings set correctly - assert model_cfg.settings.memory.eval.chunk_size == 4 + assert model_cfg.settings.memory.eval.chunk_size == 1024 assert model_cfg.settings.memory.eval.offload_inference.confidence_heads assert model_cfg.settings.memory.eval.offload_inference.token_cutoff == 0