From aa90097fb68c5d74ff22b617ea010d8ed9ff99c6 Mon Sep 17 00:00:00 2001 From: monzzzz Date: Sat, 17 Oct 2026 16:05:41 -0700 Subject: [PATCH 1/3] Add selectable full-attention execution profiles for MoT --- docs/WORLD_MODEL.md | 46 ++++++++ .../hydra_configs/model/world_model/base.yaml | 9 +- .../hydra_configs/train_zarr_cartesian.yaml | 2 + .../world_model_attention/flash4_bf16.yaml | 12 +++ .../world_model_attention/native.yaml | 7 ++ .../world_model_attention/unchunked_fp32.yaml | 12 +++ egomimic/models/world_model.py | 84 ++++++++++++++- egomimic/trainHydra.py | 11 +- scripts/benchmarks/profile_world_model.py | 8 +- tests/test_train_hydra_slurm_environment.py | 14 +++ tests/test_world_model.py | 100 ++++++++++++++++++ 11 files changed, 291 insertions(+), 14 deletions(-) create mode 100644 egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml create mode 100644 egomimic/hydra_configs/world_model_attention/native.yaml create mode 100644 egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml diff --git a/docs/WORLD_MODEL.md b/docs/WORLD_MODEL.md index 5516702ca..288a2982a 100644 --- a/docs/WORLD_MODEL.md +++ b/docs/WORLD_MODEL.md @@ -427,3 +427,49 @@ 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 + +The default execution preserves native SDPA arithmetic. +Existing checkpoint-bound sources and resolved configs remain unchanged. +For either continuous or discrete experiments, select one of: + +```text ++world_model_attention=unchunked_fp32 ++world_model_attention=flash4_bf16 +``` + +Add the execution profile after selecting the experiment so it overrides root +runtime defaults. Inconsistent explicit strict-FP32/BF16 declarations are rejected. + +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. diff --git a/egomimic/hydra_configs/model/world_model/base.yaml b/egomimic/hydra_configs/model/world_model/base.yaml index 846c0508d..3b580cc69 100644 --- a/egomimic/hydra_configs/model/world_model/base.yaml +++ b/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} diff --git a/egomimic/hydra_configs/train_zarr_cartesian.yaml b/egomimic/hydra_configs/train_zarr_cartesian.yaml index 80ddea791..2204af98d 100644 --- a/egomimic/hydra_configs/train_zarr_cartesian.yaml +++ b/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 diff --git a/egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml b/egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml new file mode 100644 index 000000000..9478274c6 --- /dev/null +++ b/egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml @@ -0,0 +1,12 @@ +# @package _global_ +defaults: + - unchunked_fp32 + - _self_ +# Opt-in 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 +# Intentional BF16 attention is not a strict-all-FP32 execution contract. +runtime: + strict_fp32: false diff --git a/egomimic/hydra_configs/world_model_attention/native.yaml b/egomimic/hydra_configs/world_model_attention/native.yaml new file mode 100644 index 000000000..7d19ef2b4 --- /dev/null +++ b/egomimic/hydra_configs/world_model_attention/native.yaml @@ -0,0 +1,7 @@ +# @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 diff --git a/egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml b/egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml new file mode 100644 index 000000000..f0f5cfc3a --- /dev/null +++ b/egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml @@ -0,0 +1,12 @@ +# @package _global_ +defaults: + - native + - _self_ +# Exact full fused efficient attention, FP32 when the trainer is 32-true. +world_model_attention: + backend: efficient +trainer: + precision: 32-true +runtime: + strict_fp32: true + disable_tf32: true diff --git a/egomimic/models/world_model.py b/egomimic/models/world_model.py index 35f73aa35..b9def84db 100644 --- a/egomimic/models/world_model.py +++ b/egomimic/models/world_model.py @@ -55,7 +55,8 @@ 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") @@ -63,6 +64,14 @@ def __init__(self, width, heads, mlp_ratio=4, attention_query_chunk_size=0): 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 diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 1eb69691e..b1e605c32 100644 --- a/egomimic/trainHydra.py +++ b/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 diff --git a/scripts/benchmarks/profile_world_model.py b/scripts/benchmarks/profile_world_model.py index 5f4fa87b4..66178eb92 100644 --- a/scripts/benchmarks/profile_world_model.py +++ b/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) diff --git a/tests/test_train_hydra_slurm_environment.py b/tests/test_train_hydra_slurm_environment.py index 5a700abef..1d166fadc 100644 --- a/tests/test_train_hydra_slurm_environment.py +++ b/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"}})) diff --git a/tests/test_world_model.py b/tests/test_world_model.py index 4684c767d..4e8bf2ad9 100644 --- a/tests/test_world_model.py +++ b/tests/test_world_model.py @@ -1214,3 +1214,103 @@ def test_pipeline_exposes_missing_track_policy_and_cost(): assert torch.isfinite(result["loss/world"]) assert result["log/world_imputed_track_tokens"] == 2 assert result["log/world_imputation_forwards"] == 3 + + +@pytest.mark.parametrize("backend,dtype", [("unknown", "native"), ("sdpa", "half"), ("flash4", "native")]) +def test_attention_execution_rejects_unknown_or_implicit_precision(backend, dtype): + with pytest.raises(ValueError): + make_model(attention_backend=backend, attention_dtype=dtype) + + +def test_attention_execution_has_no_new_weights_and_native_default_is_exact(): + torch.manual_seed(42) + original = make_model() + explicit = make_model(attention_backend="sdpa", attention_dtype="native") + explicit.load_state_dict(original.state_dict(), strict=True) + assert list(original.state_dict()) == list(explicit.state_dict()) + batch = make_batch() + values = {n: batch[n]["value"] for n in ("image", "track")} + levels = {n: torch.zeros_like(batch[n]["time"]) for n in values} + expected, actual = original(batch, "robot", values, levels), explicit(batch, "robot", values, levels) + for key in expected: + assert torch.equal(expected[key], actual[key]) + + +def test_attention_only_bf16_matches_explicit_casts_and_keeps_fp32_gradients(): + from egomimic.models.world_model import MoTBlock + import torch.nn.functional as functional + torch.manual_seed(7) + block = MoTBlock(24, 3, attention_dtype="bfloat16") + q = torch.randn(1, 3, 7, 8, requires_grad=True) + k = torch.randn(1, 3, 9, 8, requires_grad=True) + v = torch.randn_like(k, requires_grad=True) + mask = torch.ones(1, 1, 1, 9, dtype=torch.bool) + mask[..., 1::3] = False + actual = block._apply_attention(q, k, v, mask) + expected = functional.scaled_dot_product_attention(q.bfloat16(), k.bfloat16(), v.bfloat16(), attn_mask=mask).float() + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + for grad in torch.autograd.grad(actual.square().mean(), [q, k, v]): + assert grad.dtype == torch.float32 and torch.isfinite(grad).all() + + +def test_optional_gpu_attention_does_not_silently_fallback_on_cpu(): + from egomimic.models.world_model import MoTBlock + q = torch.zeros(1, 3, 2, 8) + mask = torch.ones(1, 1, 1, 2, dtype=torch.bool) + for backend, dtype in (("efficient", "native"), ("flash4", "bfloat16")): + with pytest.raises(RuntimeError, match="CUDA"): + MoTBlock(24, 3, attention_backend=backend, attention_dtype=dtype)._apply_attention(q, q, q, mask) + + +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +@pytest.mark.parametrize("profile,backend,dtype", [("native", "sdpa", "native"), ("unchunked_fp32", "efficient", "native"), ("flash4_bf16", "flash4", "bfloat16")]) +def test_attention_execution_profiles_preserve_science(mode, profile, backend, dtype): + from pathlib import Path + from hydra import compose, initialize_config_dir + from omegaconf import OmegaConf + root = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" + def config(selected): + with initialize_config_dir(config_dir=str(root), version_base=None): + return compose(config_name="train_zarr_cartesian", overrides=[f"+experiment=world_model/{mode}", "hydra/launcher=basic", f"+world_model_attention={selected}"]) + cfg, baseline = config(profile), config("native") + model = cfg.model.pipeline.stages[0].model + assert model.attention_backend == backend and model.attention_dtype == dtype + assert model.attention_query_chunk_size == 0 and not model.activation_checkpointing + assert OmegaConf.to_container(cfg.data, resolve=True) == OmegaConf.to_container(baseline.data, resolve=True) + assert cfg.model.optimizer.lr == baseline.model.optimizer.lr == .0001 + assert cfg.model.pipeline.stages[0].warmup_steps == 10000 + assert cfg.trainer.precision == "32-true" + if profile == "flash4_bf16": + assert not cfg.runtime.strict_fp32 and cfg.runtime.disable_tf32 + + +def test_attention_precision_contract_rejects_bf16_labeled_strict_fp32(): + with pytest.raises(ValueError, match="conflicts with strict_fp32"): + make_model(attention_backend="flash4", attention_dtype="bfloat16", strict_fp32=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Optional real Flash4 CUDA proof") +def test_flash4_real_padding_and_fp32_gradient_contract(): + from packaging.version import Version + if Version(torch.__version__.split("+")[0]) < Version("2.9") or torch.cuda.get_device_capability(0)[0] < 9: + pytest.skip("Optional Flash4 profile needs Torch2.9 and Hopper") + pytest.importorskip("flash_attn.cute") + from egomimic.models.world_model import MoTBlock + from torch.nn.attention import sdpa_kernel, SDPBackend + torch.manual_seed(42) + q = torch.randn(1, 6, 97, 64, device="cuda", requires_grad=True) + k = torch.randn(1, 6, 137, 64, device="cuda", requires_grad=True) + v = torch.randn_like(k, requires_grad=True) + mask = torch.rand(1, 1, 1, 137, device="cuda") > .2 + block = MoTBlock(384, 6, attention_backend="flash4", attention_dtype="bfloat16") + actual = block._apply_attention(q, k, v, mask) + with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION): + expected = torch.nn.functional.scaled_dot_product_attention(q.bfloat16(), k.bfloat16(), v.bfloat16(), attn_mask=mask).float() + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=.02, atol=.002) + actual_grads = torch.autograd.grad(actual.square().mean(), [q,k,v]) + expected_grads = torch.autograd.grad(expected.square().mean(), [q,k,v]) + for a, e in zip(actual_grads, expected_grads): + assert a.dtype == torch.float32 and torch.isfinite(a).all() + assert float((a-e).norm()/e.norm().clamp_min(1e-30)) < .02 From 01961df46140d7d807a7fe226bb3d9595b896f9d Mon Sep 17 00:00:00 2001 From: monzzzz Date: Sat, 17 Oct 2026 16:14:08 -0700 Subject: [PATCH 2/3] Normalize attention test import ordering --- tests/test_world_model.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_world_model.py b/tests/test_world_model.py index 4e8bf2ad9..db0757035 100644 --- a/tests/test_world_model.py +++ b/tests/test_world_model.py @@ -812,6 +812,7 @@ def test_validation_batch_cap_keeps_short_final_batch(): @pytest.mark.parametrize("cap", [0, -1, True, 1.5, "2"]) def test_validation_batch_cap_rejects_invalid_value(cap): from types import SimpleNamespace + from egomimic.rldb.world_model_datamodule import WorldModelDataModule with pytest.raises(ValueError, match="positive integer"): HomogeneousWorldBatchSampler(SimpleNamespace(split="val"), shuffle=False, @@ -912,6 +913,7 @@ def test_hydra_both_recipes_and_required_schema(): def test_shared_lightning_fit_validate_and_resume(tmp_path): import inspect + from lightning import Trainer from omegaconf import OmegaConf @@ -1237,8 +1239,9 @@ def test_attention_execution_has_no_new_weights_and_native_default_is_exact(): def test_attention_only_bf16_matches_explicit_casts_and_keeps_fp32_gradients(): - from egomimic.models.world_model import MoTBlock import torch.nn.functional as functional + + from egomimic.models.world_model import MoTBlock torch.manual_seed(7) block = MoTBlock(24, 3, attention_dtype="bfloat16") q = torch.randn(1, 3, 7, 8, requires_grad=True) @@ -1267,6 +1270,7 @@ def test_optional_gpu_attention_does_not_silently_fallback_on_cpu(): @pytest.mark.parametrize("profile,backend,dtype", [("native", "sdpa", "native"), ("unchunked_fp32", "efficient", "native"), ("flash4_bf16", "flash4", "bfloat16")]) def test_attention_execution_profiles_preserve_science(mode, profile, backend, dtype): from pathlib import Path + from hydra import compose, initialize_config_dir from omegaconf import OmegaConf root = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" @@ -1296,8 +1300,9 @@ def test_flash4_real_padding_and_fp32_gradient_contract(): if Version(torch.__version__.split("+")[0]) < Version("2.9") or torch.cuda.get_device_capability(0)[0] < 9: pytest.skip("Optional Flash4 profile needs Torch2.9 and Hopper") pytest.importorskip("flash_attn.cute") + from torch.nn.attention import SDPBackend, sdpa_kernel + from egomimic.models.world_model import MoTBlock - from torch.nn.attention import sdpa_kernel, SDPBackend torch.manual_seed(42) q = torch.randn(1, 6, 97, 64, device="cuda", requires_grad=True) k = torch.randn(1, 6, 137, 64, device="cuda", requires_grad=True) From cb95574e7f3dcdec8c0b0374aecbb1fefaf5f616 Mon Sep 17 00:00:00 2001 From: monzzzz Date: Sat, 17 Oct 2026 16:30:46 -0700 Subject: [PATCH 3/3] Default new MoT experiments to combined Flash4 attention --- docs/WORLD_MODEL.md | 20 +++++++++----- .../experiment/world_model/continuous.yaml | 8 +++++- .../world_model_attention/flash4_bf16.yaml | 3 ++- .../world_model_attention/native.yaml | 2 ++ .../world_model_attention/unchunked_fp32.yaml | 2 ++ tests/test_world_model.py | 27 ++++++++++++++++++- 6 files changed, 52 insertions(+), 10 deletions(-) diff --git a/docs/WORLD_MODEL.md b/docs/WORLD_MODEL.md index 288a2982a..41781012c 100644 --- a/docs/WORLD_MODEL.md +++ b/docs/WORLD_MODEL.md @@ -431,17 +431,23 @@ Do not treat this code milestone as a prepared four-domain training dataset. ## Selectable attention execution -The default execution preserves native SDPA arithmetic. -Existing checkpoint-bound sources and resolved configs remain unchanged. -For either continuous or discrete experiments, select one of: +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=flash4_bf16 +world_model_attention=unchunked_fp32 +world_model_attention=native ``` -Add the execution profile after selecting the experiment so it overrides root -runtime defaults. Inconsistent explicit strict-FP32/BF16 declarations are rejected. +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 diff --git a/egomimic/hydra_configs/experiment/world_model/continuous.yaml b/egomimic/hydra_configs/experiment/world_model/continuous.yaml index f023004ca..98b706246 100644 --- a/egomimic/hydra_configs/experiment/world_model/continuous.yaml +++ b/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} diff --git a/egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml b/egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml index 9478274c6..6377266f1 100644 --- a/egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml +++ b/egomimic/hydra_configs/world_model_attention/flash4_bf16.yaml @@ -2,11 +2,12 @@ defaults: - unchunked_fp32 - _self_ -# Opt-in attention-only BF16; keep trainer32-true and FP32 weights/FFNs/loss. +# 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 diff --git a/egomimic/hydra_configs/world_model_attention/native.yaml b/egomimic/hydra_configs/world_model_attention/native.yaml index 7d19ef2b4..631deb1c2 100644 --- a/egomimic/hydra_configs/world_model_attention/native.yaml +++ b/egomimic/hydra_configs/world_model_attention/native.yaml @@ -5,3 +5,5 @@ world_model_attention: dtype: native query_chunk_size: 0 activation_checkpointing: false + strict_fp32: false + disable_tf32: false diff --git a/egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml b/egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml index f0f5cfc3a..0d4de73c2 100644 --- a/egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml +++ b/egomimic/hydra_configs/world_model_attention/unchunked_fp32.yaml @@ -5,6 +5,8 @@ defaults: # 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: diff --git a/tests/test_world_model.py b/tests/test_world_model.py index db0757035..a23e12543 100644 --- a/tests/test_world_model.py +++ b/tests/test_world_model.py @@ -1276,7 +1276,7 @@ def test_attention_execution_profiles_preserve_science(mode, profile, backend, d root = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" def config(selected): with initialize_config_dir(config_dir=str(root), version_base=None): - return compose(config_name="train_zarr_cartesian", overrides=[f"+experiment=world_model/{mode}", "hydra/launcher=basic", f"+world_model_attention={selected}"]) + return compose(config_name="train_zarr_cartesian", overrides=[f"+experiment=world_model/{mode}", "hydra/launcher=basic", f"world_model_attention={selected}"]) cfg, baseline = config(profile), config("native") model = cfg.model.pipeline.stages[0].model assert model.attention_backend == backend and model.attention_dtype == dtype @@ -1289,6 +1289,31 @@ def config(selected): assert not cfg.runtime.strict_fp32 and cfg.runtime.disable_tf32 +@pytest.mark.parametrize("mode", ["continuous", "discrete"]) +def test_world_model_defaults_to_combined_attention_without_override(mode): + from pathlib import Path + + from hydra import compose, initialize_config_dir + from omegaconf import OmegaConf + + root = Path(__file__).resolve().parents[1] / "egomimic/hydra_configs" + overrides = [f"+experiment=world_model/{mode}", "hydra/launcher=basic"] + with initialize_config_dir(config_dir=str(root), version_base=None): + default = compose(config_name="train_zarr_cartesian", overrides=overrides) + explicit = compose(config_name="train_zarr_cartesian", overrides=[ + *overrides, "world_model_attention=flash4_bf16" + ]) + # Same complete recipe without requiring an attention command-line flag. + assert OmegaConf.to_container(default, resolve=False) == OmegaConf.to_container(explicit, resolve=False) + model = default.model.pipeline.stages[0].model + assert model.attention_backend == "flash4" + assert model.attention_dtype == "bfloat16" + assert model.attention_query_chunk_size == 0 + assert not model.activation_checkpointing and not model.strict_fp32 + assert not default.runtime.strict_fp32 and default.runtime.disable_tf32 + assert default.trainer.precision == "32-true" + + def test_attention_precision_contract_rejects_bf16_labeled_strict_fp32(): with pytest.raises(ValueError, match="conflicts with strict_fp32"): make_model(attention_backend="flash4", attention_dtype="bfloat16", strict_fp32=True)