mirror of
https://github.com/facebookresearch/fairseq.git
synced 2024-09-11 17:25:31 +03:00
optimize mixed precision (#1248)
Summary: # Before submitting - [ ] Was this discussed/approved via a Github issue? (no need for typos, doc improvements) - [x] Did you read the [contributor guideline](https://github.com/pytorch/fairseq/blob/master/CONTRIBUTING.md)? - [x] Did you make sure to update the docs? - [x] Did you write any new necessary tests? ## What does this PR do? Implements the multiply_factor optimization used in memory efficient fp16 training to mixed precision training. The methods multiply_grads and clip_grad_norm do not touch each gradient, but rather a "multiply factor" that is then factored in when unscaling gradients. ## PR review Anyone in the community is free to review the PR once the tests have passed. If we didn't discuss your PR in Github issues there's a high chance it will not be merged. ## Did you have fun? Make sure you had fun coding � Pull Request resolved: https://github.com/fairinternal/fairseq-py/pull/1248 Reviewed By: myleott Differential Revision: D23201396 Pulled By: andersonic fbshipit-source-id: 6c6f64542893e0ecac72e132464bb334dcb9874d
This commit is contained in:
parent
77983ee1a5
commit
68c87f0abf
@ -75,12 +75,8 @@ class _FP16OptimizerMixin(object):
|
||||
loss.backward()
|
||||
self._needs_sync = True
|
||||
|
||||
def _sync_fp16_grads_to_fp32(self, multiply_grads=1.):
|
||||
def _sync_fp16_grads_to_fp32(self):
|
||||
if self._needs_sync:
|
||||
if self.scaler is not None:
|
||||
# correct for dynamic loss scaler
|
||||
multiply_grads /= self.scaler.loss_scale
|
||||
|
||||
# copy FP16 grads to FP32
|
||||
if self.has_flat_params:
|
||||
offset = 0
|
||||
@ -91,20 +87,18 @@ class _FP16OptimizerMixin(object):
|
||||
numel = grad_data.numel()
|
||||
self.fp32_params.grad.data[offset:offset+numel].copy_(grad_data.view(-1))
|
||||
offset += numel
|
||||
self.fp32_params.grad.data.mul_(multiply_grads)
|
||||
else:
|
||||
for p, p32 in zip(self.fp16_params, self.fp32_params):
|
||||
if not p.requires_grad:
|
||||
continue
|
||||
if p.grad is not None:
|
||||
p32.grad.data.copy_(p.grad.data)
|
||||
p32.grad.data.mul_(multiply_grads)
|
||||
else:
|
||||
p32.grad = torch.zeros_like(p.data, dtype=torch.float)
|
||||
|
||||
self._needs_sync = False
|
||||
|
||||
def _sync_fp32_grads_to_fp16(self):
|
||||
def _sync_fp32_params_to_fp16(self):
|
||||
# copy FP32 params back into FP16 model
|
||||
if self.has_flat_params:
|
||||
offset = 0
|
||||
@ -120,36 +114,47 @@ class _FP16OptimizerMixin(object):
|
||||
continue
|
||||
p.data.copy_(p32.data)
|
||||
|
||||
def _unscale_grads(self):
|
||||
self._sync_fp16_grads_to_fp32()
|
||||
if self._multiply_factor != 1.:
|
||||
self.fp32_optimizer.multiply_grads(self._multiply_factor)
|
||||
self._multiply_factor = 1.
|
||||
|
||||
def multiply_grads(self, c):
|
||||
"""Multiplies grads by a constant ``c``."""
|
||||
if self._needs_sync:
|
||||
self._sync_fp16_grads_to_fp32(c)
|
||||
elif self.has_flat_params:
|
||||
self.fp32_params.grad.data.mul_(c)
|
||||
else:
|
||||
for p32 in self.fp32_params:
|
||||
p32.grad.data.mul_(c)
|
||||
self._multiply_factor *= c
|
||||
|
||||
def clip_grad_norm(self, max_norm, aggregate_norm_fn=None):
|
||||
"""Clips gradient norm and updates dynamic loss scaler."""
|
||||
self._sync_fp16_grads_to_fp32()
|
||||
grad_norm = utils.clip_grad_norm_(self.fp32_params, max_norm, aggregate_norm_fn)
|
||||
|
||||
# detect overflow and adjust loss scale
|
||||
grad_norm = self._multiply_factor * self.fp32_optimizer.clip_grad_norm(0, aggregate_norm_fn)
|
||||
|
||||
if self.scaler is not None:
|
||||
if grad_norm > max_norm > 0.0:
|
||||
self._multiply_factor *= max_norm / grad_norm
|
||||
|
||||
self.scaler.check_overflow(grad_norm)
|
||||
else:
|
||||
clip_coef = (max_norm / (grad_norm + 1e-6)).clamp_(max=1)
|
||||
self._multiply_factor *= clip_coef
|
||||
|
||||
return grad_norm
|
||||
|
||||
def step(self, closure=None):
|
||||
"""Performs a single optimization step."""
|
||||
self._sync_fp16_grads_to_fp32()
|
||||
self.fp32_optimizer.step(closure)
|
||||
|
||||
if self.supports_step_with_scale:
|
||||
self.fp32_optimizer.step(closure, scale=(1. / self._multiply_factor))
|
||||
else:
|
||||
self._unscale_grads()
|
||||
self.fp32_optimizer.step(closure)
|
||||
|
||||
if self.scaler is not None:
|
||||
self.scaler.update()
|
||||
|
||||
self._sync_fp32_grads_to_fp16()
|
||||
self._sync_fp32_params_to_fp16()
|
||||
|
||||
def zero_grad(self):
|
||||
"""Clears the gradients of all optimized parameters."""
|
||||
@ -162,6 +167,9 @@ class _FP16OptimizerMixin(object):
|
||||
p32.grad.zero_()
|
||||
self._needs_sync = False
|
||||
|
||||
if self.scaler is not None:
|
||||
self._multiply_factor = 1. / float(self.scaler.loss_scale)
|
||||
|
||||
|
||||
class FP16Optimizer(_FP16OptimizerMixin, optim.FairseqOptimizer):
|
||||
"""
|
||||
|
77
tests/test_fp16_optimizer.py
Normal file
77
tests/test_fp16_optimizer.py
Normal file
@ -0,0 +1,77 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from fairseq.optim.fp16_optimizer import FP16Optimizer, MemoryEfficientFP16Optimizer
|
||||
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), 'test requires a GPU')
|
||||
class TestGradientScaling(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.x = torch.tensor([2.0]).cuda().half()
|
||||
weight = 3.0
|
||||
bias = 5.0
|
||||
self.error = 1.0
|
||||
self.target = torch.tensor([self.x * weight + bias + self.error]).cuda().half()
|
||||
self.loss_fn = torch.nn.L1Loss()
|
||||
|
||||
self.model = torch.nn.Linear(1, 1)
|
||||
self.model.weight.data = torch.tensor([[weight]])
|
||||
self.model.bias.data = torch.tensor([bias])
|
||||
self.model.cuda().half()
|
||||
self.params = list(self.model.parameters())
|
||||
|
||||
self.namespace_dls = argparse.Namespace(
|
||||
optimizer='adam',
|
||||
lr=[0.1],
|
||||
adam_betas='(0.9, 0.999)',
|
||||
adam_eps=1e-8,
|
||||
weight_decay=0.0,
|
||||
fp16_init_scale=1,
|
||||
fp16_scale_window=1,
|
||||
fp16_scale_tolerance=1,
|
||||
threshold_loss_scale=1,
|
||||
min_loss_scale=1e-4
|
||||
)
|
||||
|
||||
def run_iter(self, model, params, optimizer):
|
||||
optimizer.zero_grad()
|
||||
y = model(self.x)
|
||||
loss = self.loss_fn(y, self.target)
|
||||
optimizer.backward(loss)
|
||||
self.assertEqual(loss, torch.tensor(1., device='cuda:0', dtype=torch.float16))
|
||||
|
||||
grad_norm = optimizer.clip_grad_norm(0)
|
||||
self.assertAlmostEqual(grad_norm.item(), 2.2361, 4)
|
||||
|
||||
optimizer.step()
|
||||
self.assertEqual(model.weight, torch.tensor([[3.0996]], device='cuda:0', dtype=torch.float16, requires_grad=True))
|
||||
self.assertEqual(model.bias, torch.tensor([5.1016], device='cuda:0', dtype=torch.float16, requires_grad=True))
|
||||
self.assertEqual(optimizer.scaler.loss_scale, 2.)
|
||||
|
||||
def test_mixed_precision(self):
|
||||
model = copy.deepcopy(self.model)
|
||||
params = list(model.parameters())
|
||||
optimizer = FP16Optimizer.build_optimizer(self.namespace_dls, params)
|
||||
|
||||
self.run_iter(model, params, optimizer)
|
||||
self.assertTrue(torch.all(optimizer.fp32_params.eq(torch.tensor([3.1000, 5.1000], device='cuda:0', requires_grad=True))))
|
||||
|
||||
def test_memory_efficient(self):
|
||||
model = copy.deepcopy(self.model)
|
||||
params = list(model.parameters())
|
||||
optimizer = MemoryEfficientFP16Optimizer.build_optimizer(self.namespace_dls, params)
|
||||
|
||||
self.run_iter(model, params, optimizer)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
Loading…
Reference in New Issue
Block a user