Transformers - Strong New Nets

KataGo is continuing to improve at https://katagotraining.org/ and if you'd like to donate your spare GPU cycles and support it, it could use your help there!

New Transformer Models

Included in this release for download are three very strong transformer models of various sizes that newly work with this 1.17.0 release:

• b10c384h6nbttflrs.bin.gz - smallest, stronger per visit than the strongest b18 main run models but usually as fast or faster than them
• b10c512h8nbt3tflrs-fson-silu-rsnh.bin.gz - medium, stronger per visit than the strongest b28 main run models but usually as fast or faster than them
• b11c768h12nbt3tflrs-fson-silu.bin.gz - largest transformer so far, stronger than the b40 zhizhi main run models but usually as fast or faster than them

(see here for pytorch .ckpt versions of these transformers)

See https://katagotraining.org/networks/ for the main run models so far which still all work too. The main run as of the time of this release has not switched to transformers since this is the first version supporting them, but will switch soon once enough contributors upgrade to v1.17.0. See also https://katagotraining.org/extra_networks/ for other unusual models that can be used with KataGo.

Getting Started / Choosing a Backend

If you're a new user, this section has tips for getting started and basic usage! For choosing a backend:

• If you have an NVIDIA GPU, use CUDA+CUDNN or TensorRT. You'll also have to install CUDA and one of CUDNN or TensorRT from nvidia depending on your choice.

  • TensorRT - Usually best performance but longer startup times.
    • For transformers, i.e. the best models, prefer CUDA 13 + TensorRT 10.16, older versions may be outperformed by CUDA+CUDNN.
  • CUDA+CUDNN - Faster startup times than TensorRT, still pretty good performance, occasionally actually fastest.
    • For transformers, i.e. the best models, usually slightly slower than TRT 10.16 (but not always), and faster than older TRT versions.
    • Use CUDNN >= 9.8.0 if at all possible. The CUDNN 8.9.7 builds will be a LOT slower when running transformer models.
  • (CUDA 13, CUDNN 9.24, TensorRT 10.16 builds are new in this release. TensorRT < 10 support has been dropped)

• If you have a non-NVIDIA GPU, or have an NVIDIA GPU but really don't want to install CUDNN/TensorRT, OpenCL works across GPUs but will be much slower. (Future releases will likely also add faster and better support for other GPUs, stay tuned)

• For MacOS, use the Metal backend, which you can generally get by installing KataGo from homebrew, which usually updates not too long after KataGo's own release.

• If you need a pure-CPU version of KataGo, use Eigen AVX2, or if your CPU is also ancient and Eigen AVX2 doesn't work, then try Eigen.

Other notes

• +bs50 - these are just for fun, and don't support distributed training but DO support board sizes up to 50x50. They may also be slightly slower and will use much more memory, even when only playing on 19x19, so use them only when you really want to try large boards.

• Linux executables were compiled on a 22.04 Ubuntu machine using AppImage. You will still need to install e.g. correct versions of Cuda/TensorRT or have drivers for OpenCL, etc. on your own. Compiling from source is also not so hard on Linux, see the "TLDR" instructions for Linux here.

Major Changes this Release

Transformer support! (yay)

This release adds support for transformer neural nets in the C++ engine on all major backends (thanks @ChinChangYang for Metal/CoreML, and thanks @zsqdx for performance optimizations and @hzyhhzy for various training improvements, along with many others). Transformer models are generally much stronger given the same compute cost, and the main training run at katagotraining.org will soon be switching to a strong transformer and will therefore require v1.17.x to contribute to.

GPU Backend Fixes and Performance Improvements:

• Added CUDA config option cudaUse1x1Matmul controlling whether 1x1 convolutions run as a cuBLAS matrix multiply, which is slightly faster in FP16. New default uses it when FP16 is enabled, which may result in a few percent speedup on some machines.
• Although OpenCL is still not as optimized as other backends, made a variety of robustness improvements, and fixed a bug that prevented the use of tensor cores for 1x1 convs, speeding up models noticeably on some GPUs.
• Fixed a major memory leak in Metal backend that could eventually exhaust memory in long-running processes.
• Greatly reduced the memory usage of the on-device CoreML model conversion and ANE (peak memory during load roughly halved, and steady-state reduced also by a large factor).
• Various improvements to error handling, checks, and robustness in the Metal/CoreML backend.

(Metal/CoreML fixes by thanks to @ChinChangYang)

Other Feature additions/changes for KataGo users:

