google-deepmind / google-deepmind/open_spiel

[Feature] Support bfloat16 mixed-precision in AlphaZero_Torch

Open
#1,548 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.