pyro-ppl/pyro

[FR] Support Automatic Mixed Precision training

Open

#3,316 opened on Jan 31, 2024

View on GitHub
 (7 comments) (0 reactions) (0 assignees)Python (8,211 stars) (981 forks)batch import
enhancementhelp wanted

Description

Issue Description

Better support for mixed precision training would be extremely helpful, at least for SVI. I can manually cast data into float16 or bfloat16 but I am unable to leverage PyTorch's automatic mixed precision training. This is because it requires the use of the GradScaler class during the optimization loop to properly scale gradients in a mixed-precision-aware manner. See the documentation for more info: https://pytorch.org/docs/stable/amp.html

It would be nice to have support for using this class within pyro optimizers to allow for amp support.

Contributor guide

Tech stack
pythonpytorch
Domain
machine learningperformance
Issue type
feature
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-2 days
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
active
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
clear
Prerequisites
PyTorchPyro SVIAMP basics
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
40
Research direction
Examine Pyro's optimizer classes in pyro/optim/ to understand gradient application. Review issue comments for proposed approaches. Implement GradScaler integration into the optimizer's step method, following PyTorch's AMP guidelines. Consider a flag to enable AMP in SVI.