Use fp32 for DeMemWM Plucker geometry

.exp_artifact/dememwm_remaining_fix_plan.md

@@ -285,7 +285,7 @@ Validate DeMemWM DiT stream boundaries
 
 ## Substep 5: Compute Plucker Geometry In fp32 Under Mixed Precision
 
-Status: `[ ]`
+Status: `[x]`
 
 Bug:
 

algorithms/dememwm/models/dit.py

@@ -645,7 +645,7 @@ class DiT(nn.Module):
             )
             return None
 
-        pose = frame_memory_pose.to(device=device, dtype=dtype)
+        pose = frame_memory_pose.to(device=device)
         if pose.ndim != 3 or pose.shape[-1] != 5:
             _warn_frame_memory_geometry_once(
                 "malformed_pose_metadata",
@@ -665,7 +665,7 @@ class DiT(nn.Module):
             )
             return None
         image_hw = image_hw if torch.is_tensor(image_hw) else torch.as_tensor(image_hw)
-        image_hw = image_hw.to(device=device, dtype=dtype)
+        image_hw = image_hw.to(device=device)
         if image_hw.ndim == 1:
             image_hw = image_hw.unsqueeze(0)
         if image_hw.ndim != 2 or image_hw.shape[-1] < 2:
@@ -683,54 +683,64 @@ class DiT(nn.Module):
             )
             return None
         image_hw = image_hw.contiguous()
