google-deepmind / google-deepmind/open_spiel
[Feature] Support bfloat16 mixed-precision in AlphaZero_Torch
- Dominant language
- C++
- Stars
- 5.5k
- Forks
- 1.2k
- Avg merge
- 2d 8h
- Merged PRs (30d)
- 4
Description
Description:
Motivation
Currently, the AlphaZero PyTorch implementation (open_spiel/algorithms/alpha_zero_torch) uses float32 exclusively. For board games with large observation spaces (e.g., Go is 19×19×17=6137, Chess is 8×8×119=7616 floats), this creates significant memory overhead and leaves potential Tensor Core performance gains on the table during both training and inference.
Implementing mixed-precision (bfloat16) support would effectively halve the GPU memory footprint for activations and yield a 30-50% training speedup, enabling larger batch sizes.
Current Implementation Limitations
Observation tensors are instantiated and copied directly from std::vector to torch::Tensor using from_blob().clone() with no precision scaling.
There is currently no Automatic Mixed Precision (AMP) support in the learning loop or inference methods.
Proposed Strategy
I would like to open a PR to implement bfloat16 mixed-precision support with the following approach:
Configuration: Add a precision field (e.g., "float32", "bfloat16") to ModelConfig in open_spiel/algorithms/alpha_zero_torch/model.h.
Autocast Integration: In ModelImpl::losses, wrap the forward pass with torch::autocast when half-precision is configured.
Tensor Casting: In vpnet.cc (Inference and Learn methods), cast the observation and target tensors to kBFloat16 immediately after .to(device).
Numerical Stability: Maintain float32 for the value head output calculating MSE loss, restricting bfloat16 to the activations and policy logits.
Loss Scaling: Integrate WeightDecay masking within the AMP loss scaler.
Having implemented similar float16/bfloat16 hardware-level support in other open-source ML frameworks, I am happy to build this out and run the necessary performance benchmarks.
Would the team be open to a PR implementing this? Are there any specific edge cases in the current vpnet architecture I should be mindful of before starting?
Contributor guide
Research direction
Start with open_spiel/algorithms/alpha_zero_torch/model.h, vpnet.cc, and ModelImpl::losses to map the existing configuration, inference, and learning paths. Done means bfloat16 can be selected for activations in both training and inference while the value-head MSE remains float32, with the proposed behavior validated through performance benchmarks.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100