Skip to content

fix: call .eval() on tokenizer and model in KronosPredictor - #382

Open
hexonal wants to merge 1 commit into
shiyu-coder:masterfrom
hexonal:fix/predictor-eval-mode
Open

fix: call .eval() on tokenizer and model in KronosPredictor#382
hexonal wants to merge 1 commit into
shiyu-coder:masterfrom
hexonal:fix/predictor-eval-mode

Conversation

@hexonal

@hexonal hexonal commented Aug 4, 2026

Copy link
Copy Markdown

KronosPredictor never switches the modules out of training mode, so inference runs with dropout active.

Impact

Kronos-small ships attn_dropout_p = 0.1. Because KronosPredictor.__init__ only does .to(device), the modules stay in train() mode for the whole lifetime of the predictor:

  • On MPS it crashes outright. F.scaled_dot_product_attention with a non-zero dropout_p is not implemented for the MPS backend, so any predict() on Apple Silicon raises rather than returning a forecast.
  • On CUDA/CPU it is worse than a crash — it is silent. Predictions come back looking fine, but every forward pass drops 10% of the attention weights at random. Two calls with identical inputs give different answers for a reason that has nothing to do with the models sampling behaviour, and the extra variance is indistinguishable from the intended sample_count sampling noise.

Measured on this fork: with sample_count=5 on 60m US equity bars, adding .eval() cut MAE by roughly a fifth on an otherwise unchanged pipeline.

Fix

self.tokenizer = self.tokenizer.to(self.device).eval()
self.model = self.model.to(self.device).eval()

.eval() is the conventional place for this — the predictor is inference-only, and nothing in the class ever calls .train(). Callers who genuinely want stochastic modules can still flip them back after construction.

Two lines, no behavioural change for anyone who was already calling .eval() themselves.

`KronosPredictor.__init__` 把 tokenizer 与 model 搬上设备后没有切到 eval 模式,
而 Kronos-small 的配置里 `attn_dropout_p=0.1`。后果分两种,都不好:

- Apple Silicon(MPS)上直接崩。`auto_regressive_inference` 整个包在
  `torch.no_grad()` 里,而 MPS 的 SDPA 不支持「dropout_p>0 且 grad 关闭」这个组合。
- CPU/CUDA 上不报错,但**带着 dropout 做推理** —— 输出带随机噪声、同一输入两次
  调用结果不同,而使用者完全看不出来。

改法是构造时就 `.eval()`。这不改变任何训练路径:`finetune/` 下的训练脚本自己
调 `.train()`,而 KronosPredictor 从设计上就只用于推理。
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant