Fix AMP precision policy

README.md

@@ -69,10 +69,11 @@ encoder.
   `num_steps=1` means one NFE.
 
 The export ships weights in `float32`. The recommended and default runtime path
-is `bfloat16` AMP for the main encoder, decoder, and class-token path, with
-`float32` retained for sensitive operations such as whitening/dewhitening,
-normalization math, RoPE frequency construction, and VP diffusion schedule
-helpers.
+is `bfloat16` AMP for the main encoder, decoder, and class-token path. The
+loader retains normalization affine parameters, GRN/residual gates, the final
+pixel projection, latent statistics, RoPE/time frequencies, sampler state, and
+whitening/dewhitening in `float32`. These tensors are loaded from the original
+FP32 weights before ordinary parameters are converted to BF16.
 
 ## Usage
 

dinac_ae/encoder.py

@@ -37,7 +37,7 @@ def _resolve_encoder_mlp_type(name: str) -> MLPType:
             return MLPType.RELU
         case _ as unreachable:
             raise ValueError(
-                "Unsupported encoder_mlp_type for DinacAE export: " f"{unreachable!r}"
+                f"Unsupported encoder_mlp_type for DinacAE export: {unreachable!r}"
             )
 
 
@@ -65,7 +65,7 @@ class EncoderPosterior:
     def mode(self) -> Tensor:
         """Return the posterior mode in token space."""
 
-        return (self.alpha * self.mean.to(torch.float32)).to(dtype=self.mean.dtype)
+        return self.alpha * self.mean.to(torch.float32)
 
     def sample(self, *, generator: torch.Generator | None = None) -> Tensor:
         """Sample from the posterior."""
@@ -77,7 +77,7 @@ class EncoderPosterior:
             dtype=torch.float32,
             generator=generator,
         )
-        return (self.alpha * mean_fp32 + self.sigma * eps).to(dtype=self.mean.dtype)
+        return self.alpha * mean_fp32 + self.sigma * eps
 
 
 class Encoder(nn.Module):

dinac_ae/model.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from pathlib import Path
 
 import torch
@@ -13,6 +14,7 @@ from dit.repa_projection import DinoTokenAlignmentHead
 from .config import DinacAEConfig, DinacAEInferenceConfig
 from .decoder import Decoder
 from .encoder import Encoder, EncoderPosterior
+from .precision import parameter_storage_dtype, validate_inference_storage
 from .samplers import run_ddim, run_dpmpp_2m
 from .vp_diffusion import get_schedule, make_initial_state, sample_noise
 
@@ -58,8 +60,7 @@ def _resolve_class_head_mlp_type(name: str) -> MLPType:
             return MLPType.RELU
         case _ as unreachable:
             raise ValueError(
-                "Unsupported class_head_mlp_type for DinacAE export: "
-                f"{unreachable!r}"
+                f"Unsupported class_head_mlp_type for DinacAE export: {unreachable!r}"
             )
 
 
@@ -111,27 +112,50 @@ class DinacAE(nn.Module):
             register_token_count=int(config.class_head_register_token_count),
         )
 