-        image_h, image_w = image_hw[:, 0].view(B, 1, 1), image_hw[:, 1].view(B, 1, 1)
-
-        # Token grid sets position count; image_hw keeps camera scaling in original image coordinates.
-        y = (torch.arange(grid_h, device=device, dtype=dtype).view(1, grid_h, 1) + 0.5) * (image_h / grid_h)
-        x = (torch.arange(grid_w, device=device, dtype=dtype).view(1, 1, grid_w) + 0.5) * (image_w / grid_w)
-        fx, fy = 0.35 * image_w, 0.35 * image_h
-        cx, cy = 0.5 * image_w, 0.5 * image_h
-        directions = torch.stack(
-            (
-                (-(x - cx) / fx).expand(-1, grid_h, -1),
-                (-(y - cy) / fy).expand(-1, -1, grid_w),
-                torch.ones((B, grid_h, grid_w), device=device, dtype=dtype),
-            ),
-            dim=-1,
-        )
-        directions = F.normalize(directions, dim=-1)
-        query_rays = torch.cat([torch.zeros_like(directions), directions], dim=-1)[:, None].expand(-1, target_frames, -1, -1, -1)
-        intrinsics = torch.stack((fx[:, 0, 0], fy[:, 0, 0], cx[:, 0, 0], cy[:, 0, 0]), dim=-1)
-
-        c2w = self._frame_memory_pose_to_c2w(pose)
-        target_pose = pose[:, :target_frames]
-        target_c2w = c2w[:, :target_frames]
-        target_rot = target_c2w[:, :, :3, :3].transpose(-1, -2)
-        target_t = target_c2w[:, :, :3, 3]
-        target_w2c = torch.eye(4, device=device, dtype=dtype).view(1, 1, 4, 4).repeat(B, target_frames, 1, 1)
-        target_w2c[:, :, :3, :3] = target_rot
-        target_w2c[:, :, :3, 3] = -torch.matmul(target_rot, target_t.unsqueeze(-1)).squeeze(-1)
-
-        cache = {"query_rays": query_rays, "target_pose": target_pose, "image_hw": image_hw, "intrinsics": intrinsics}
-        cursor = target_frames
-        for stream_name in ("anchor", "dynamic", "revisit"):
-            stream_len = int(frame_memory_segments.get(stream_name, 0))
-            stream_pose = pose[:, cursor:cursor + stream_len]
-            relative_c2w = torch.matmul(target_w2c[:, :, None], c2w[:, None, cursor:cursor + stream_len])
-            rays_d = torch.einsum("bhwj,btmij->btmhwi", directions, relative_c2w[..., :3, :3])
-            rays_d = F.normalize(rays_d, dim=-1)
-            rays_o = relative_c2w[..., :3, 3][:, :, :, None, None, :].expand(-1, -1, -1, grid_h, grid_w, -1)
-            stream_mask = None if frame_memory_masks is None or frame_memory_masks.get(stream_name) is None else frame_memory_masks[stream_name].to(device=device, dtype=torch.bool)
-            cache[stream_name] = {
-                "query_rays": query_rays,
-                "relative_pose": stream_pose[:, None] - target_pose[:, :, None],
-                "relative_c2w": relative_c2w,
-                "relative_rays": torch.cat([torch.linalg.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1),
-                "image_hw": image_hw,
-                "intrinsics": intrinsics,
-                "mask": stream_mask,
+        # Mixed precision can round camera offsets before relative poses/rays are
+        # formed, so keep this local geometry block in fp32 and only cast caches.
+        with torch.autocast(device_type=device.type, enabled=False):
+            pose = pose.to(dtype=torch.float32)
+            image_hw = image_hw.to(dtype=torch.float32)
+            image_h, image_w = image_hw[:, 0].view(B, 1, 1), image_hw[:, 1].view(B, 1, 1)
+
+            # Token grid sets position count; image_hw keeps camera scaling in original image coordinates.
+            y = (torch.arange(grid_h, device=device, dtype=torch.float32).view(1, grid_h, 1) + 0.5) * (image_h / grid_h)
+            x = (torch.arange(grid_w, device=device, dtype=torch.float32).view(1, 1, grid_w) + 0.5) * (image_w / grid_w)
+            fx, fy = 0.35 * image_w, 0.35 * image_h
+            cx, cy = 0.5 * image_w, 0.5 * image_h
+            directions = torch.stack(
+                (
+                    (-(x - cx) / fx).expand(-1, grid_h, -1),
+                    (-(y - cy) / fy).expand(-1, -1, grid_w),
+                    torch.ones((B, grid_h, grid_w), device=device, dtype=torch.float32),
+                ),
+                dim=-1,
+            )
+            directions = F.normalize(directions, dim=-1)
+            query_rays = torch.cat([torch.zeros_like(directions), directions], dim=-1)[:, None].expand(-1, target_frames, -1, -1, -1)
+            intrinsics = torch.stack((fx[:, 0, 0], fy[:, 0, 0], cx[:, 0, 0], cy[:, 0, 0]), dim=-1)
+
+            c2w = self._frame_memory_pose_to_c2w(pose)
+            target_pose = pose[:, :target_frames]
+            target_c2w = c2w[:, :target_frames]
+            target_rot = target_c2w[:, :, :3, :3].transpose(-1, -2)
+            target_t = target_c2w[:, :, :3, 3]
+            target_w2c = torch.eye(4, device=device, dtype=torch.float32).view(1, 1, 4, 4).repeat(B, target_frames, 1, 1)
+            target_w2c[:, :, :3, :3] = target_rot
+            target_w2c[:, :, :3, 3] = -torch.matmul(target_rot, target_t.unsqueeze(-1)).squeeze(-1)
+
+            cache = {
+                "query_rays": query_rays.to(dtype=dtype),
+                "target_pose": target_pose.to(dtype=dtype),
+                "image_hw": image_hw.to(dtype=dtype),
+                "intrinsics": intrinsics.to(dtype=dtype),
             }
-            cursor += stream_len
+            cursor = target_frames
+            for stream_name in ("anchor", "dynamic", "revisit"):
+                stream_len = int(frame_memory_segments.get(stream_name, 0))
+                stream_pose = pose[:, cursor:cursor + stream_len]
+                relative_c2w = torch.matmul(target_w2c[:, :, None], c2w[:, None, cursor:cursor + stream_len])
+                rays_d = torch.einsum("bhwj,btmij->btmhwi", directions, relative_c2w[..., :3, :3])
+                rays_d = F.normalize(rays_d, dim=-1)
+                rays_o = relative_c2w[..., :3, 3][:, :, :, None, None, :].expand(-1, -1, -1, grid_h, grid_w, -1)
+                stream_mask = None if frame_memory_masks is None or frame_memory_masks.get(stream_name) is None else frame_memory_masks[stream_name].to(device=device, dtype=torch.bool)
+                cache[stream_name] = {
+                    "query_rays": query_rays.to(dtype=dtype),
+                    "relative_pose": (stream_pose[:, None] - target_pose[:, :, None]).to(dtype=dtype),
+                    "relative_c2w": relative_c2w.to(dtype=dtype),
+                    "relative_rays": torch.cat([torch.linalg.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1).to(dtype=dtype),
+                    "image_hw": image_hw.to(dtype=dtype),
+                    "intrinsics": intrinsics.to(dtype=dtype),
+                    "mask": stream_mask,
+                }
+                cursor += stream_len
         return cache
 
     def forward(self, x, t, action_cond=None, pose_cond=None, current_frame=None, mode=None, 

tests/test_dememwm_temporal_attention.py

@@ -373,6 +373,56 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         self.assertTrue(torch.equal(anchor_cache["relative_pose"][0, :, 0, 0], torch.tensor([2.0, 1.0])))
         self.assertTrue(torch.equal(spies[0][1].geometry[0]["mask"], masks["dynamic"]))
 
+    def test_frame_memory_geometry_uses_fp32_math_then_returns_activation_dtype(self):
+        model = DiT(
+            input_h=2,
+            input_w=2,
+            patch_size=1,
+            in_channels=1,
+            hidden_size=8,
+            depth=1,
+            num_heads=1,
+            mlp_ratio=1.0,
+            action_cond_dim=3,
+            pose_cond_dim=0,
+            reference_length=0,
+            use_plucker=True,
+            use_memory_attention=True,
+        )
+        device = torch.device("cpu")
+        segments = {"target": 1, "anchor": 1, "dynamic": 0, "revisit": 0}
+        frame_memory_pose = torch.zeros((1, 2, 5), device=device, dtype=torch.float32)
+        frame_memory_pose[0, 0, 0] = 10000.0
+        frame_memory_pose[0, 1, 0] = 10001.0
+        image_hw = torch.tensor([[4, 8]], device=device, dtype=torch.long)
+
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+            cache = model._build_frame_memory_geometry(
+                segments,
+                None,
+                frame_memory_pose,
+                image_hw,
+                grid_h=2,
+                grid_w=2,
+                device=device,
+                dtype=torch.float16,
+            )
+
+        anchor_cache = cache["anchor"]
+        for tensor in (
+            cache["query_rays"],
+            cache["target_pose"],
+            cache["image_hw"],
+            cache["intrinsics"],
+            anchor_cache["relative_pose"],
+            anchor_cache["relative_c2w"],
+            anchor_cache["relative_rays"],
+        ):
+            self.assertEqual(tensor.dtype, torch.float16)
+            self.assertTrue(torch.isfinite(tensor).all().item())
+        self.assertEqual(anchor_cache["relative_pose"][0, 0, 0, 0].item(), 1.0)
+        self.assertEqual(anchor_cache["relative_c2w"][0, 0, 0, 0, 3].item(), 1.0)
+
     def test_frame_memory_reference_attention_masks_padded_keys(self):
         attn = FrameMemoryReferenceAttention(hidden_size=1, num_heads=1)
         with torch.no_grad():