Skip to content

GitLab

Unverified Commit 13445c55 authored by Benjamin Lefaudeux's avatar Benjamin Lefaudeux Committed by GitHub
Browse files

[feature-fix-refactor][ShardedDDP] Make it possible to change trainability graph on the fly (#369) 

* Better unit testing
* Make it possible to refresh the DDP assumptions when the model has changed. Make it optional so that you save some time
* Enabling accumulation tests
parent 1a636557
Hide whitespace changes
Inline Side-by-side
Showing with 409 additions and 259 deletions
CHANGELOG.md
View file @ 13445c55
......@@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [next rel] - TBD
### Fixed
- ShardedDDP and OSS handle model trainability changes during training ([#369](https://github.com/facebookresearch/fairscale/issues/369))
## [0.1.6] - 2021-02-10
### Added
......
fairscale/nn/data_parallel/sharded_ddp.py
View file @ 13445c55
......@@ -8,10 +8,12 @@ A nn.Module wrapper to go with a Sharded Optimizer in order to handle targeted g
reduction automatically.
"""
from collections import deque
import contextlib
import functools
from itertools import chain
import logging
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
from typing import Any, Callable, Deque, Dict, Generator, List, Optional, Tuple, Union
import torch
from torch import nn
......@@ -22,6 +24,10 @@ from fairscale.optim import OSS
from fairscale.optim.utils import Bucket, Workhandle
def _trainable(param: torch.Tensor) -> bool:
return param.requires_grad
class ShardedDataParallel(nn.Module):
""" Wrap the model, and reduce the gradients to the right rank during the backward pass.
......@@ -45,9 +51,13 @@ class ShardedDataParallel(nn.Module):
Synchronize the models in between the ranks when starting up. Not needed if each rank has the same seed,
or the training restarts from a saved state
reduce_buffer_size (int):
the max size of the buffer used to batch the small parameter tensors, in number of elements (default 8M).
The max size of the buffer used to batch the small parameter tensors, in number of elements (default 8M).
this will impact the long term memory consumption, because these buckets correspond to parameters which will not be sharded.
Set to 0 to remove all bucketing.
auto_refresh_trainable (bool):
(default: True) Check whether the parameters trainability (`requires_grad`) has changed and update both ShardedDDP
and OSS automatically if this is the case. If set to False, `refresh_trainable()` needs to be called anytime
a parameter is frozen or unfrozen.
.. warning:
......@@ -71,6 +81,14 @@ class ShardedDataParallel(nn.Module):
handled. In that case ShardedDDP will raise an exception and suggest to either remove the unused parameters from your model
(https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html?highlight=unused_parameters is helpful)
or set `reduce_buffer_size` to 0
.. warning:
If `auto_refresh_trainable` is set to `True` (this is the default) then any trainability change in the model graph will be handled
automatically.
If `auto_refresh_trainable` is set to `False`, ShardedDDP will not refresh its assumptions with respect to trainable parameters
for every forward pass, in the hope of saving some time. If some parameters are frozen or unfrozen over time, please refresh
ShardedDDP assumptions by calling `refresh_trainable()` just after said change (before the next forward pass).
"""
def __init__(
......@@ -81,12 +99,14 @@ class ShardedDataParallel(nn.Module):
broadcast_buffers: bool = True,
sync_models_at_startup: bool = True,
reduce_buffer_size: int = 2 ** 23,
auto_refresh_trainable: bool = True,
):
super().__init__()
self.module = module
self.sharded_optimizers = [sharded_optimizer] if isinstance(sharded_optimizer, OSS) else sharded_optimizer
self.enable_broadcast_buffers = broadcast_buffers
self.auto_refresh_trainable = auto_refresh_trainable
# Handle a no_sync() context which prevents the gradient synchronization,
# accumulate in place
......@@ -117,14 +137,19 @@ class ShardedDataParallel(nn.Module):
# several optimizers can be present each working on seperate parameter set which is spread across multiple ranks
# - we build an iterator which goes through all the parameters involved globally
all_param_iterator = chain(
*[sum([sum(p, []) for p in optim.per_device_params.values()], []) for optim in self.sharded_optimizers]
self._all_params = list(
chain(
*[sum([sum(p, []) for p in optim.per_device_params.values()], []) for optim in self.sharded_optimizers]
)
)
self._grad_to_be_reduced = [True for _ in filter(lambda x: x.requires_grad, all_param_iterator)]
self._trainable_params: List[torch.Tensor] = []
self._grad_to_be_reduced: List[bool] = []
self._trainable_param_to_rank: Dict[torch.Tensor, int] = {}
self._reference_trainable_mask = list(map(_trainable, self._all_params))
# - keep track of the grads which have already been reduced
self._reduced_grads: Dict[OSS, int] = {}
self._reduced_grads_max = {o: len(o.param_to_rank.values()) for o in self.sharded_optimizers}
self._reduced_grads = 0
self._reduced_grads_max = 0
# - setup buckets and tensor views
model_size = sum([p.numel() for p in self.module.parameters()])
......@@ -140,14 +165,12 @@ class ShardedDataParallel(nn.Module):
)
self.use_buckets = self.buffer_max_size > 0
self.buckets: Dict[OSS, Dict[torch.device, List[Bucket]]] = {o: {} for o in self.sharded_optimizers}
self.buckets: Dict[torch.device, List[Bucket]] = {}
self._should_bucket_grad: List[bool] = []
self._bucket_list: Optional[List[Bucket]] = None
self._setup_bucket_strategy()
# - setup backward hooks which will be called by Torch's autograd in due time
self._grad_accs: List[Callable] = []
self._setup_backward_hooks()
# passing a handle to torch.nn.SyncBatchNorm layer
self._passing_sync_batchnorm_handle(self.module)
......@@ -156,13 +179,25 @@ class ShardedDataParallel(nn.Module):
if sync_models_at_startup:
self._sync_params_and_buffers()
self._clear_counters()
self._work_handles: Deque[Workhandle] = deque()
self.refresh_trainable()
def forward(self, *inputs: Any, **kwargs: Any) -> Any:
"""
Module forward pass, handles any DDP-specific work in the background. Primes the
backward pass for gradient reduction to the proper ranks.
"""
# Optionally check whether the trainable parameters have changed
if self.auto_refresh_trainable:
trainable_mask = list(map(_trainable, self._all_params))
if trainable_mask != self._reference_trainable_mask:
logging.warning("ShardedDDP detected that the trainable params changed, updating the partitioning")
self.refresh_trainable()
self._reference_trainable_mask = trainable_mask
if self.enable_broadcast_buffers:
# NCCL communications are on a different stream, needs to be blocking
# for the subsequent FW to be correct
......@@ -205,13 +240,38 @@ class ShardedDataParallel(nn.Module):
"""
for optimizer in self.buckets.keys():
for device in self.buckets[optimizer].keys():
for bucket in self.buckets[optimizer][device]:
bucket.buffer.to(device=device, dtype=dtype, non_blocking=non_blocking)
for device in self.buckets.keys():
for bucket in self.buckets[device]:
bucket.buffer.to(device=device, dtype=dtype, non_blocking=non_blocking)
self.module.to(device)
def refresh_trainable(self) -> None:
""" If the module trainability has changed, update all the assumptions """
# Make sure that this is not done while gradients are waiting to be reduced (if no_sync context for instance)
assert not functools.reduce(lambda x, y: x or y, self._grad_to_be_reduced, False), "Grads waiting to be reduced"
self._trainable_params = list(filter(lambda x: x.requires_grad, self._all_params))
self._trainable_params.sort(key=lambda x: x.numel())
self._grad_to_be_reduced = [True for _ in self._trainable_params]
self._trainable_param_to_rank = {}
for optim in self.sharded_optimizers:
# OSS may need to change the communication pattern
optim.refresh_trainable()
# Update ShardedDDP given the new partitions
for (
device_per_rank_params
) in optim.per_device_params.values(): # all the params on this device (inc all ranks)
for device_params in device_per_rank_params:
for param in filter(lambda x: x.requires_grad, device_params):
self._trainable_param_to_rank[param] = optim.param_to_rank[param]
self._setup_bucket_strategy()
self._setup_backward_hooks()
def reduce(self) -> None:
""".. deprecated:: 0.0.4
......@@ -223,6 +283,9 @@ class ShardedDataParallel(nn.Module):
def sync_buffers(self, blocking: bool = False) -> None:
"""
Sync all the param buffers in between ranks (including for instance batch norm statistics).
Arguments:
blocking (bool): wait for the operation to conclude.
"""
last_work_handle = None
......@@ -236,6 +299,21 @@ class ShardedDataParallel(nn.Module):
# Only wait for the last coms, they're inlined on the same CUDA stream
last_work_handle.wait()
def zero_grad(self, set_to_none: bool = False) -> None:
r"""Sets gradients of all model parameters to zero. See similar function
under :class:`torch.optim.Optimizer` for more context.
Arguments:
set_to_none (bool): instead of setting to zero, set the grads to None.
See :meth:`torch.optim.Optimizer.zero_grad` for details.
"""
for index, trainable_param in enumerate(self._trainable_params):
if set_to_none and not self._should_bucket_grad[index]:
trainable_param.grad = None
elif trainable_param.grad is not None:
trainable_param.grad.zero_()
def __getattr__(self, name: str) -> Any:
"""Forward missing attributes to wrapped module."""
try:
......@@ -254,21 +332,20 @@ class ShardedDataParallel(nn.Module):
@torch.no_grad()
def _clear_counters(self) -> None:
"""Reset all the grad reduce and call counters"""
if not self.should_accumulate_grads:
self._grad_to_be_reduced = [True for _ in self._grad_to_be_reduced]
self._reduced_grads = 0
self._grad_to_be_reduced = [True for _ in self._grad_to_be_reduced]
self._reduced_grads = {o: 0 for o in self.sharded_optimizers}
if self.use_buckets:
assert self._bucket_list is not None
# Do not reset the buckets
if self.use_buckets:
assert self._bucket_list is not None
for bucket in self._bucket_list:
assert bucket.sent, (
"A bucket failed to be sent, probably unused parameters."
+ "Either remove the unused parameter or de-activate ShardedDDP buckets -set reduce_buffer_size to 0-"
)
for bucket in self._bucket_list:
assert self.should_accumulate_grads or bucket.sent, (
"A bucket failed to be sent, probably unused parameters."
+ "Either remove the unused parameter or de-activate ShardedDDP buckets -set reduce_buffer_size to 0-"
)
bucket.reset()
bucket.reset()
def _find_rank(self, param: Parameter) -> Tuple[OSS, int]:
""" Look up where this parameter belongs to """
......@@ -279,7 +356,7 @@ class ShardedDataParallel(nn.Module):
assert False, "This parameter is not present in an optimizer, this should not happen"
return (None, -1)
def _get_reduce_fn(self, index: int, param: torch.Tensor, dst_rank: int, optimizer: OSS) -> Callable:
def _get_reduce_fn(self, index: int, param: torch.Tensor, dst_rank: int) -> Callable:
"""
Two possible backward hooks for a given parameter: either directly reduce to the appropriate rank,
or contribute to a bucket and reduce when the bucket is full.
......@@ -287,16 +364,16 @@ class ShardedDataParallel(nn.Module):
Either way a delayed action is necessary and is passed as a callback.
"""
@torch.no_grad()
def reduce(*_: Any) -> None:
# Skip gradient reduction, do not alter status flags
if not self.should_accumulate_grads and self._grad_to_be_reduced[index]:
assert param.grad is not None, "Reducing gradients during backward pass, cannot be None"
if not self.use_buckets or not self._should_bucket_grad[index]:
# Direct reduction
@torch.no_grad()
def reduce(*_: Any) -> None:
# Skip gradient reduction, do not alter status flags
if not self.should_accumulate_grads and self._grad_to_be_reduced[index]:
assert param.grad is not None, "Reducing gradients during backward pass, cannot be None"
# Make sure that this is not fired twice
self._grad_to_be_reduced[index] = False
if not self.use_buckets or not self._should_bucket_grad[index]:
# Make sure that this is not fired twice
self._grad_to_be_reduced[index] = False
param.grad.mul_(self.world_size_scaling)
# Future work includes clearing up the buffer if possible
......@@ -305,7 +382,7 @@ class ShardedDataParallel(nn.Module):
param.grad = None
# Async reduce for this buffer, log the future
optimizer.work_handles.append(
self._work_handles.append(
Workhandle(
handle=dist.reduce(
tensor=param.grad.data, dst=dst_rank, group=self.process_group, async_op=True
......@@ -313,9 +390,28 @@ class ShardedDataParallel(nn.Module):
callback=cleanup,
)
)
self._reduced_grads[optimizer] += 1
else:
bucket = self.buckets[optimizer][param.device][dst_rank]
self._reduced_grads += 1
# Opportunistically try to empty the queue
self._try_consume_work_handle()
# If all the reduce operations have been called,
# make sure that all the asynchronous calls have concluded before moving on
# and execute the delayed actions (release gradients, unroll the buckets)
if self._reduced_grads == self._reduced_grads_max:
self._consume_work_handles()
else:
@torch.no_grad()
def reduce(*_: Any) -> None:
# Skip gradient reduction, do not alter status flags
if not self.should_accumulate_grads and self._grad_to_be_reduced[index]:
assert param.grad is not None, "Reducing gradients during backward pass, cannot be None"
# Make sure that this is not fired twice
self._grad_to_be_reduced[index] = False
bucket = self.buckets[param.device][dst_rank]
bucket.params_checked_in += 1
if bucket.full():
......@@ -324,7 +420,7 @@ class ShardedDataParallel(nn.Module):
# Reduce the bucket
bucket.sent = True
optimizer.work_handles.append(
self._work_handles.append(
Workhandle(
handle=dist.reduce(
tensor=bucket.buffer, dst=dst_rank, group=self.process_group, async_op=True,
......@@ -332,16 +428,16 @@ class ShardedDataParallel(nn.Module):
callback=None,
)
)
self._reduced_grads[optimizer] += 1
self._reduced_grads += 1
# Opportunistically try to empty the queue
optimizer._try_consume_work_handle()
# Opportunistically try to empty the queue
self._try_consume_work_handle()
# If all the reduce operations have been called,
# make sure that all the asynchronous calls have concluded before moving on
# and execute the delayed actions (release gradients, unroll the buckets)
if self._reduced_grads[optimizer] == self._reduced_grads_max[optimizer]:
optimizer._consume_work_handles()
# If all the reduce operations have been called,
# make sure that all the asynchronous calls have concluded before moving on
# and execute the delayed actions (release gradients, unroll the buckets)
if self._reduced_grads == self._reduced_grads_max:
self._consume_work_handles()
return reduce
......@@ -352,33 +448,27 @@ class ShardedDataParallel(nn.Module):
"""
# Go through the parameters, attach the hook
for sharded_optimizer in self.sharded_optimizers:
for (
device_per_rank_params
) in sharded_optimizer.per_device_params.values(): # all the params on this device (inc all ranks)
for device_params in device_per_rank_params:
for param in filter(lambda x: x.requires_grad, device_params):
if param.grad is not None and param.grad.requires_grad:
raise RuntimeError("ShardedDataParallel only works with gradients that don't require grad")
self._grad_accs = []
for index, param in enumerate(self._trainable_params):
if param.grad is not None and param.grad.requires_grad:
raise RuntimeError("ShardedDataParallel only works with gradients that don't require grad")
# Register the hook to the next function in line,
# so that the hook is fired when this grad has properly been computed
p_tmp = param.expand_as(param)
assert p_tmp.grad_fn is not None
grad_acc = p_tmp.grad_fn.next_functions[0][0]
dst_rank = sharded_optimizer.param_to_rank[param]
index = len(self._grad_accs)
# Register the hook to the next function in line,
# so that the hook is fired when this grad has properly been computed
p_tmp = param.expand_as(param)
assert p_tmp.grad_fn is not None
grad_acc = p_tmp.grad_fn.next_functions[0][0]
dst_rank = self._trainable_param_to_rank[param]
grad_acc.register_hook(self._get_reduce_fn(index, param, dst_rank, sharded_optimizer))
self._grad_accs.append(grad_acc) # keep this function in scope
grad_acc.register_hook(self._get_reduce_fn(index, param, dst_rank))
self._grad_accs.append(grad_acc) # keep this function in scope
# Add a hook on the module to flush the buckets, if needed
if self.use_buckets:
def bucket_flush(*unused: Any) -> None:
handle = None
def bucket_flush(*_: Any) -> None:
assert self._bucket_list is not None
handle = None
for bucket in self._bucket_list:
if not bucket.sent:
......@@ -425,62 +515,80 @@ class ShardedDataParallel(nn.Module):
layer._specify_ddp_gpu_num(1) # type: ignore
def _setup_bucket_strategy(self) -> None:
"""Devise a bucketing strategy on a per-rank ownership level. These buckets will not be sharded, since the gradients would be re-allocated during the backward in that case.
"""Devise a bucketing strategy on a per-rank ownership level.
These buckets will not be sharded, since the gradients would be re-allocated during the backward in that case.
This method can be a slow for big models, but it it not typically called often (not for every forward for instance)
"""
# A priori, one reduce call per param
self._reduced_grads_max = len(self._trainable_params)
if not self.use_buckets:
return
# Devise the bucketing strategy
for sharded_optimizer in self.sharded_optimizers:
for device, per_rank_params in sharded_optimizer.per_device_params.items():
self.buckets[sharded_optimizer][device] = []
for dst_rank, params in enumerate(per_rank_params):
offset = 0
self.buckets[sharded_optimizer][device].append(
Bucket(
buffer=torch.zeros(self.buffer_max_size, dtype=per_rank_params[0][0].dtype, device=device)
)
)
bucket = self.buckets[sharded_optimizer][device][dst_rank]
bucket.destination = dst_rank
for param in filter(lambda x: x.requires_grad is True, params):
# Criteria to decide whether this parameter is to be bucketed or not:
# - enough room in the bucket
if (offset + param.numel()) < self.buffer_max_size:
self._should_bucket_grad.append(True)
# This parameter gradients becomes a view of the bucket
offset_next = offset + param.numel()
# Devise the bucketing strategy. Parameters are already sorted, in that:
# - these are only the trainable parameters, so they should produce grads
# - they are sorted by increasing size
self.buckets = {}
if param.grad is None:
# will be overwritten just below, see next line
param.grad = torch.zeros_like(param)
for param in self._trainable_params:
device = param.device
dst_rank = self._trainable_param_to_rank[param]
param.grad.data = bucket.buffer[offset:offset_next].view_as(param.data)
offset = offset_next
# Update the bucket
self._reduced_grads_max[sharded_optimizer] -= 1 # one less reduce call per bucketed grad
self.buckets[sharded_optimizer][device][dst_rank].max_params_checked_in += 1
else:
self._should_bucket_grad.append(False)
# Resize the bucket to remove lost space in the end
bucket.buffer.resize_(offset)
if bucket.max_params_checked_in > 0:
self._reduced_grads_max[sharded_optimizer] += 1 # one reduce call per bucket
self._bucket_list = list(
chain(
*[
self.buckets[sharded_optimizer][device]
for sharded_optimizer in self.sharded_optimizers
for device in sharded_optimizer.per_device_params.keys()
if param.device not in self.buckets.keys():
self.buckets[param.device] = [
Bucket(buffer=torch.zeros(self.buffer_max_size, dtype=param.dtype, device=device))
for _ in range(dist.get_world_size(self.process_group))
]
)
)
bucket = self.buckets[device][dst_rank]
bucket.destination = dst_rank
# Criteria to decide whether this parameter is to be bucketed or not:
# - enough room in the bucket
if (bucket.fill + param.numel()) < self.buffer_max_size:
self._should_bucket_grad.append(True)
# This parameter gradients becomes a view of the bucket
fill_next = bucket.fill + param.numel()
if param.grad is None:
# will be overwritten just below, see next line
param.grad = torch.zeros_like(param)
param.grad.data = bucket.buffer[bucket.fill : fill_next].view_as(param.data)
bucket.fill = fill_next
# Update the bucket
self._reduced_grads_max -= 1 # one less reduce call per bucketed grad
self.buckets[device][dst_rank].max_params_checked_in += 1
else:
self._should_bucket_grad.append(False)
self._bucket_list = list(chain(*[self.buckets[device] for device in self.buckets.keys()]))
# Resize the buckets to remove lost space in the end
for bucket in self._bucket_list:
bucket.buffer.resize_(bucket.fill)
bucket.sent = True
if bucket.max_params_checked_in > 0:
self._reduced_grads_max += 1 # one reduce call per bucket
def _consume_work_handles(self) -> None:
"""Consume all the futures which are tied to this optimizer's buckets.
We start from the first/older ones, since they are the most likely to be ready and non-blocking
"""
while len(self._work_handles) > 0:
work_handle = self._work_handles.popleft()
work_handle.handle.wait()
if work_handle.callback is not None:
work_handle.callback()
def _try_consume_work_handle(self) -> None:
"""Try to consume the oldest future. This is non blocking, if not ready we'll pass"""
while len(self._work_handles) > 0 and self._work_handles[0].handle.is_completed():
work_handle = self._work_handles.popleft()
if work_handle.callback is not None:
work_handle.callback()
fairscale/optim/oss.py
View file @ 13445c55
......@@ -3,19 +3,19 @@
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.
from collections import OrderedDict, deque
from collections import OrderedDict
import copy
from itertools import chain
import logging
from math import inf
from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional, Type, Union
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
import torch
import torch.distributed as dist
from torch.nn import Parameter
from torch.optim import SGD, Optimizer
from .utils import Workhandle, broadcast_object, recursive_copy_to_device
from .utils import broadcast_object, recursive_copy_to_device
__all__ = ["OSS"]
......@@ -52,6 +52,14 @@ class OSS(Optimizer):
torch.distributed group (default: group.WORLD)
broadcast_buffer_size (int):
(deprecated) used to cap the size of the broadcast buffers, not being used anymore.
.. warning: the communication patterns that OSS use depend on the "trainability" graph,
meaning that all the parameters which `require_grad` are handled differently. This is
not reevaluated at every step, please use `refresh_trainable()` if your model changed
(freeze or unfreeze for instance).
If used with :class:<fairscale.nn.ShardedDDP> then an automatic change detection is possible,
via the `auto_refresh_trainable` parameter.
"""
#: The optimizer used for a given shard
......@@ -81,27 +89,22 @@ class OSS(Optimizer):
self._param_to_index: Dict[int, int] = {}
self._local_params: Optional[List[torch.Tensor]] = None
# Build the wrapped optimizer, responsible for a shard of the params
# Default empty values + immutables
self._optim_defaults = default
self._optim_constructor = optim
self.group = group if group is not None else dist.group.WORLD
self.world_size = dist.get_world_size(self.group)
self.rank = dist.get_rank(self.group)
self.global_rank = self.get_global_rank(self.group, self.rank)
self.optim = optim(self.partition_parameters()[self.rank], **default)
# - Sync local and global param_groups keys
for global_group, local_group in zip(self.param_groups, self.optim.param_groups):
for key, value in local_group.items():
if key != "params":
global_group[key] = value
self.buckets: Dict[torch.device, List[torch.Tensor]] = {}
# Optional consolidated optimizer state
self._all_states: List[Dict[str, Any]] = []
self._all_states: List[Dict[str, Any]] = [] # Optional consolidated optimizer state
self._default_device = torch.device("cpu")
# Current default device is set by the parameters allocated to this rank
self._device = list(self.per_device_params.keys())[0]
self.work_handles: Deque[Workhandle] = deque()
self.buckets: Dict[torch.device, List[torch.Tensor]] = {}
self._setup_flat_buffers()
# Setup everything which is related to the parameters to be trained
# (partition and optimizer for the shard)
self.refresh_trainable()
# Partition helpers
def partition_parameters(self) -> List[List[dict]]:
......@@ -277,12 +280,12 @@ class OSS(Optimizer):
# Compute the norm on this grad set,
# then sync all the norms from all ranks
if norm_type == inf:
total_norm = max(p.grad.detach().abs().max().to(self._device) for p in local_params)
total_norm = max(p.grad.detach().abs().max().to(self._default_device) for p in local_params)
# all reduce over data parallel and model parallel workers
dist.all_reduce(total_norm, op=torch.distributed.ReduceOp.MAX, group=dist.group.WORLD)
else:
local_norm = torch.norm(
input=torch.stack([torch.norm(input=p.grad.detach(), p=norm_type, dtype=torch.float32).to(self._device) for p in local_params]), # type: ignore
input=torch.stack([torch.norm(input=p.grad.detach(), p=norm_type, dtype=torch.float32).to(self._default_device) for p in local_params]), # type: ignore
p=norm_type,
)
......@@ -412,16 +415,25 @@ class OSS(Optimizer):
OSS._sync_param_groups(state_dict["param_groups"], self.param_groups)
OSS._sync_param_groups(self.param_groups, self.optim.param_groups)
def refresh_trainable(self) -> None:
""" Updates the partitioning and communication patterns if the trainability (`requires_grad`)
of some parameters changed
"""
# Create the optim which will work on the param shard
if not hasattr(self, "optim"):
self._clear_cache()
self._default_device = list(self.per_device_params.keys())[0]
self.optim = self._optim_constructor(self.partition_parameters()[self.rank], **self._optim_defaults)
OSS._sync_param_groups(self.optim.param_groups, self.param_groups)
self._setup_flat_buffers()
def _broadcast_state_dict(self) -> None:
"""Broadcast this rank's state shard, discard others"""
# Default to CPU space to gain some memory headroom
local_cpu_state = recursive_copy_to_device(
self.optim.state_dict(), non_blocking=True, device=torch.device("cpu")
)
# Tensor cannot be really empty, even if its size is meaningless
dummy_sync_tensor = torch.tensor([1], device=self._device)
dummy_sync_tensor = torch.tensor([1], device=self._default_device)
for rank in range(self.world_size):
if rank == self.rank:
......@@ -431,17 +443,20 @@ class OSS(Optimizer):
)
# legacy compatibility for old torch versions
broadcast_object(
self.local_state_dict(), src_rank=self.global_rank, group=self.group, dist_device=self._device
self.local_state_dict(),
src_rank=self.global_rank,
group=self.group,
dist_device=self._default_device,
)
else:
global_rank = self.get_global_rank(self.group, rank)
# Discard this tensor/rank, broadcast necessary for syncing and because NCCL does not support gather
broadcast_object(
torch.tensor([dummy_sync_tensor], dtype=torch.uint8, device=self._device),
torch.tensor([dummy_sync_tensor], dtype=torch.uint8, device=self._default_device),
src_rank=global_rank,
group=self.group,
dist_device=self._device,
dist_device=self._default_device,
)
def _collect_sharded_states(self) -> List[Dict[str, Any]]:
......@@ -457,19 +472,19 @@ class OSS(Optimizer):
# Sync with other replicas
broadcast_object(
torch.tensor([0], dtype=torch.uint8, device=self._device),
torch.tensor([0], dtype=torch.uint8, device=self._default_device),
src_rank=self.global_rank,
group=self.group,
dist_device=self._device,
dist_device=self._default_device,
)
else:
# Fetch the optim state from the other replicas
global_rank = self.get_global_rank(self.group, rank)
replica_state = broadcast_object(
torch.tensor([0], dtype=torch.uint8, device=self._device),
torch.tensor([0], dtype=torch.uint8, device=self._default_device),
src_rank=global_rank,
group=self.group,
dist_device=self._device,
dist_device=self._default_device,
)
all_states.append(
......@@ -546,23 +561,6 @@ class OSS(Optimizer):
if last_work_handle:
last_work_handle.wait()
def _consume_work_handles(self) -> None:
"""Consume all the futures which are tied to this optimizer's buckets.
We start from the first/older ones, since they are the most likely to be ready and non-blocking
"""
while len(self.work_handles) > 0:
work_handle = self.work_handles.popleft()
work_handle.handle.wait()
if work_handle.callback is not None:
work_handle.callback()
def _try_consume_work_handle(self) -> None:
"""Try to consume the oldest future. This is non blocking, if not ready we'll pass"""
while len(self.work_handles) > 0 and self.work_handles[0].handle.is_completed():
work_handle = self.work_handles.popleft()
if work_handle.callback is not None:
work_handle.callback()
def _setup_flat_buffers(self) -> None:
"""Make all params which are on the same device and tied to the same rank views of a single buffer.
This is used at construction time, and anytime parameter trainability is changed (frozen or unfrozen) and
......@@ -570,19 +568,35 @@ class OSS(Optimizer):
"""
for device, per_rank_params in self.per_device_params.items():
self.buckets[device] = []
# Only wipe the existing buckets if there are none
# (could be that this is called twice, when trainability changes)
if device not in self.buckets.keys():
self.buckets[device] = []
# Make parameters a view of the bucket
for dst_rank, params in enumerate(per_rank_params):
if len(params) > 0:
# Clone the non-trainable params, if in a bucket it will get destroyed
for param in filter(lambda x: not x.requires_grad, params):
param.data = param.data.detach().clone()
# Merge all the trainable params in a single bucket
trainable_params = list(filter(lambda x: x.requires_grad, params))
buffer_size = sum(map(lambda x: x.numel(), trainable_params))
self.buckets[device].append(torch.empty(buffer_size, dtype=params[0].dtype, device=device))
bucket = torch.empty(buffer_size, dtype=params[0].dtype, device=device)
offset = 0
for param in trainable_params:
# This parameter becomes a view of the bucket
offset_next = offset + param.numel()
self.buckets[device][dst_rank][offset:offset_next].copy_(param.data.flatten())
param.data = self.buckets[device][dst_rank][offset:offset_next].view_as(param.data)
bucket[offset:offset_next].copy_(param.data.flatten())
param.data = bucket[offset:offset_next].view_as(param.data)
offset = offset_next
# Either replace the existing bucket, or create it
if len(self.buckets[device]) == dst_rank:
self.buckets[device].append(bucket)
else:
self.buckets[device][dst_rank] = bucket
else:
self.buckets[device].append(torch.zeros(1, device=device))
fairscale/optim/utils.py
View file @ 13445c55
......@@ -3,11 +3,11 @@
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.
import collections
import io
from typing import Any, Callable, Dict, Optional
import torch
from torch._six import container_abcs
import torch.distributed as dist
......@@ -38,7 +38,7 @@ def recursive_copy_to_device(value: Any, non_blocking: bool, device: torch.devic
return values if isinstance(value, list) else tuple(values)
if isinstance(value, container_abcs.Mapping):
if isinstance(value, collections.abc.Mapping):
device_val: Dict[str, Any] = {}
for key, val in value.items():
device_val[key] = recursive_copy_to_device(val, non_blocking=non_blocking, device=device)
......@@ -89,6 +89,7 @@ class Bucket:
self.max_size = buffer.numel()
# Current status for this buffer
self.fill = 0
self.params_checked_in = 0
self.max_params_checked_in = 0 # atttribute present for convenience purposes
self.destination = -1
......
fairscale/utils/testing.py
View file @ 13445c55
......@@ -406,3 +406,11 @@ def objects_are_equal(a: Any, b: Any, raise_exception: bool = False) -> bool:
return False
else:
return a == b
def check_same_model_params(model_a: torch.nn.Module, model_b: torch.nn.Module, message: str = "") -> None:
for p_a, p_b in zip(model_a.parameters(), model_b.parameters()):
assert torch.allclose(p_a, p_b, atol=1e-3), f"Model parameters differ\n{p_a} {p_b}\n" + message
for b_a, b_b in zip(model_a.buffers(), model_b.buffers()):
assert torch.allclose(b_a, b_b), f"Model buffers differ {b_a} - {b_b}\n" + message
stubs/torch/nn/parallel/distributed.pyi
View file @ 13445c55
......@@ -18,12 +18,13 @@ class DistributedDataParallel(Module[T_co]):
check_reduction: bool = ...
broadcast_bucket_size: float = ...
bucket_bytes_cap: float = ...
find_unused_parameters: bool = ...
# TODO type process_group once `distributed` module is stubbed
def __init__(self, module: Module[T_co], device_ids: Optional[_devices_t] = ...,
output_device: Optional[_device_t] = ..., dim: int = ...,
broadcast_buffers: bool = ..., process_group: Optional[Any] = ..., bucket_cap_mb: float = ...,
check_reduction: bool = ...) -> None: ...
check_reduction: bool = ..., find_unused_parameters: bool = ...) -> None: ...
def forward(self, *inputs: Any, **kwargs: Any) -> T_co: ...
......
tests/nn/data_parallel/test_sharded_ddp.py
View file @ 13445c55
......@@ -23,7 +23,7 @@ from torch.nn.parallel import DistributedDataParallel as DDP
from fairscale.nn.data_parallel import ShardedDataParallel
from fairscale.optim import OSS
from fairscale.optim.grad_scaler import ShardedGradScaler
from fairscale.utils.testing import GPT2, skip_if_no_cuda, skip_if_py38, skip_if_single_gpu
from fairscale.utils.testing import GPT2, check_same_model_params, skip_if_no_cuda, skip_if_py38, skip_if_single_gpu
def run_one_step(rank, world_size, backend, device, temp_file_name):
......@@ -133,67 +133,66 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
torch.cuda.set_device(rank)
torch.manual_seed(rank)
np.random.seed(rank)
NUMBER_BATCHS = 5
INPUTS = 2
BATCH_SIZE = 32
def check_parity(amp: bool, accumulate: bool, change_train_graph: bool):
# The API should be the exact same in between the sharded and non-sharded variants, generic closure
def closure(model, scaler, input_tensor, should_accumulate):
accumulate_steps = 3 if should_accumulate else 1
model.zero_grad()
def step():
if scaler is not None:
with torch.cuda.amp.autocast():
loss = model(input_tensor).abs().sum()
scaler.scale(loss).backward()
else:
loss = model(input_tensor).abs().sum()
loss.backward()
with model.no_sync() if should_accumulate else suppress():
for _ in range(accumulate_steps - 1):
step()
step()
def check_parity(amp: bool):
# Any model works. Add one different buffer per rank
model = Sequential(Linear(2, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3))
model = Sequential(Linear(INPUTS, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3), Linear(3, 3))
model.register_buffer("test_buffer", torch.ones((1)) * rank)
model.to(device)
sharded_optimizer = OSS(params=model.parameters(), optim=torch.optim.SGD, lr=1e-3, momentum=0.99)
# Make sure that the model starts with non-trainable, so that we check for the buckets to be
# properly reassigned when/if this changes
next(model.parameters()).requires_grad = False
sharded_optimizer = OSS(params=model.parameters(), optim=torch.optim.SGD, lr=1e-5, momentum=0.99)
sharded_ddp_model = ShardedDataParallel(
module=model, sharded_optimizer=sharded_optimizer, broadcast_buffers=True
)
ddp_model_single = copy.deepcopy(model)
ddp_optimizer = torch.optim.SGD(ddp_model_single.parameters(), lr=1e-3, momentum=0.99)
ddp_model = DDP(ddp_model_single, device_ids=[rank], broadcast_buffers=True)
ddp_optimizer = torch.optim.SGD(ddp_model_single.parameters(), lr=1e-5, momentum=0.99)
ddp_model = DDP(ddp_model_single, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
ddp_scaler = TorchGradScaler() if amp else None
sharded_ddp_scaler = ShardedGradScaler() if amp else None
def check_same_model_params():
for pg, ddp_pg in zip(sharded_optimizer.param_groups, ddp_optimizer.param_groups):
for p, ddp_p in zip(pg["params"], ddp_pg["params"]):
assert torch.allclose(
p, ddp_p, atol=1e-3
), f"Model parameters differ in between DDP and ShardedDDP {p} {ddp_p}"
for b, ddp_b in zip(sharded_ddp_model.buffers(), ddp_model.buffers()):
assert torch.allclose(
b, ddp_b, atol=1e-3
), f"Model buffers differ in between DDP and ShardedDDP. AMP {amp}"
# The model should be synchronized in between the ranks at construction time, check that
check_same_model_params()
check_same_model_params(sharded_ddp_model, ddp_model)
# The models should stay the same in between the ranks
for i in range(10):
input_tensor = torch.rand((64, 2)).to(device)
# Typical training loop, check that we get the exact same results as DDP
for i in range(NUMBER_BATCHS):
input_tensor = torch.rand((BATCH_SIZE, INPUTS)).to(device)
def closure_ddp(input_tensor=input_tensor):
ddp_optimizer.zero_grad()
if ddp_scaler is not None:
with torch.cuda.amp.autocast():
ddp_loss = ddp_model(input_tensor).abs().sum()
ddp_scaler.scale(ddp_loss).backward()
else:
ddp_loss = ddp_model(input_tensor).abs().sum()
ddp_loss.backward()
return ddp_loss
return closure(ddp_model, ddp_scaler, input_tensor, accumulate)
def closure_sharded(input_tensor=input_tensor):
sharded_optimizer.zero_grad()
if sharded_ddp_scaler is not None:
with torch.cuda.amp.autocast():
sharded_loss = sharded_ddp_model(input_tensor).abs().sum()
sharded_ddp_scaler.scale(sharded_loss).backward()
else:
sharded_loss = sharded_ddp_model(input_tensor).abs().sum()
sharded_loss.backward()
return sharded_loss
return closure(sharded_ddp_model, sharded_ddp_scaler, input_tensor, accumulate)
# Step/scale both
if ddp_scaler is not None:
......@@ -210,13 +209,28 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
else:
sharded_optimizer.step(closure=closure_sharded)
check_same_model_params()
check_same_model_params(sharded_ddp_model, ddp_model, f"Rank: {rank} - Step {i} broke")
check_parity(amp=False)
# Flip the trainability of the first parameter back and forth
if i == 0 and change_train_graph:
next(sharded_ddp_model.parameters()).requires_grad = not next(
sharded_ddp_model.parameters()
).requires_grad
next(ddp_model.parameters()).requires_grad = not next(ddp_model.parameters()).requires_grad
check_same_model_params(sharded_ddp_model, ddp_model, f"Rank: {rank} - Trainability refresh {i} broke")
# Catch a version of pytorch which would not support AMP
# Test all combinations: AMP, Accumulate, Change train graph
amp_tests = [False]
if hasattr(torch.cuda.amp, "autocast"):
check_parity(amp=True)
amp_tests.append(True)
for accumulate in [False, True]:
for change_train_graph in [False, True]:
for amp in amp_tests:
print(
f"Checking configuration: accumulate {accumulate} - change train graph {change_train_graph} - amp {amp}"
)
check_parity(amp=amp, accumulate=accumulate, change_train_graph=change_train_graph)
dist.destroy_process_group()
......@@ -417,6 +431,8 @@ def run_test_ddp_sync_batch_norm(rank, world_size, backend, device, temp_file_na
model = Sequential(Linear(2, 3), torch.nn.BatchNorm1d(3), Linear(3, 3)).to(device)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
model.to(device) # in pytorch 1.5 syncBN switches to the default device/cpu
optimizer = OSS(params=model.parameters(), optim=torch.optim.SGD, lr=0.01, momentum=0.99)
ddp_model = ShardedDataParallel(model, optimizer)
......
tests/optim/test_oss.py
View file @ 13445c55
......@@ -11,7 +11,7 @@
import copy
from math import inf
import tempfile
from typing import Any, Type, cast
from typing import Any, Dict, Type, cast
import unittest
import numpy as np
......@@ -22,7 +22,7 @@ import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
import fairscale.optim as optim
from fairscale.utils.testing import skip_if_no_cuda, skip_if_py39_no_cuda, skip_if_single_gpu
from fairscale.utils.testing import check_same_model_params, skip_if_no_cuda, skip_if_py39_no_cuda, skip_if_single_gpu
BACKEND = dist.Backend.NCCL if torch.cuda.is_available() else dist.Backend.GLOO # type: ignore
DEVICE = "cuda" if torch.cuda.is_available() else torch.device("cpu")
......@@ -688,10 +688,6 @@ def run_state_dict_distributed(rank, world_size, tempfile_name):
model.zero_grad()
outputs = head(model(inputs))
def check_equal_models(message: str):
for param1, param2 in zip(model_oss1.parameters(), model_oss2.parameters()):
assert torch.allclose(param1, param2), message
# pull the current state, broadcast it to all ranks
sharded_optimizer2.consolidate_state_dict(recipient_rank=RECIPIENT_RANK) # all ranks
state_dict2 = sharded_optimizer2.state_dict() if rank == RECIPIENT_RANK else {}
......@@ -701,12 +697,16 @@ def run_state_dict_distributed(rank, world_size, tempfile_name):
sharded_optimizer2 = optim.OSS(model_oss2.parameters(), lr=1e6, momentum=0.0001)
sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
sharded_optimizer2.load_state_dict(state_dict2)
check_equal_models("parameters of the two identical models have diverged (before any steps)")
check_same_model_params(
model_oss1, model_oss2, "parameters of the two identical models have diverged (before any steps)"
)
# now take a step and check that parameters are equal
run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
check_equal_models("parameters of the two identical models have diverged (after stepping)")
check_same_model_params(
model_oss1, model_oss2, "parameters of the two identical models have diverged (after stepping)"
)
# save the state dict for one model only, then distribute to the other ranks
sharded_optimizer2.consolidate_state_dict(recipient_rank=RECIPIENT_RANK) # all ranks
......@@ -722,7 +722,9 @@ def run_state_dict_distributed(rank, world_size, tempfile_name):
# take a step
run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
check_equal_models("parameters of the two identical models have diverged (after consolidating)")
check_same_model_params(
model_oss1, model_oss2, "parameters of the two identical models have diverged (after consolidating)"
)
# save again for one rank, then distribute to the others
sharded_optimizer2.consolidate_state_dict(recipient_rank=RECIPIENT_RANK) # all ranks
......@@ -737,7 +739,9 @@ def run_state_dict_distributed(rank, world_size, tempfile_name):
# take a step
run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
check_equal_models("parameters of the two identical models have diverged (after reloading)")
check_same_model_params(
model_oss1, model_oss2, "parameters of the two identical models have diverged (after reloading)"
)
dist.destroy_process_group()
......@@ -768,7 +772,7 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
out_channels = 3
batch = 64
def check_optimizer_equivalence(optimizer: Type[torch.optim.Optimizer]):
def check_optimizer_equivalence(optimizer: Type[torch.optim.Optimizer], change_train_graph: bool = False):
# Any model works. Add one different buffer per rank
trunk = torch.nn.Sequential(torch.nn.Linear(in_channels, hidden), torch.nn.Linear(hidden, hidden))
trunk.register_buffer("test_buffer", torch.ones((1)) * rank)
......@@ -777,14 +781,14 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
head = torch.nn.Linear(hidden, out_channels).to(device)
# Define a model to be trained by OSS
oss_model = torch.nn.Sequential(trunk, head)
oss_module = torch.nn.Sequential(trunk, head)
oss_trainable_params = [
{"params": trunk.parameters(), "lr": 1e-5},
{"params": head.parameters(), "lr": 1e-4},
]
optimizer_settings = {}
if isinstance(optim, torch.optim.SGD):
optimizer_settings: Dict[Any, Any] = {}
if isinstance(optimizer, torch.optim.SGD):
optimizer_settings["momentum"] = 0.9
sharded_optimizer = optim.OSS(
......@@ -795,7 +799,7 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
**optimizer_settings,
)
oss_ddp_model = DDP(module=oss_model, device_ids=[rank], broadcast_buffers=True)
oss_ddp_model = DDP(module=oss_module, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
# Define a model to be trained by normal pytorch + DDP
ddp_trunk = copy.deepcopy(trunk)
......@@ -807,19 +811,7 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
{"params": ddp_head.parameters(), "lr": 1e-4},
]
ddp_optimizer = optimizer(ddp_trainable_params, **optimizer_settings) # type: ignore
ddp_model = DDP(module=ddp_module, device_ids=[rank], broadcast_buffers=True)
def check_same_model_params():
for pg, ddp_pg in zip(sharded_optimizer.param_groups, ddp_optimizer.param_groups):
for p, ddp_p in zip(pg["params"], ddp_pg["params"]):
assert torch.allclose(
p, ddp_p, atol=1e-3
), f"Model parameters differ in between Pytorch optim and OSS \n{p} {ddp_p}\nworld size {world_size}"
for b, ddp_b in zip(oss_ddp_model.buffers(), ddp_model.buffers()):
assert torch.allclose(
b, ddp_b
), f"Model buffers differ in between Pytorch optim and OSS\nworld size {world_size}"
ddp_model = DDP(module=ddp_module, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
def check_step():
input_tensor = torch.rand((batch, in_channels)).to(device)
......@@ -843,13 +835,21 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
loss_ddp, loss_sharded_optim
), f"Losses differ in between Pytorch optim and OSS\nworld size {world_size}"
check_same_model_params(oss_ddp_model, ddp_model)
# The model should be synchronized in between the ranks at construction time, check that
check_same_model_params()
check_same_model_params(oss_ddp_model, ddp_model)
# The models should stay the same in between ddp and sharded optimizer
for i in range(5):
check_step()
check_same_model_params()
# Check that altering the trainable parameters does not cause DDP and OSS to diverge
if change_train_graph:
# Flip the first parameter from trainable to non-trainable and vice-versa
next(ddp_module.parameters()).requires_grad = not next(ddp_module.parameters()).requires_grad
next(oss_module.parameters()).requires_grad = not next(oss_module.parameters()).requires_grad
# sharded_optimizer.refresh_trainable()
# Check that the checkpoints are compatible
# - get states
......@@ -864,10 +864,10 @@ def run_ddp_parity(rank, world_size, backend, temp_file_name):
# - run one step and check that the models are still the same
check_step()
check_same_model_params()
for opt in [torch.optim.SGD, torch.optim.Adam]:
check_optimizer_equivalence(opt)
for opt in [torch.optim.Adam, torch.optim.SGD]:
check_optimizer_equivalence(opt, change_train_graph=False)
check_optimizer_equivalence(opt, change_train_graph=True)
dist.destroy_process_group()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment