Skip to content

v0.17.1: Fix in-training val metric for MTP

Choose a tag to compare

View all tags
@mrothroc mrothroc released this 01 May 19:51
· 521 commits to main since this release
v0.17.1
c6dc13d

v0.17.1: Fix in-training val metric for MTP

Patch release fixing a metric-selection bug introduced with v0.17.0's Multi-Token Prediction feature.

The bug

When MTP was active (mtp.n > 1), the val number printed during -mode arch training was the MTP weighted-sum-over-horizons training-objective loss, not the next-1-only horizon cross-entropy that's directly comparable to non-MTP runs.

Effect: a user comparing val numbers between MTP-on and MTP-off runs would see what looked like a +0.1 to +0.4 nat regression with MTP enabled — but the model's actual next-1-token BPB was unchanged or improved. The training was correct; only the displayed metric was wrong.

The IR layer was already correct: arch/builder.go::emitLanguageModelLossIR emits two distinct outputs when MTP is active — loss (training objective) and eval_loss (next-1-only). The C++ trainer (gpu/ir_trainer.cpp) had an existing helper preferring eval_loss for evaluation. The gap was at the Go boundary: train/gpu_trainer_mlx.go::EvaluateGPU didn't make the contract explicit, so different downstream callers could read either output.

The fix

mlxGPUTrainer now tracks the preferred eval scalar from declared IR outputs at construction time and returns eval_loss from EvaluateGPU when declared, falling back to loss for non-MTP configs.

Verified

  • Kitchen-sink config (with MTP n=3), 500 steps:
    • Fixed: val=7.2790
    • Forced-old (weighted-loss display): val=7.3863 (Δ +0.107 nat)
  • Non-MTP plain_3L: val=6.1426 before and after — fallback to loss works correctly.
  • Two new MLX-tagged tests in train/mtp_mlx_test.go lock the contract:
    • TestMTPValidationLossPrefersEvalLossOutputMLX
    • TestValidationLossFallsBackToLossOutputWithoutMTPMLX

Affected users

Anyone running -mode arch with mtp.n > 1 and reading the val number during or after training. The training was always correct (the optimizer minimizes loss, the MTP-weighted training objective). Only the val display was reading the wrong tensor. Users running -mode eval were already getting the correct number via BuildEvalIRProgramFromConfig, which emits horizon-0 only.

Install

brew tap mrothroc/mixlab https://github.com/mrothroc/mixlab
brew upgrade mixlab

Or build from source:

CGO_ENABLED=1 go build -tags mlx -o mixlab ./cmd/mixlab/
Assets 2
Loading