Add graph break option for Mammut to work around torch.compile issue. Add timm ViT Block discovery if possible.

Author: rwightman

src/open_clip/mammut_model.py

@@ -130,6 +130,14 @@ def __init__(
 
         self.context_length = multimodal_cfg.context_length
 
+        # Runtime knob (--torchcompile-pass-break): under full-graph compile, split the graph
+        # between the two traversals of the shared decoder (contrastive pass, caption pass).
+        # One graph holding both checkpointed traversals makes the compiled backward retain a
+        # multi-block recompute working set (2-3.5x eager peak memory, torch 2.9-2.13); the
+        # break restores eager-like memory at no measured speed cost. Plain attribute rather
+        # than cfg: it is a compile-time training concern, not part of the model definition.
+        self.pass_graph_break = False
+
     def set_grad_checkpointing(self, enable: bool = True, impl: str = 'inline'):
         self.visual.set_grad_checkpointing(enable, impl=impl)
         self.text.set_grad_checkpointing(enable, impl=impl)
@@ -228,6 +236,11 @@ def forward(
         if image_latent is None:
             return {"text_features": text_latent}
 
+        if self.pass_graph_break:
+            # see ctor note: keep the contrastive and caption decoder traversals in separate
+            # compiled graphs; a no-op in eager and under the 'blocks' compile strategy
+            torch._dynamo.graph_break()
+
         # caption pass: causal self-attention w/ cross-attention over projected image tokens
         image_kv = image_embs @ self.map_viz2txt_kv
 

src/open_clip/task/base_task.py

@@ -250,7 +250,9 @@ def _get_fsdp_shard_modules(self) -> List[Tuple[str, nn.Module]]:
         """Discover modules to shard with FSDP2.
 
         Default: finds all ResidualAttentionBlock, CustomResidualAttentionBlock,
-        ModernTextBlock, and Bottleneck instances within the trainable module.
+        ModernTextBlock, Bottleneck, and timm ViT ``Block`` (timm towers, incl. NaFlexVit
+        trunks) instances within the trainable module. Also drives the 'blocks' compile
+        strategy, so timm vision trunks get per-block compile rather than staying eager.
         Models can override this by defining a ``fsdp_shard_modules()`` method.
         """
         model = unwrap_model(self.trainable_module)
@@ -260,7 +262,13 @@ def _get_fsdp_shard_modules(self) -> List[Tuple[str, nn.Module]]:
         from open_clip.transformer import ResidualAttentionBlock, CustomResidualAttentionBlock, ModernTextBlock
         from open_clip.modified_resnet import Bottleneck
 
-        shard_types = (ResidualAttentionBlock, CustomResidualAttentionBlock, ModernTextBlock, Bottleneck)
+        shard_types = [ResidualAttentionBlock, CustomResidualAttentionBlock, ModernTextBlock, Bottleneck]
+        try:
+            from timm.models.vision_transformer import Block as TimmViTBlock
+            shard_types.append(TimmViTBlock)
+        except ImportError:
+            pass
+        shard_types = tuple(shard_types)
 
         modules = []
         for name, mod in model.named_modules():

src/open_clip_train/main.py

@@ -412,6 +412,15 @@ def main(args):
         _logger.info(f'Disabling DDP dynamo optimizer ({reason}).')
         torch._dynamo.config.optimize_ddp = False
 
+    if args.torchcompile and args.torchcompile_pass_break:
+        inner = unwrap_model(model)
+        if hasattr(inner, 'pass_graph_break'):
+            inner.pass_graph_break = True
+            _logger.info('Enabling graph break between contrastive and caption decoder passes.')
+        else:
+            _logger.warning(
+                '--torchcompile-pass-break: model has no dual-pass graph break support; ignoring.')
+
     if args.torchcompile and args.torchcompile_strategy in ('model', 'blocks'):
         if args.fsdp:
             _logger.info(

src/open_clip_train/params.py

@@ -555,6 +555,17 @@ def parse_args(args):
              "automatic dynamic otherwise switches to a symbolic graph on the second shape, "
              "with worse memory planning.",
     )
+    parser.add_argument(
+        "--torchcompile-pass-break",
+        default=False,
+        action='store_true',
+        help="Insert a dynamo graph break between the contrastive and caption passes of dual-pass "
+             "decoder models (MaMMUT) under full-graph compile. One graph holding two checkpointed "
+             "traversals of the same decoder blocks makes the compiled backward retain a multi-block "
+             "recompute working set (2-3.5x eager peak memory); breaking between the passes restores "
+             "eager-like memory at no measured speed cost. Not needed with "
+             "--torchcompile-strategy blocks (recompute already stays eager there).",
+    )
     parser.add_argument(
         "--accum-freq", type=int, default=1, help="Update the model every --acum-freq steps."
     )

tests/test_naflex_mammut.py

@@ -281,3 +281,29 @@ def test_registered_config_builds(name, patch, seq_len):
     assert model.visual.trunk.get_patch_size() == (patch, patch)
     assert not isinstance(model.visual.trunk.norm_pre, torch.nn.Identity)
     assert model.visual.output_tokens
+
+
+# ---------------------------------------------------------------- task integration extras
+
+def test_blocks_discovery_includes_timm_trunk():
+    """The FSDP/blocks-compile discovery must cover the timm vision trunk, not just the text
+    stacks -- otherwise --torchcompile-strategy blocks leaves the whole vision tower eager."""
+    model = _tiny_model().train()
+    task = CoCaTask(model, verbose=False)
+    names = [n for n, _ in task._get_fsdp_shard_modules()]
+    trunk = [n for n in names if n.startswith('visual.trunk.blocks')]
+    assert len(trunk) == TINY_VISION_CFG['timm_model_kwargs']['depth'], \
+        f'timm trunk blocks not discovered: {names}'
+
+
+def test_mammut_pass_graph_break_eager_noop():
+    """pass_graph_break must be a pure compile-time hint: eager outputs identical either way."""
+    model = _tiny_model()
+    batch, text = _patch_batch(), _text_batch()
+    with torch.no_grad():
+        base = model(image=batch, text=text)
+        model.pass_graph_break = True
+        split = model(image=batch, text=text)
+    for k, v in base.items():
+        if torch.is_tensor(v):
+            torch.testing.assert_close(v, split[k], rtol=0, atol=0)