Optimize MoT attention and enable Flash4 BF16 defaults

docs/WORLD_MODEL.md

@@ -427,3 +427,55 @@ for each domain, synchronized physical clocks, measured-pose schemas, train-only
 normalization, dataset splits and complete preprocessing receipts. Then verify
 an actual batch and a scheduled full-size GPU optimizer/validation/reload check.
 Do not treat this code milestone as a prepared four-domain training dataset.
+
+
+## Selectable attention execution
+
+New continuous and discrete MoT experiments default to the combined Flash4 BF16
+attention profile. No attention override is needed: select the experiment normally.
+Weights, FFNs and losses remain FP32; TF32 stays disabled. Whole-block activation
+checkpointing is off and queries are unchunked.
+
+Explicit alternatives remain available through the already-selected config group:
+
+```text
+world_model_attention=unchunked_fp32
+world_model_attention=native
+```
+
+The shared continuous recipe selects the default and the discrete recipe inherits
+it. Historical checkpoint-bound sources/resolved configs and direct Python model
+constructor defaults remain unchanged. Inconsistent explicit strict-FP32/BF16
+declarations are rejected. Existing explicit per-model overrides still take
+precedence and must be inspected during a future launch preflight.
+
+Both optimized profiles disable whole-block activation checkpointing and set
+query chunk size zero. Every query still attends to every valid key across
+all modalities and timestamps. Image patchification, token counts, parameter
+counts, data sampling, normalization and loss weights are unchanged. The FP32
+profile requires fused efficient CUDA SDPA and rejects a math fallback.
+
+The Flash4 profile explicitly casts only attention Q/K/V to BF16 and converts
+its output back to the incoming dtype. With the existing `trainer.precision=32-true`,
+parameters, the remaining model and losses stay FP32. The Flash4 profile explicitly
+sets `runtime.strict_fp32=false` and `runtime.disable_tf32=true`: attention is
+intentionally BF16 while FP32 backend TF32 shortcuts remain disabled. It currently requires
+batch one and Hopper-or-newer CUDA hardware. Boolean padding (including holes)
+is handled by gathering all valid keys/values; all queries are retained.
+No implicit fallback, reduced attention window or token subsampling is used.
+
+The tested optional runtime is Torch2.9.1+cu126, torchvision0.24.1+cu126,
+FlashAttention4.0.0b35 (beta), CuTeDSL4.8.1, quack0.6.5, apache-tvm-ffi0.1.14.post1
+and cuda-bindings12.9.0. The ordinary project environment remains Torch2.7.1;
+**do not install Flash4 into it or replace its shared dependencies.** The
+optional backend is lazy imported only when selected and requires the separately
+verified runtime. The base package's Torch pins have not been relaxed; this is
+an explicit source/runtime execution profile rather than a claim that the
+standard dependency resolver installs that optional environment.
+
+The external beta API is checked for output container, shape, dtype and device.
+Numerical BF16 differences are intentional; finite gradients and short timings
+do not prove unchanged convergence or final task quality. A requested future
+full run must use its own resolved-source/config/runtime optimizer, validation
+and checkpoint compatibility preflight. No existing canceled run is resumed
+by selecting or adding these profiles.

egomimic/hydra_configs/experiment/world_model/continuous.yaml

@@ -1,6 +1,7 @@
 # @package _global_
 # Offline world-model recipe. Requires a separately verified prepared-data manifest.
 defaults:
+  - /world_model_attention: flash4_bf16
   - override /model: world_model/base
   - override /data: world_model/prepared
   - override /trainer: default
@@ -18,7 +19,7 @@ evaluator:
   sample_seed: 42
 trainer:
   accelerator: gpu
-  precision: 32-true  # User-selected FP32; inherited by the discrete comparison.
+  precision: 32-true  # FP32 weights/loss; attention precision comes from the profile.
   devices: 1
   min_epochs: 0
   max_epochs: -1
@@ -36,3 +37,8 @@ callbacks:
     save_on_train_epoch_end: false
     save_top_k: -1
     save_last: true
+
+# Resolve after root defaults so profile selection controls the precision contract.
+runtime:
+  strict_fp32: ${world_model_attention.strict_fp32}
+  disable_tf32: ${world_model_attention.disable_tf32}

egomimic/hydra_configs/model/world_model/base.yaml

@@ -38,7 +38,10 @@ pipeline:
         max_cameras: 3
         max_tracks: 1120
         max_tokens: 8192
-        # Exact full attention; opt-in memory controls pending full-size GPU proof.
-        attention_query_chunk_size: 0
-        activation_checkpointing: false
+        # Execution profiles do not change patch tokens, dimensions or objective.
+        attention_backend: ${oc.select:world_model_attention.backend,sdpa}
+        attention_dtype: ${oc.select:world_model_attention.dtype,native}
+        strict_fp32: ${runtime.strict_fp32}
+        attention_query_chunk_size: ${oc.select:world_model_attention.query_chunk_size,0}
+        activation_checkpointing: ${oc.select:world_model_attention.activation_checkpointing,false}
         time_mode: ${data.mode}

egomimic/hydra_configs/train_zarr_cartesian.yaml

@@ -43,6 +43,8 @@ runtime:
   # Opt in to disabling TF32 as well as autocast; requires precision=32-true
   # and NVIDIA_TF32_OVERRIDE=0 in the launch environment.
   strict_fp32: false
+  # Disable TF32 without declaring every model operation FP32.
+  disable_tf32: false
   # Lightning owns SIGUSR1/scontrol by default. An external ICE runner must
   # opt in explicitly and prove that child/framework requeue is disabled.
   slurm_requeue_owner: lightning

egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml

@@ -0,0 +1,13 @@
+# @package _global_
+defaults:
+  - unchunked_fp32
+  - _self_
+# Default attention-only BF16; keep trainer32-true and FP32 weights/FFNs/loss.
+# Requires the separately verified optional Torch2.9 + FlashAttention4 runtime.
+world_model_attention:
+  backend: flash4
+  dtype: bfloat16
+  strict_fp32: false
+# Intentional BF16 attention is not a strict-all-FP32 execution contract.
+runtime:
+  strict_fp32: false

egomimic/hydra_configs/world_model_attention/native.yaml

@@ -0,0 +1,9 @@
+# @package _global_
+# Historical default model arithmetic; per-run memory overrides remain supported.
+world_model_attention:
+  backend: sdpa
+  dtype: native
+  query_chunk_size: 0
+  activation_checkpointing: false
+  strict_fp32: false
+  disable_tf32: false

egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml

@@ -0,0 +1,14 @@
+# @package _global_
+defaults:
+  - native
+  - _self_
+# Exact full fused efficient attention, FP32 when the trainer is 32-true.
+world_model_attention:
+  backend: efficient
+  strict_fp32: true
+  disable_tf32: true
+trainer:
+  precision: 32-true
+runtime:
+  strict_fp32: true
+  disable_tf32: true

egomimic/models/world_model.py

@@ -55,14 +55,23 @@ def temporal_rope(x, times):
 class MoTBlock(nn.Module):
     """Private modality projections/FFNs, one joint attention operation."""
 