• GTP extensions kata-raw-nn and kata-raw-human-nn now accept an optional color argument to specify the player to move, and fixed the parsing of the optional optimism argument of kata-raw-nn. See GTP_Extensions.md.
• Analysis engine allowMoves can now specify one entry for each player separately, instead of only a single entry total.
• Increased the max komi that users can specify from 150 to 400.
• KataGo now logs the number of parameters of the neural net being loaded.
• Dropped support for TensorRT versions < 10.

Bugfixes for KataGo users:

• Fixed bug where the experimental eval cache (added v1.16.4, off by default) could reuse cached search results across changed search parameters in GTP and the analysis engine, causing major errors as very different winrates/scores could get blended together.
• Fixed New Zealand rules to use a default komi of 7 instead of 7.5, matching the actual New Zealand rules.
• Fixed OpenCL tuner bug where performing the initial tuning with FP16 disabled would incorrectly cache the GPU as not supporting FP16, silently degrading performance of later FP16 runs.

C++ selfplay/training-data changes and other internal changes:

• TensorRT backend now works by emitting ONNX in memory and parsing it with TensorRT's ONNX parser. Set trtDisableOnnx=true to fall back to the previous hand-constructed network graph (convnets only).
• For TensorRT, added a trtDumpDebugPlanToDir config option that dumps the emitted ONNX, the serialized engine, and per-layer engine info for debugging precision or fusion issues.
• CUDA backend now requires cudnn-frontend library, for fused attention kernels for transformers.
• Transformer nets are exported using a new model version 17.
• Added experimental option to reanalyze positions during selfplay data generation (useReanalyze, reanalyzeProp, and related options). After a game finishes, a random subset of the cheap-search positions are re-searched fully and recorded as training data, favoring positions where the cheap search result was surprising.
• Added experimental option useSearchValueSurprise to compute value surprise for data weighting/selection from the search result of the turn itself relative to the raw net evaluation, rather than from the realized game outcome, so that data weighting does not condition on the game's actual result.
• Added experimental options to support future migration work for pass-alive handling in trained nets and self-play to fix a pathological corner case caused by KataGo's rules, including new support for trained models to flag themselves as supporting it.
• Selfplay now logs some additional statistics about game generation.
• Adjusted handling of illegal moves when loading starting positions and sgfs, adding a new option to tolerate them.
• Fixed a minor concurrent access bug in training data writing.

Python scripts feature additions/changes:

• Major training performance optimizations, as much as 1.7x faster on some hardware setups on some model architectures (probably a bit less for convnets). Most optimizations are on by default and individually controllable via KATAGO_* environment variables.
• Added experimental support for the Aurora optimizer, a variant of Muon.
• Added -lr-schedule argument to train.py to specify an explicit piecewise-constant learning rate scale schedule from the command line.
• Major cleanup of shuffle.py arguments and defaults. NOTE: deliberate compatibility break - some formerly-required arguments now have sensible defaults, -keep-target-rows is now required (and accepts 'all'), and the deprecated -add-to-window-size was removed in favor of -add-to-data-rows.
• shuffle.py has a -num-waves option for very large whole-dataset shuffles that still maintains perfect permutation, a -dry-run-print-resource-cost mode to estimate disk/memory costs without running, and better logging.
• Added new script refresh_batchnorm_stats.py that exactly recomputes batchnorm running statistics of a checkpoint (for both the plain and SWA model) over a sample of data, improving models for checkpoints whose batchnorm stats are stale (particularly those at higher LR or large batch size).
• Improved training data loading: added a -data-prefetch-depth option, reduced CPU and memory overhead, improved failure handling.
• benchmark_fresh_model.py improvements, including a mode that benchmarks the full training loop step rather than only the model forward/backward.
• Model export now computes a bound on the possible attention logit magnitudes of transformer nets and refuses to export nets whose logits could grow large enough to compromise FP16 attention masking in inference backends (overridable), and train.py now has an optional penalty to discourage large attention logits during training.
• Cleaned up model configs and added configs for some larger transformer models.
• Fixed optimizer name not being set in the training state when loading a checkpoint without optimizer state (thanks @loker404!).
• Fixed a barrier issue in distributed training that could cause validation to time out.
• Fixed shuffle.py logged data row ranges to account for -add-to-data-rows.
• Minor fixes and updates to various other scripts.

Build/CI fixes:

• Fixed Windows CI to work with newer Visual Studio versions rather than requiring VS 2022 (thanks @ChinChangYang).
• Fixed stale build cache issues in GitHub Actions builds.