Preserve precision when unscaling Chronos-2 and Chronos-Bolt forecasts

src/chronos/chronos2/model.py

@@ -759,7 +759,12 @@ def forward(
             q=self.num_quantiles,
             h=num_output_patches * self.chronos_config.output_patch_size,
         )
-        quantile_preds = self.instance_norm.inverse(quantile_preds, loc_scale)
+        # Preserve precision when unscaling large-magnitude forecasts.
+        quantile_preds = self.instance_norm.inverse(
+            quantile_preds,
+            loc_scale,
+            output_dtype=torch.float32,
+        )
         quantile_preds = rearrange(
             quantile_preds,
             "b (q h) -> b q h",

src/chronos/chronos_bolt.py

@@ -121,7 +121,12 @@ def forward(
 
         return scaled_x.to(orig_dtype), (loc, scale)
 
-    def inverse(self, x: torch.Tensor, loc_scale: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    def inverse(
+        self,
+        x: torch.Tensor,
+        loc_scale: tuple[torch.Tensor, torch.Tensor],
+        output_dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
         orig_dtype = x.dtype
         x = x.to(torch.float32)
         loc, scale = loc_scale
@@ -131,7 +136,7 @@ def inverse(self, x: torch.Tensor, loc_scale: tuple[torch.Tensor, torch.Tensor])
 
         x = x * scale + loc
 
-        return x.to(orig_dtype)
+        return x.to(orig_dtype if output_dtype is None else output_dtype)
 
 
 class ResidualBlock(nn.Module):
@@ -385,6 +390,7 @@ def forward(
         quantile_preds = self.instance_norm.inverse(
             quantile_preds.view(batch_size, -1),
             loc_scale,
+            output_dtype=torch.float32,
         ).view(*quantile_preds_shape)
 
         return ChronosBoltOutput(

test/test_chronos2.py

@@ -44,6 +44,14 @@ def test_chronos2_encoder_accepts_config_without_is_decoder():
     assert len(encoder.block) == config.num_layers
 
 
+def test_when_chronos2_model_uses_bfloat16_then_unscaled_predictions_use_float32(pipeline):
+    pipeline.model.to(torch.bfloat16)
+
+    output = pipeline.model(context=torch.rand(1, 16), num_output_patches=1)
+
+    assert output.quantile_preds.dtype == torch.float32
+
+
 def test_base_chronos2_pipeline_loads_from_s3():
     BaseChronosPipeline.from_pretrained("s3://autogluon/chronos-2", device_map="cpu")
 

test/test_chronos_bolt.py

@@ -22,6 +22,14 @@ def pipeline() -> ChronosBoltPipeline:
     return BaseChronosPipeline.from_pretrained(DUMMY_MODEL_PATH, device_map="cpu")
 
 
+def test_when_chronos_bolt_model_uses_bfloat16_then_unscaled_predictions_use_float32():
+    pipeline = ChronosBoltPipeline.from_pretrained(DUMMY_MODEL_PATH, device_map="cpu", torch_dtype=torch.bfloat16)
+
+    output = pipeline.model(context=torch.rand(1, 16))
+
+    assert output.quantile_preds.dtype == torch.float32
+
+
 def test_base_chronos_pipeline_loads_from_huggingface():
     BaseChronosPipeline.from_pretrained("amazon/chronos-bolt-tiny", device_map="cpu")
 
@@ -355,6 +363,25 @@ def test_when_instancenorm_applied_and_reversed_then_output_correct():
     assert torch.allclose(output, input_)
 
 
+def test_when_instancenorm_reversed_to_float32_then_precision_is_preserved():
+    inorm = InstanceNorm()
+    normalized = torch.tensor([[0.125]], dtype=torch.bfloat16)
+    loc = torch.tensor([[1_000_000.0]], dtype=torch.float32)
+    scale = torch.tensor([[100.0]], dtype=torch.float32)
+
+    output = inorm.inverse(
+        normalized,
+        (loc, scale),
+        output_dtype=torch.float32,
+    )
+
+    assert output.dtype == torch.float32
+    torch.testing.assert_close(
+        output,
+        torch.tensor([[1_000_012.5]], dtype=torch.float32),
+    )
+
+
 @pytest.mark.parametrize("task_kwargs", [{}, {"eval_metric": "WQL", "quantile_levels": [0.1, 0.2]}])
 def test_pipeline_can_evaluate_on_dummy_fev_task(task_kwargs):
     pipeline = BaseChronosPipeline.from_pretrained(