-    def __init__(self, width, heads, mlp_ratio=4, attention_query_chunk_size=0):
+    def __init__(self, width, heads, mlp_ratio=4, attention_query_chunk_size=0,
+                 attention_backend="sdpa", attention_dtype="native"):
         super().__init__()
         if width % heads or (width // heads) % 2:
             raise ValueError("RoPE requires an even attention head width")
         self.heads = heads
         if type(attention_query_chunk_size) is not int or attention_query_chunk_size < 0:
             raise ValueError("attention_query_chunk_size must be a nonnegative integer")
         self.attention_query_chunk_size = attention_query_chunk_size
+        if attention_backend not in ("sdpa", "efficient", "flash4"):
+            raise ValueError("attention_backend must be sdpa, efficient or flash4")
+        if attention_dtype not in ("native", "bfloat16"):
+            raise ValueError("attention_dtype must be native or bfloat16")
+        if attention_backend == "flash4" and attention_dtype != "bfloat16":
+            raise ValueError("flash4 requires explicit bfloat16 attention")
+        self.attention_backend = attention_backend
+        self.attention_dtype = attention_dtype
         self.norm1 = nn.ModuleList(nn.LayerNorm(width) for _ in MODALITIES)
         self.qkv = nn.ModuleList(nn.Linear(width, 3 * width) for _ in MODALITIES)
         self.out = nn.ModuleList(nn.Linear(width, width) for _ in MODALITIES)
@@ -100,13 +109,13 @@ def forward(self, streams, times, valid):
             parts = []
             for q in query.split(chunk, dim=2):
                 if torch.is_grad_enabled() and any(t.requires_grad for t in (q, key, value)):
-                    part = checkpoint(self._attend, q, key, value, mask, use_reentrant=False)
+                    part = checkpoint(self._apply_attention, q, key, value, mask, use_reentrant=False)
                 else:
-                    part = self._attend(q, key, value, mask)
+                    part = self._apply_attention(q, key, value, mask)
                 parts.append(part)
             result = torch.cat(parts, dim=2)
         else:
-            result = self._attend(query, key, value, mask)
+            result = self._apply_attention(query, key, value, mask)
         result = result.transpose(1, 2).flatten(2).split(sizes, 1)
         outputs = []
         for i, (x, y) in enumerate(zip(streams, result)):
@@ -115,6 +124,63 @@ def forward(self, streams, times, valid):
             outputs.append(x.masked_fill(~valid[i][..., None], 0))
         return outputs
 
+    def _apply_attention(self, query, key, value, mask):
+        # Preserve historical default arithmetic and autocast behavior exactly.
+        if self.attention_backend == "sdpa" and self.attention_dtype == "native":
+            return self._attend(query, key, value, mask)
+        original_dtype = query.dtype
+        with torch.autocast(device_type=query.device.type, enabled=False):
+            if self.attention_dtype == "bfloat16":
+                query, key, value = (x.to(torch.bfloat16) for x in (query, key, value))
+            if self.attention_backend == "flash4":
+                result = self._flash4_attend(query, key, value, mask)
+            elif self.attention_backend == "efficient":
+                if query.device.type != "cuda":
+                    raise RuntimeError("efficient attention requires CUDA; no math fallback")
+                from torch.nn.attention import SDPBackend, sdpa_kernel
+                with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
+                    result = self._attend(query, key, value, mask)
+            else:
+                result = self._attend(query, key, value, mask)
+        return result.to(original_dtype)
+
+    @staticmethod
+    def _flash4_attend(query, key, value, mask):
+        # Batch-one support is explicit until packed variable-length batches are
+        # validated. Removing padded keys preserves full joint attention,
+        # including noncontiguous padding; no queries or valid keys are dropped.
+        if query.device.type != "cuda":
+            raise RuntimeError("flash4 attention requires a CUDA Hopper/Blackwell GPU")
+        if query.shape[0] != 1 or mask.dtype != torch.bool or mask.shape != (1, 1, 1, key.shape[2]):
+            raise ValueError("flash4 currently requires batch1 and a boolean key-padding mask")
+        if any(x.dtype != torch.bfloat16 for x in (query, key, value)):
+            raise ValueError("flash4 requires bfloat16 Q/K/V")
+        from packaging.version import Version
+        if Version(torch.__version__.split("+")[0]) < Version("2.9"):
+            raise RuntimeError("flash4 profile requires the separately verified Torch>=2.9 runtime; see docs/WORLD_MODEL.md")
+        if torch.cuda.get_device_capability(query.device)[0] < 9:
+            raise RuntimeError("flash4 profile requires Hopper or newer CUDA hardware")
+        try:
+            from flash_attn.cute import flash_attn_func
+        except (ImportError, AttributeError) as error:
+            raise RuntimeError("Install the verified optional FlashAttention4 runtime; no silent backend fallback") from error
+        valid_keys = mask[0, 0, 0].nonzero(as_tuple=True)[0]
+        if valid_keys.numel() == 0:
+            raise ValueError("flash4 attention requires at least one valid key")
+        key = key.index_select(2, valid_keys)
+        value = value.index_select(2, valid_keys)
+        q, k, v = (x.transpose(1, 2).contiguous() for x in (query, key, value))
+        result = flash_attn_func(q, k, v, causal=False, return_lse=False)
+        # The tested 4.0.0b35 API returns (output, optional_lse). Reject unknown
+        # containers/shapes instead of assuming the README illustration's API.
+        if isinstance(result, tuple):
+            if len(result) != 2:
+                raise RuntimeError("Unsupported FlashAttention4 result tuple")
+            result = result[0]
+        if not torch.is_tensor(result) or result.shape != q.shape or result.dtype != q.dtype or result.device != q.device:
+            raise RuntimeError("Unsupported FlashAttention4 output shape/dtype/device")
+        return result.transpose(1, 2)
+
     @staticmethod
     def _attend(query, key, value, mask):
         return F.scaled_dot_product_attention(
@@ -154,11 +220,18 @@ def __init__(
         image_prediction_type="v",
         attention_query_chunk_size=0,
         activation_checkpointing=False,
+        attention_backend="sdpa",
+        attention_dtype="native",
+        strict_fp32=False,
     ):
         super().__init__()
         if type(activation_checkpointing) is not bool:
             raise ValueError("activation_checkpointing must be bool")
         self.activation_checkpointing = activation_checkpointing
+        if type(strict_fp32) is not bool:
+            raise ValueError("strict_fp32 must be a boolean")
+        if strict_fp32 and attention_dtype != "native":
+            raise ValueError("Explicit BF16 attention conflicts with strict_fp32; select the mixed-attention execution profile")
         if image_prediction_type not in ("v", "x0"):
             raise ValueError("image_prediction_type must be v or x0")
         self.image_prediction_type = image_prediction_type
@@ -203,7 +276,8 @@ def __init__(
             nn.Linear(2, width), nn.SiLU(), nn.Linear(width, width)
         )
         self.blocks = nn.ModuleList(
-            MoTBlock(width, heads, mlp_ratio, attention_query_chunk_size)
+            MoTBlock(width, heads, mlp_ratio, attention_query_chunk_size,
+                     attention_backend, attention_dtype)
             for _ in range(depth)
         )
         # One private image adapter per embodiment per block. Zero disables the

egomimic/trainHydra.py

@@ -43,17 +43,20 @@
 
 
 def _configure_strict_fp32(cfg: DictConfig) -> None:
-    """Opt in to full FP32 math, including CUDA/cuDNN backend precision."""
+    """Configure FP32/TF32 backend policy; explicit model dtype stays declared."""
     enabled = OmegaConf.select(cfg, "runtime.strict_fp32", default=False)
     if not isinstance(enabled, bool):
         raise ValueError("runtime.strict_fp32 must be a boolean")
-    if not enabled:
+    disable_tf32 = OmegaConf.select(cfg, "runtime.disable_tf32", default=False)
+    if not isinstance(disable_tf32, bool):
+        raise ValueError("runtime.disable_tf32 must be a boolean")
+    if not enabled and not disable_tf32:
         return
     precision = OmegaConf.select(cfg, "trainer.precision")
-    if str(precision) not in {"32", "32-true"}:
+    if enabled and str(precision) not in {"32", "32-true"}:
         raise ValueError("runtime.strict_fp32 requires trainer.precision=32-true")
     if os.environ.get("NVIDIA_TF32_OVERRIDE") != "0":
-        raise ValueError("runtime.strict_fp32 requires NVIDIA_TF32_OVERRIDE=0")
+        raise ValueError("FP32/disable_tf32 backend policy requires NVIDIA_TF32_OVERRIDE=0")
     torch.set_float32_matmul_precision("highest")
     torch.backends.cuda.matmul.allow_tf32 = False
     torch.backends.cudnn.allow_tf32 = False

scripts/benchmarks/profile_world_model.py

@@ -53,6 +53,9 @@ def main():
     parser.add_argument("--pose-dims", type=json.loads, required=True)
     parser.add_argument("--embodiment", required=True)
     parser.add_argument("--query-chunk", type=int, default=512)
+    parser.add_argument("--activation-checkpointing", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--attention-backend", choices=["sdpa", "efficient", "flash4"], default="sdpa")
+    parser.add_argument("--attention-dtype", choices=["native", "bfloat16"], default="native")
     parser.add_argument("--width", type=int, default=384)
     parser.add_argument("--depth", type=int, default=8)
     parser.add_argument("--heads", type=int, default=6)
@@ -76,7 +79,7 @@ def main():
     report = {
         "scope": "Synthetic memory/backward diagnostic; not actual data, optimizer or reload proof",
         "torch": torch.__version__, "device": torch.cuda.get_device_name(0) if cuda else "cpu",
-        "precision": "FP32; autocast off; TF32 off", "results": [],
+        "precision": f"FP32 parameters/model/loss; attention={args.attention_dtype}; autocast/TF32 off", "results": [],
         "arguments": {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()},
         "model_sha256": hashlib.sha256(Path(inspect.getfile(MoTWorldModel)).read_bytes()).hexdigest(),
     }
@@ -95,7 +98,8 @@ def main():
             model = MoTWorldModel(
                 args.pose_dims, width=args.width, depth=args.depth, heads=args.heads,
                 max_cameras=args.cameras, max_tracks=args.tracks, max_tokens=tokens,
-                attention_query_chunk_size=args.query_chunk, activation_checkpointing=True,
+                attention_query_chunk_size=args.query_chunk, activation_checkpointing=args.activation_checkpointing,
+                attention_backend=args.attention_backend, attention_dtype=args.attention_dtype,
             ).float().to(args.device)
             batch = synthetic_batch(frames, args.cameras, patches, args.tracks,
                                     args.pose_dims[args.embodiment], args.device)

tests/test_train_hydra_slurm_environment.py

@@ -321,3 +321,17 @@ def test_lightning_owned_restart_retains_last_checkpoint_fallback(
     expected = str(tmp_path / "lightning-root" / "checkpoints" / "last.ckpt")
     assert selected == expected
     assert cfg.ckpt_path == expected
+
+
+def test_disable_tf32_is_independent_of_explicit_model_precision(monkeypatch, preserve_precision_flags):
+    monkeypatch.setenv("NVIDIA_TF32_OVERRIDE", "0")
+    torch.set_float32_matmul_precision("high")
+    torch.backends.cudnn.allow_tf32 = True
+    train_hydra._configure_strict_fp32(OmegaConf.create({"runtime": {"strict_fp32": False, "disable_tf32": True}, "trainer": {"precision": "32-true"}}))
+    assert torch.get_float32_matmul_precision() == "highest"
+    assert not torch.backends.cuda.matmul.allow_tf32 and not torch.backends.cudnn.allow_tf32
+
+
+def test_disable_tf32_rejects_nonboolean():
+    with pytest.raises(ValueError, match="disable_tf32"):
+        train_hydra._configure_strict_fp32(OmegaConf.create({"runtime": {"disable_tf32": "true"}}))