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?
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?