-    def _restore_float32_norm_buffers(self) -> None:
-        """Keep latent running stats in float32 after device/dtype moves."""
-
-        self.latent_norm_running_mean = self.latent_norm_running_mean.to(
-            dtype=torch.float32
-        )
-        self.latent_norm_running_var = self.latent_norm_running_var.to(
-            dtype=torch.float32
+    def _apply(
+        self,
+        fn: Callable[[Tensor], Tensor],
+        recurse: bool = True,
+    ) -> DinacAE:
+        """Move the model without rounding FP32 parameter or buffer islands."""
+
+        preserved_parameters = {
+            name: parameter.detach().float()
+            for name, parameter in self.named_parameters()
+            if parameter_storage_dtype(name, torch.bfloat16) is torch.float32
+        }
+        preserved_buffers = {
+            name: buffer.detach().float()
+            for name, buffer in self.named_buffers()
+            if buffer.is_floating_point()
+        }
+        super()._apply(fn, recurse=recurse)
+        for name, parameter in self.named_parameters():
+            preserved = preserved_parameters.get(name)
+            if preserved is not None:
+                parameter.data = preserved.to(device=parameter.device)
+        for name, buffer in self.named_buffers():
+            preserved = preserved_buffers.get(name)
+            if preserved is not None:
+                buffer.data = preserved.to(device=buffer.device)
+        return self
+
+    @property
+    def ordinary_weight_storage_dtype(self) -> torch.dtype:
+        """Return the ordinary-weight storage dtype."""
+
+        return self.decoder.patchify.proj.weight.dtype
+
+    def validate_inference_storage(self) -> None:
+        """Require the loaded model's complete mixed-precision policy."""
+
+        device = self.decoder.patchify.proj.weight.device
+        validate_inference_storage(
+            self,
+            base_dtype=self.ordinary_weight_storage_dtype,
+            device=device,
         )
 
-    def to(self, *args: object, **kwargs: object) -> DinacAE:
-        """Move the model while preserving float32 latent stats buffers."""
-
-        moved = super().to(*args, **kwargs)
-        if not isinstance(moved, DinacAE):
-            raise RuntimeError(
-                f"Expected DinacAE after nn.Module.to(), got {type(moved).__name__}"
-            )
-        moved._restore_float32_norm_buffers()
-        return moved
-
     @classmethod
     def from_pretrained(
         cls,
@@ -144,6 +168,8 @@ class DinacAE(nn.Module):
     ) -> DinacAE:
         """Load a pretrained export from a local directory or the Hub."""
 
+        if dtype is not torch.float32 and dtype is not torch.bfloat16:
+            raise ValueError(f"Unsupported DINAC-AE parameter storage dtype: {dtype}")
         model_dir = _resolve_model_dir(
             path_or_repo_id,
             revision=revision,
@@ -167,6 +193,7 @@ class DinacAE(nn.Module):
         model.load_state_dict(state_dict, strict=True)
         model = model.to(dtype=dtype, device=torch.device(device))
         model.eval()
+        model.validate_inference_storage()
         return model
 
     def _latent_norm_stats(self) -> tuple[Tensor, Tensor]:
@@ -210,9 +237,13 @@ class DinacAE(nn.Module):
             height=int(images.shape[2]),
             width=int(images.shape[3]),
         )
-        model_dtype = next(self.parameters()).dtype
-        latents = self.encoder(images.to(dtype=model_dtype))
-        return self.whiten(latents).to(dtype=model_dtype)
+        device = self.decoder.patchify.proj.weight.device
+        with torch.autocast(
+            device_type=device.type,
+            dtype=torch.bfloat16,
+        ):
+            latents = self.encoder(images.to(device=device, dtype=torch.bfloat16))
+        return self.whiten(latents)
 
     def encode_posterior(self, images: Tensor) -> EncoderPosterior:
         """Encode images and return the raw posterior."""
@@ -221,23 +252,34 @@ class DinacAE(nn.Module):
             height=int(images.shape[2]),
             width=int(images.shape[3]),
         )
-        model_dtype = next(self.parameters()).dtype
-        return self.encoder.encode_posterior(images.to(dtype=model_dtype))
+        device = self.decoder.patchify.proj.weight.device
+        with torch.autocast(
+            device_type=device.type,
+            dtype=torch.bfloat16,
+        ):
+            return self.encoder.encode_posterior(
+                images.to(device=device, dtype=torch.bfloat16)
+            )
 
     def predict_class(self, latents: Tensor) -> Tensor:
         """Predict the exported DINO class token from whitened latents."""
 
-        dewhitened = self.dewhiten(latents)
+        device = self.decoder.patchify.proj.weight.device
+        dewhitened = self.dewhiten(latents).to(
+            device=device,
+            dtype=torch.float32,
+        )
         t_zero = torch.zeros(
             (int(latents.shape[0]),),
-            device=latents.device,
+            device=device,
             dtype=torch.float32,
         )
-        head_dtype = self.dino_token_alignment_head.in_proj.weight.dtype
-        device_type = "cuda" if latents.device.type == "cuda" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):
+        with torch.autocast(
+            device_type=device.type,
+            dtype=torch.bfloat16,
+        ):
             out = self.dino_token_alignment_head(
-                dewhitened.to(device=latents.device, dtype=head_dtype),
+                dewhitened,
                 t=t_zero,
             )
         return out.class_token.to(torch.float32)
