From 7bea39d9027f6591204f648180438f2c62e03ca6 Mon Sep 17 00:00:00 2001 From: tritsystem Date: Fri, 16 Oct 2026 02:12:47 -0700 Subject: [PATCH] fix(layers): upcast the EMAQuantizer codebook with its inputs EMAQuantizer.quantize() casts the inputs to float32 but used the embedding weight in its own dtype, so a VQVAE cast to float16/bfloat16 with .half() or .to(dtype) failed in torch.mm ("expected m1 and m2 to have the same dtype"). Autocast was unaffected because the weights stay float32 there. Upcast the codebook as well; float32 models are unchanged. Assisted-by: Claude Opus 5.5 Signed-off-by: tritsystem --- monai/networks/layers/vector_quantizer.py | 6 ++++-- tests/networks/layers/test_vector_quantizer.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/monai/networks/layers/vector_quantizer.py b/monai/networks/layers/vector_quantizer.py index 388f93fe2df..30571ed652d 100644 --- a/monai/networks/layers/vector_quantizer.py +++ b/monai/networks/layers/vector_quantizer.py @@ -105,6 +105,8 @@ def quantize(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, to del encoding_indices_view[1] inputs = inputs.float() + # The codebook is upcast too, so a model cast to float16/bfloat16 (not just autocast) still works. + embedding_weight = self.embedding.weight.float() # Converting to channel last format flat_input = inputs.permute(self.flatten_permutation).contiguous().view(-1, self.embedding_dim) @@ -112,8 +114,8 @@ def quantize(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, to # Calculate Euclidean distances distances = ( (flat_input**2).sum(dim=1, keepdim=True) - + (self.embedding.weight.t() ** 2).sum(dim=0, keepdim=True) - - 2 * torch.mm(flat_input, self.embedding.weight.t()) + + (embedding_weight.t() ** 2).sum(dim=0, keepdim=True) + - 2 * torch.mm(flat_input, embedding_weight.t()) ) # Mapping distances to indexes diff --git a/tests/networks/layers/test_vector_quantizer.py b/tests/networks/layers/test_vector_quantizer.py index 43533d03771..54719642f66 100644 --- a/tests/networks/layers/test_vector_quantizer.py +++ b/tests/networks/layers/test_vector_quantizer.py @@ -68,6 +68,19 @@ def test_ema(self): self.assertTrue(all(layer.embedding.weight[0] != original_weight_0)) self.assertTrue(all(layer.embedding.weight[1] == original_weight_1)) + @parameterized.expand([[torch.float16], [torch.bfloat16]]) + def test_ema_half_precision_module(self, dtype): + """A quantizer cast to half precision (not just autocast) runs and picks the same codes as float32.""" + layer = EMAQuantizer(spatial_dims=2, num_embeddings=16, embedding_dim=8).eval() + x = torch.randn(1, 8, 4, 4) + expected_indices = layer.quantize(x)[2] + + layer = layer.to(dtype) + quantized, _, encoding_indices = layer(x.to(dtype)) + + self.assertEqual(quantized.dtype, dtype) + torch.testing.assert_close(encoding_indices, expected_indices) + class TestVectorQuantizer(unittest.TestCase): @parameterized.expand(TEST_CASES)