@@ -260,8 +302,10 @@ class DinacAE(nn.Module):
         self._require_image_size_divisible(height=int(height), width=int(width))
         batch = int(latents.shape[0])
         device = latents.device
-        model_dtype = next(self.parameters()).dtype
-        decoder_latents = self.dewhiten(latents).to(device=device, dtype=model_dtype)
+        decoder_latents = self.dewhiten(latents).to(
+            device=device,
+            dtype=torch.float32,
+        )
         noise = sample_noise(
             (batch, int(self.config.in_channels), int(height), int(width)),
             noise_std=float(self.config.pixel_noise_std),
@@ -277,7 +321,10 @@ class DinacAE(nn.Module):
             logsnr_max=float(self.config.logsnr_max),
         )
         device_type = "cuda" if device.type == "cuda" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):
+        with torch.autocast(
+            device_type=device_type,
+            dtype=torch.bfloat16,
+        ):
 
             def _forward_fn(
                 x_t: Tensor,
@@ -289,9 +336,9 @@ class DinacAE(nn.Module):
             ) -> Tensor:
                 _ = mask_latent_tokens
                 return self.decoder(
-                    x_t.to(dtype=model_dtype),
+                    x_t.to(dtype=torch.float32),
                     t,
-                    latents_in.to(dtype=model_dtype),
+                    latents_in.to(dtype=torch.float32),
                     drop_middle_blocks=bool(drop_middle_blocks),
                 )
 

dinac_ae/precision.py

@@ -0,0 +1,83 @@
+"""Frozen mixed-precision storage policy for DINAC-AE inference."""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+def _is_norm_parameter(name: str) -> bool:
+    """Return whether a parameter belongs to a normalization module."""
+
+    return any("norm" in component for component in name.split("."))
+
+
+def _is_gate_parameter(name: str) -> bool:
+    """Return whether a parameter is a learned residual or GRN gate."""
+
+    return (
+        name.endswith(".layer_scale")
+        or name.endswith(".grn.gamma")
+        or name.endswith(".grn.beta")
+        or name.endswith("decoder.path_drop_mask_feature")
+    )
+
+
+def _is_decoder_output_parameter(name: str) -> bool:
+    """Return whether a parameter belongs to the final pixel projection."""
+
+    return ".decoder.out_proj." in f".{name}."
+
+
+def parameter_storage_dtype(name: str, base_dtype: torch.dtype) -> torch.dtype:
+    """Return storage dtype for one frozen DINAC-AE parameter."""
+
+    if base_dtype is torch.float32:
+        return torch.float32
+    if base_dtype is not torch.bfloat16:
+        raise ValueError(f"Unsupported DINAC-AE parameter storage dtype: {base_dtype}")
+    preserve = (
+        _is_norm_parameter(name)
+        or _is_gate_parameter(name)
+        or _is_decoder_output_parameter(name)
+    )
+    return torch.float32 if preserve else torch.bfloat16
+
+
+def validate_inference_storage(
+    module: nn.Module,
+    *,
+    base_dtype: torch.dtype,
+    device: torch.device,
+) -> None:
+    """Require the exact mixed parameter and FP32-buffer inference policy."""
+
+    def _device_matches(actual: torch.device) -> bool:
+        """Return whether an actual device satisfies the requested device."""
+
+        return actual.type == device.type and (
+            device.index is None or actual.index == device.index
+        )
+
+    parameter_mismatches = tuple(
+        name
+        for name, parameter in module.named_parameters()
+        if not _device_matches(parameter.device)
+        or parameter.dtype is not parameter_storage_dtype(name, base_dtype)
+    )
+    if parameter_mismatches:
+        raise RuntimeError(
+            "DINAC-AE parameter storage does not match its inference policy: "
+            f"{parameter_mismatches}"
+        )
+    buffer_mismatches = tuple(
+        name
+        for name, buffer in module.named_buffers()
+        if not _device_matches(buffer.device)
+        or (buffer.is_floating_point() and buffer.dtype is not torch.float32)
+    )
+    if buffer_mismatches:
+        raise RuntimeError(
+            "DINAC-AE floating inference buffers must retain FP32 storage: "
+            f"{buffer_mismatches}"
+        )