Skip to content

GitLab

Commit 61e92904 authored by chenzk's avatar chenzk
Browse files

v1.0

parents
Hide whitespace changes
Inline Side-by-side
Showing with 1808 additions and 0 deletions
optim/base.py 0 → 100644
View file @ 61e92904
from abc import ABC, abstractmethod
from collections import defaultdict
from copy import deepcopy
from itertools import chain
from typing import (
Any,
Callable,
DefaultDict,
Dict,
Hashable,
Iterable,
List,
Optional,
Set,
Tuple,
TypeVar,
Union,
)
import torch
from typing_extensions import TypeAlias
Args: TypeAlias = Tuple[Any, ...]
Kwargs: TypeAlias = Dict[str, Any]
StateDict: TypeAlias = Dict[str, Any]
class BaseOptimizer(ABC):
id_to_name: Dict[int, str]
param_groups: List[Dict[str, Any]]
@abstractmethod
def __getstate__(self):
...
@abstractmethod
def __setstate__(self, state):
...
@abstractmethod
def __repr__(self):
...
@abstractmethod
def zero_grad(self):
...
@abstractmethod
def state_dict_additional_keys(self) -> Set[str]:
"""Additional states we store in `state_dict`. It has to be a dictionary using parameter name as key, and a tensor as value."""
...
@abstractmethod
def state_dict(self) -> dict:
...
@abstractmethod
def load_state_dict(self, state_dict: dict, map_location: Optional[Union[str, torch.device]] = None) -> None:
...
@abstractmethod
def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
...
def inherit_from(self, cls) -> bool:
...
Optimizer = TypeVar("Optimizer", BaseOptimizer, torch.optim.Optimizer)
# Modified from torch.optim.Optimizer._process_value_according_to_param_policy
@staticmethod
def _process_value_according_to_param_policy(
param: torch.Tensor,
value: torch.Tensor,
param_id: int,
param_groups: List[Dict[Any, Any]],
map_location: Optional[Union[str, torch.device]],
key: Hashable = None,
) -> torch.Tensor:
# If map_location is specified, use it instead of param.device
target_device = map_location if map_location is not None else param.device
fused = False
capturable = False
assert param_groups is not None
for pg in param_groups:
if param_id in pg["params"]:
fused = pg["fused"] if "fused" in pg else False
capturable = pg["capturable"] if "capturable" in pg else False
break
if key == "step":
if capturable or fused:
return value.to(dtype=torch.float32, device=target_device)
else:
return value
else:
if param.is_floating_point():
return value.to(dtype=param.dtype, device=target_device)
else:
return value.to(device=target_device)
# Modified from torch.optim.Optimizer.load_state_dict
@torch._disable_dynamo
def custom_load_state_dict(self, state_dict: StateDict, map_location: Union[str, torch.device]) -> None:
r"""Loads the optimizer state.
Args:
state_dict (dict): optimizer state. Should be an object returned
from a call to :meth:`state_dict`.
map_location (str or torch.device, optional): Device where to load the optimizer states.
If None, states will be loaded to the same device as their corresponding parameters.
Default: None
"""
# shallow copy, to be consistent with module API
state_dict = state_dict.copy()
for pre_hook in self._optimizer_load_state_dict_pre_hooks.values():
hook_result = pre_hook(self, state_dict)
if hook_result is not None:
state_dict = hook_result
# Validate the state_dict
groups = self.param_groups
# Deepcopy as we write into saved_groups later to update state
saved_groups = deepcopy(state_dict["param_groups"])
if len(groups) != len(saved_groups):
raise ValueError("loaded state dict has a different number of " "parameter groups")
param_lens = (len(g["params"]) for g in groups)
saved_lens = (len(g["params"]) for g in saved_groups)
if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):
raise ValueError(
"loaded state dict contains a parameter group " "that doesn't match the size of optimizer's group"
)
# Update the state
id_map = dict(
zip(chain.from_iterable(g["params"] for g in saved_groups), chain.from_iterable(g["params"] for g in groups))
)
def _cast(param, value, param_id=None, param_groups=None, key=None):
r"""Make a deep copy of value, casting all tensors to device of param."""
if isinstance(value, torch.Tensor):
return _process_value_according_to_param_policy(param, value, param_id, param_groups, map_location, key)
elif isinstance(value, dict):
return {k: _cast(param, v, param_id=param_id, param_groups=param_groups, key=k) for k, v in value.items()}
elif isinstance(value, Iterable):
return type(value)(_cast(param, v, param_id=param_id, param_groups=param_groups) for v in value)
else:
return value
# Copy state assigned to params (and cast tensors to appropriate types).
# State that is not assigned to params is copied as is (needed for
# backward compatibility).
state: DefaultDict[torch.Tensor, Dict[Any, Any]] = defaultdict(dict)
for k, v in state_dict["state"].items():
if k in id_map:
param = id_map[k]
state[param] = _cast(param, v, param_id=k, param_groups=state_dict["param_groups"])
else:
state[k] = v
# Update parameter groups, setting their 'params' value
def update_group(group: Dict[str, Any], new_group: Dict[str, Any]) -> Dict[str, Any]:
new_group["params"] = group["params"]
return new_group
param_groups = [update_group(g, ng) for g, ng in zip(groups, saved_groups)]
self.__setstate__({"state": state, "param_groups": param_groups})
for post_hook in self._optimizer_load_state_dict_post_hooks.values():
post_hook(self)
optim/clip_grads.py 0 → 100644
View file @ 61e92904
from typing import Iterable, Optional, Tuple
import torch
import nanotron.distributed as dist
from nanotron import logging
from nanotron.optim.gradient_accumulator import GradientAccumulator
from nanotron.parallel.parameters import NanotronParameter
logger = logging.get_logger(__name__)
def clip_grad_norm(
mp_pg: dist.ProcessGroup,
named_parameters: Iterable[Tuple[str, NanotronParameter]],
max_norm: float,
grad_accumulator: Optional[GradientAccumulator],
norm_type: float = 2.0,
) -> torch.Tensor:
"""Clips gradients. Adapted from torch.nn.utils.clip_grad_norm_.
Norms are computed in fp32 precision to retain most accuracy.
Args:
mp_pg (dist.ProcessGroup): Process group for model parallel, ie all the ranks part of the same model replica (TP x PP)
named_parameters (Iterable[(str, Parameter)]): an iterable of named Parameters that will have gradients normalized.
grad_accumulator (GradientAccumulator): grad accumulator. If not None, in case of Zero1, we need to clip all fp32 grads
max_norm (float or int): max norm of the gradients
norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for infinity norm.
.. note:: In case parameters contains tied weights, we keep only a single copy of the gradient, but modify the
gradient of all tied weights.
"""
named_parameters = list(named_parameters)
world_rank = dist.get_rank()
# assert that all params require grad
for _, p in named_parameters:
assert p.requires_grad, "clip_grad_norm_ only supports Tensors that require grad"
if grad_accumulator is None:
grads = [
p.grad for _, p in named_parameters if not p.is_tied or world_rank == p.get_tied_info().global_ranks[0]
]
else:
# In case of FP32 Grad Accum, We need to clip all fp32 grads
grads = [
grad_accumulator.get_grad_buffer(name)
for name, p in named_parameters
if not p.is_tied or world_rank == p.get_tied_info().global_ranks[0]
]
# Calculate gradient norm
if norm_type == torch.inf:
if len(grads) > 0:
total_norm = torch.max(
torch.stack([torch.linalg.vector_norm(g.detach(), ord=torch.inf, dtype=torch.float) for g in grads])
)
else:
total_norm = torch.zeros([], dtype=torch.float, device=torch.device("cuda"))
dist.all_reduce(total_norm, group=mp_pg, op=dist.ReduceOp.MAX)
else:
if len(grads) > 0:
# TODO @nouamanetazi: Check if we should calculate norm per parameter (remove .pow(norm_type)
total_norm = torch.linalg.vector_norm(
torch.stack([torch.linalg.vector_norm(g.detach(), ord=norm_type, dtype=torch.float) for g in grads]),
ord=norm_type,
dtype=torch.float,
).pow(norm_type)
else:
total_norm = torch.zeros([], dtype=torch.float, device=torch.device("cuda"))
dist.all_reduce(total_norm, group=mp_pg, op=dist.ReduceOp.SUM)
total_norm.pow_(1.0 / norm_type)
# Scale gradients
clip_coef = max_norm / (total_norm + 1.0e-6)
# Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so
# avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization
# when the gradients do not reside in CPU memory.
clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
devices = {
param.grad.device if grad_accumulator is None else grad_accumulator.get_grad_buffer(name).device
for name, param in named_parameters
}
device_to_clip_coef_clamped = {device: clip_coef_clamped.to(device) for device in devices}
for name, param in named_parameters:
if grad_accumulator is None:
param.grad.detach().mul_(device_to_clip_coef_clamped[param.grad.device])
else:
grad_accumulator.get_grad_buffer(name).detach().mul_(
device_to_clip_coef_clamped[grad_accumulator.get_grad_buffer(name).device]
)
return total_norm
optim/gradient_accumulator.py 0 → 100644
View file @ 61e92904
import dataclasses
from abc import ABC, abstractmethod
from collections import OrderedDict
from contextlib import contextmanager
from typing import Callable, Dict, Iterator, Optional, Tuple
import torch
from torch.distributed import GradBucket
import nanotron.distributed as dist
from nanotron import logging
from nanotron.parallel.parameters import NanotronParameter
from nanotron.utils import get_untyped_storage, tensor_from_untyped_storage
logger = logging.get_logger(__name__)
class GradientAccumulator(ABC):
fp32_grads_allreduce_handle: Optional[torch.futures.Future]
@abstractmethod
def __init__(self, named_parameters: Iterator[Tuple[str, NanotronParameter]]):
...
@abstractmethod
def backward(self, loss: torch.Tensor):
...
@abstractmethod
def step(self):
...
@abstractmethod
def sync_gradients_across_dp(self, dp_pg: dist.ProcessGroup, reduce_op: dist.ReduceOp, reduce_scatter: bool):
...
@abstractmethod
def zero_grad(self):
...
@abstractmethod
def get_parameter_for_optimizer(self, name: str) -> NanotronParameter:
...
@abstractmethod
def get_grad_buffer(self, name: str) -> torch.Tensor:
...
@abstractmethod
def state_dict(self) -> Dict[str, torch.Tensor]:
...
@abstractmethod
def load_state_dict(self, state_dict: torch.Tensor):
...
class FP32GradientAccumulator(GradientAccumulator):
def __init__(
self,
named_parameters: Iterator[Tuple[str, NanotronParameter]],
grad_buckets_named_params: Optional[Iterator[Tuple[str, NanotronParameter]]] = None,
):
"""Create a gradient accumulator that will accumulate gradients in fp32.
Args:
named_parameters: The parameters that will be updated by the optimizer. In case of Zero 1, this is the parameters that will be updated in this DP rank.
grad_buckets_named_params: The parameters to accumulate gradients for. If None it defaults to `named_parameters`. In case of Zero 1, this should be all the parameters in the model.
Note: We use `grad_buckets_named_params` to keep grad buffers for all parameters even when Zero 1 is used. This is because we need to accumulate gradients for all parameters without having to reduce in every accumulation step.
Note: We make a fp32 copy of parameters during initialization. Therefore parameters need to be initialized or loaded from a checkpoint before constructing this gradient accumulator
"""
if grad_buckets_named_params is None:
named_parameters = list(named_parameters)
grad_buckets_named_params = named_parameters
# Initialize grad bucket
self.fp32_grad_buffers, self._contiguous_fp32_grad_buffer = self.build_grad_buffers(
named_parameters=grad_buckets_named_params
)
# Assign big buffer for weights + grad in fp32
segment_index = {}
length = 0
for name, param in named_parameters:
if not param.requires_grad:
continue
start = length
end_weight = start + param.numel()
assert name not in segment_index
segment_index[name] = (start, end_weight, param)
length = end_weight
big_flat_buffer = torch.empty(length, dtype=torch.float, device="cuda")
self.parameters = {
name: {
"fp32": big_flat_buffer[start_weight:end_weight].view_as(param),
"half": param,
}
for name, (start_weight, end_weight, param) in segment_index.items()
}
with torch.inference_mode():
for _, elt in self.parameters.items():
fp32_param = elt["fp32"]
half_param = elt["half"]
# Check that fp32 weights have the same memory representation as half precision weights
assert fp32_param.stride() == half_param.stride()
# Copy weights from half precision to full precision
fp32_param.copy_(half_param)
# Set requires_grad=True
fp32_param.requires_grad = True
self._is_accumulation_sync_step = False
# We need the last allreduce handle to make sure it finishes before the optimizer step
self.fp32_grads_allreduce_handle: Optional[torch.futures.Future] = None
def assign_param_offsets(self, param_name_to_offsets: Dict[str, Dict[int, Tuple[int, int]]], dp_rank: int):
"""To use only when you use with ZeRODistributedOptimizer"""
self.param_name_to_offsets = {
name: elt[dp_rank] for name, elt in param_name_to_offsets.items() if dp_rank in elt
}
def sync_gradients_across_dp(self, dp_pg: dist.ProcessGroup, reduce_op: dist.ReduceOp, reduce_scatter: bool):
if dp_pg.size() == 1:
# They are already synced
return
if reduce_scatter:
# Usually you need to run `all_reduce` in order for all gradients to be synced.
# However when the optimizer state are sharded, you really just need to scatter to ranks that are going to run the optimizer state.
# Effectively you replace a `all_reduce` with a `reduce_scatter` which should save an `all_gather` when using RING algorithm.
assert hasattr(self, "param_name_to_offsets")
named_offsets = sorted(self.param_name_to_offsets.items(), key=lambda x: x[0])
flat_grad_buffers = [self.fp32_grad_buffers[name]["fp32_grad"].view(-1) for name, _ in named_offsets]
dist.reduce_scatter_coalesced(
output_tensor_list=[
flat_grad_buffer[start_offset:end_offset]
for (_, (start_offset, end_offset)), flat_grad_buffer in zip(named_offsets, flat_grad_buffers)
],
input_tensor_lists=[
torch.split(
flat_grad_buffer,
split_size_or_sections=len(self.fp32_grad_buffers[name]["fp32_grad"].view(-1)) // dp_pg.size(),
)
for (name, _), flat_grad_buffer in zip(named_offsets, flat_grad_buffers)
],
group=dp_pg,
)
else:
dist.all_reduce(self._contiguous_fp32_grad_buffer, op=reduce_op, group=dp_pg)
@staticmethod
def build_grad_buffers(
named_parameters: Iterator[Tuple[str, NanotronParameter]],
) -> Tuple[Dict[str, Dict], torch.Tensor]:
"""Builds grad buffers for all model's parameters, independently of ZeRO sharding
Args:
named_parameters: Parameters to build buckets for. In case of Zero1, this should be all parameters.
Note:
In ZeRO-1, we need to accumulate grads for all parameters, because we need to allreduce all parameters' grads across DP at each sync step.
"""
named_parameters = [(name, param) for name, param in named_parameters if param.requires_grad]
needed_buffer_size = sum(param.numel() for _, param in named_parameters)
# important to have grads zeroed initially (see `self._accumulate_grad`)
contiguous_buffer_f32_gradients = torch.zeros(needed_buffer_size, dtype=torch.float, device="cuda")
untyped_storage = get_untyped_storage(contiguous_buffer_f32_gradients)
element_size = contiguous_buffer_f32_gradients.element_size()
# NOTE: Although `bias` can only exist on TP=0. It shouldn't be a problem here, because we only sync across DP
fp32_grad_buffers = OrderedDict() # keeps order of insertion
offset = 0
for name, param in named_parameters:
if not param.requires_grad:
continue
assert param.dtype != torch.float, f"Expected {name} not to be float"
assert param.is_contiguous(), f"Expected {name} to be contiguous"
next_offset = offset + param.numel() * element_size
fp32_grad_buffer = tensor_from_untyped_storage(
untyped_storage=untyped_storage[offset:next_offset], dtype=torch.float
)
fp32_grad_buffers[name] = {
"half": param,
# We create sliced tensors by also slicing storage.
# We need to specify "cuda" in order to share the same data storage, otherwise it build the tensor in "cpu" and copies over the data
"fp32_grad": fp32_grad_buffer.view_as(param),
}
offset = next_offset
return fp32_grad_buffers, contiguous_buffer_f32_gradients
def backward(self, loss: torch.Tensor):
result = loss.backward()
for name, elt in self.fp32_grad_buffers.items():
self._accumulate_grad(name=name, half_param=elt["half"])
return result
def _accumulate_grad(self, name: str, half_param: NanotronParameter) -> None:
"""Accumulate grad in fp32 and set the fp32 grad to the fp32 grad buffer, so that optimizer can update fp32 weights afterwards"""
assert half_param.grad is not None, f"Expected param {name} to have gradient."
fp32_grad = self.get_grad_buffer(name=name)
if self._is_accumulation_sync_step is False:
# WARNING: We assume fp32_grad_bucket is already zeroed
fp32_grad.add_(half_param.grad)
# In case _is_accumulation_sync_step = True: no need to add half gradients, because it's done in the allreduce hook
# TODO @thomasw21: Is it better to set to zero instead?
half_param.grad = None
# In the case an optimizer decides to set it to None, we need to re-assign previous buffer
if name in self.parameters:
fp32_param = self.parameters[name]["fp32"]
if hasattr(self, "param_name_to_offsets"):
if name not in self.param_name_to_offsets:
# When `name` isn't in `param_name_to_offsets` it means the slice is empty.
return
start_offset, end_offset = self.param_name_to_offsets[name]
grad = fp32_grad.view(-1)[start_offset:end_offset]
else:
grad = fp32_grad
fp32_param.grad = grad
@contextmanager
def no_sync(self):
"""A context manager to disable gradient synchronizations across
data-parallel ranks.
Note: if we use `no_sync` once, that means we're in DDP mode, and we switch the default of self._is_accumulation_sync_step to True.
"""
old_is_accumulation_sync_step = self._is_accumulation_sync_step
self._is_accumulation_sync_step = False
try:
yield
finally:
self._is_accumulation_sync_step = old_is_accumulation_sync_step
@torch.inference_mode()
def step(self):
"""Updates fp32 weights from fp32 grads.
In case where OptimizerFromGradientAccumulator and gradient_accumulator_builder are using different parameters (e.g ZeRO).
We need to update only the parameters that were updated by the optimizer.
"""
for name in self.parameters.keys():
fp32_param = self.parameters[name]["fp32"]
half_param = self.parameters[name]["half"]
# TODO @nouamane: should we use a fused kernel to copy?
# Copy weights from full precision to half precision
half_param.copy_(fp32_param)
def zero_grad(self):
# Full precision gradients are reset to zero/none after the underlying `optimiser.step`, so no need to reset.
for elt in self.fp32_grad_buffers.values():
half_param = elt["half"]
if half_param.grad is None:
continue
half_param.grad = None
# in case where self.parameters and self.fp32_grad_buffers are not the same (e.g we want to accumulate all DPs grads, and only sync at sync step)
self._contiguous_fp32_grad_buffer.zero_()
def get_parameter_for_optimizer(self, name: str) -> NanotronParameter:
return self.parameters[name]["fp32"]
def get_grad_buffer(self, name: str) -> torch.Tensor:
"""Returns the gradient of the parameter from the appropriate grad bucket."""
return self.fp32_grad_buffers[name]["fp32_grad"]
def state_dict(self) -> Dict[str, torch.Tensor]:
# We consider `fp32` parameters as a state of the gradient accumulator
return {name: elt["fp32"] for name, elt in self.parameters.items()}
def load_state_dict(self, state_dict: Dict[str, torch.Tensor]):
assert set(state_dict.keys()) == set(self.parameters.keys())
with torch.inference_mode():
for name, elt in self.parameters.items():
elt["fp32"].copy_(state_dict[name])
@dataclasses.dataclass
class FP32GradBucketManager:
"""Manages the fp32 gradient buckets.
Attributes:
dp_pg: The process group to allreduce gradients across.
accumulator: The gradient accumulator which keeps the gradient buffers.
bucket_id_to_fp32_grad_buckets_and_dependencies: A dictionary mapping bucket ids to:
- fp32 grad bucket (torch.Tensor)
- set of param ids that are in the bucket -> used to know when to delete the buffer
param_id_to_bucket_id: A dictionary mapping param ids to bucket ids."""
dp_pg: dist.ProcessGroup
accumulator: FP32GradientAccumulator
param_id_to_name: Dict[int, str]
def __post_init__(self):
self.accumulator._is_accumulation_sync_step = True
def get_fp32_accum_hook(
reduce_scatter: bool,
reduce_op: dist.ReduceOp = dist.ReduceOp.AVG,
) -> Callable:
"""Returns a DDP communication hook that performs gradient accumulation in fp32.
Args:
reduce_op: The reduction operation to perform.
"""
# s = torch.cuda.Stream()
def fp32_accum_hook(state: FP32GradBucketManager, bucket: GradBucket) -> torch.futures.Future[torch.Tensor]:
# nonlocal s
# DDP groups grads in GradBuckets. This hook is called throughout the bwd pass, once each bucket is ready to overlap communication with computation.
# See https://pytorch.org/docs/stable/ddp_comm_hooks.html#what-does-a-communication-hook-operate-on for more details.
dp_pg = state.dp_pg
accumulator = state.accumulator
param_id_to_name = state.param_id_to_name
# Add new incoming gradient
# with torch.cuda.stream(s):
for param, grad in zip(bucket.parameters(), bucket.gradients()):
name = param_id_to_name[id(param)]
fp32_grad_buffer = accumulator.get_grad_buffer(name)
fp32_grad_buffer.add_(grad.view_as(fp32_grad_buffer))
# sync across dp
if dp_pg.size() == 1:
fut = torch.futures.Future()
fut.set_result(bucket.buffer())
return fut
if reduce_scatter:
assert hasattr(accumulator, "param_name_to_offsets")
grad_buffer_tensor_list = [
accumulator.get_grad_buffer(param_id_to_name[id(param)]).view(-1) for param in bucket.parameters()
]
device = grad_buffer_tensor_list[0].device
dtype = grad_buffer_tensor_list[0].dtype
output_tensor_list = [
grad_buffer[slice(*accumulator.param_name_to_offsets[param_id_to_name[id(param)]])]
if param_id_to_name[id(param)] in accumulator.param_name_to_offsets
else torch.empty(0, dtype=dtype, device=device)
for grad_buffer, param in zip(grad_buffer_tensor_list, bucket.parameters())
]
input_tensor_lists = [
torch.split(grad_buffer, split_size_or_sections=len(grad_buffer) // dp_pg.size())
for grad_buffer in grad_buffer_tensor_list
]
dist.reduce_scatter_coalesced(
output_tensor_list=output_tensor_list,
input_tensor_lists=input_tensor_lists,
op=reduce_op,
group=dp_pg,
async_op=True,
)
else:
grad_buffer_tensor_list = [
accumulator.get_grad_buffer(param_id_to_name[id(param)]).view(-1) for param in bucket.parameters()
]
accumulator.fp32_grads_allreduce_handle = dist.all_reduce_coalesced(
grad_buffer_tensor_list, group=dp_pg, async_op=True, op=reduce_op
)
# we shouldn't wait for this future for the rest of the backward
# with torch.cuda.stream(s):
fut: torch.futures.Future[torch.Tensor] = torch.futures.Future()
half_grad_bucket = bucket.buffer()
fut.set_result(half_grad_bucket)
return fut # We don't care about the new half grad values, so we return the old ones
return fp32_accum_hook
optim/inherit_from_other_optimizer.py 0 → 100644
View file @ 61e92904
from functools import cache
from typing import Callable, Dict, Optional, Set, Union
import torch
from nanotron.optim.base import BaseOptimizer, Optimizer, custom_load_state_dict
class InheritFromOtherOptimizer(BaseOptimizer):
def __init__(self, optimizer: Optimizer, id_to_name: Dict[int, str]):
self.id_to_name = id_to_name
# if self.optimizer is from torch we replace load_state_dict with the one from torch
if isinstance(optimizer, torch.optim.Optimizer):
# Replace the load_state_dict method with our custom implementation that enables CPU offload
original_load_state_dict = optimizer.load_state_dict
optimizer.load_state_dict = (
lambda state_dict, map_location=None: custom_load_state_dict(
optimizer, state_dict, map_location=map_location
)
if map_location is not None
else original_load_state_dict(state_dict)
)
self.optimizer: Optimizer = optimizer
def __getstate__(self):
return self.optimizer.__getstate__()
def __setstate__(self, state):
return self.optimizer.__setstate__(state)
def __repr__(self):
return f"{self.__class__.__name__}({self.optimizer.__repr__()})"
def zero_grad(self):
return self.optimizer.zero_grad()
@cache
def state_dict_additional_keys(self) -> Set[str]:
if isinstance(self.optimizer, BaseOptimizer):
return self.optimizer.state_dict_additional_keys()
else:
return set()
def state_dict(self) -> dict:
return self.optimizer.state_dict()
def load_state_dict(self, state_dict: dict, map_location: Optional[Union[str, torch.device]] = None) -> None:
return self.optimizer.load_state_dict(state_dict, map_location=map_location)
def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
return self.optimizer.step(closure=closure)
def get_base_optimizer(self):
if isinstance(self.optimizer, torch.optim.Optimizer):
return self.optimizer
else:
return self.optimizer.get_base_optimizer()
@property
def param_groups(self):
return self.optimizer.param_groups
def inherit_from(self, cls):
if isinstance(self, cls):
return True
if isinstance(self.optimizer, InheritFromOtherOptimizer):
return self.optimizer.inherit_from(cls)
return False
optim/named_optimizer.py 0 → 100644
View file @ 61e92904
from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union
import torch
from nanotron import logging
from nanotron.optim.inherit_from_other_optimizer import InheritFromOtherOptimizer
logger = logging.get_logger(__name__)
class NamedOptimizer(InheritFromOtherOptimizer):
"""Mimics somewhat the torch optimizer API"""
def __init__(
self,
named_params_or_groups: Iterable[Union[Tuple[str, torch.Tensor], Dict[str, Any]]],
optimizer_builder: Callable[[Iterable[Dict[str, Any]]], torch.optim.Optimizer],
):
named_param_groups = list(named_params_or_groups)
if len(named_param_groups) == 0 or not isinstance(named_param_groups[0], dict):
named_param_groups = [{"named_params": named_param_groups}]
id_to_name = {}
params = []
for named_param_group in named_param_groups:
assert "named_params" in named_param_group
# Don't need to check that param_groups are overlapping since the optimizer will do it for me.
# https://github.com/pytorch/pytorch/blob/88b3810c94b45f5982df616e2bc4c471d173f491/torch/optim/optimizer.py#L473
id_to_name.update(
{id(param): name for name, param in named_param_group["named_params"] if id(param) not in id_to_name}
)
params.append(
{
**{k: v for k, v in named_param_group.items() if k != "named_params"},
"params": [param for _, param in named_param_group["named_params"]],
}
)
name_to_id = {v: k for k, v in id_to_name.items()}
assert len(id_to_name) == len(name_to_id)
# Sanity check
for param_group in params:
_params = param_group["params"]
for param in _params:
# https://github.com/pytorch/pytorch/issues/100701
assert param.numel() > 0
super().__init__(optimizer=optimizer_builder(params), id_to_name=id_to_name)
def state_dict(self) -> dict:
optim_state_dict = super().state_dict()
assert "names" not in optim_state_dict
state_id_to_name = {id(state): self.id_to_name[id(param)] for param, state in self.optimizer.state.items()}
optim_state_dict["names"] = {
index: state_id_to_name[id(state)] for index, state in optim_state_dict["state"].items()
}
return optim_state_dict
def load_state_dict(self, state_dict: dict, map_location: Optional[Union[str, torch.device]] = None) -> None:
assert set(self.id_to_name.values()) == set(
state_dict["names"].values()
), f"Elements don't match:\n - Elements in `self.id_to_name` that aren't in the other one: {set(self.id_to_name.values()) - set(state_dict['names'].values())}\n - Elements in `state_dict[\"names\"]` that aren't in the other one: {set(state_dict['names'].values()) - set(self.id_to_name.values())}"
assert len(state_dict["state"]) == len(
state_dict["names"]
), f"Number of params in loaded state dict ({len(state_dict['state'])}) doesn't match number of names ({len(state_dict['names'])})"
assert len(state_dict["state"]) > 0, "Loading empty state dict"
OPTIMIZER_STATE_KEYS = sorted(state_dict["state"][0].keys() - {"step"})
for key in OPTIMIZER_STATE_KEYS:
for k, state in state_dict["state"].items():
assert (
key in state
), f"Key {key} not found in state dict: {state} which corresponds to param_name: {state_dict['names'][k]}"
return super().load_state_dict(state_dict, map_location=map_location)
optim/optimizer_from_gradient_accumulator.py 0 → 100644
View file @ 61e92904
from functools import cache
from typing import Any, Callable, Dict, Iterable, Optional, Set, Tuple, Union
import torch
from nanotron.optim.base import BaseOptimizer
from nanotron.optim.gradient_accumulator import GradientAccumulator
from nanotron.optim.inherit_from_other_optimizer import InheritFromOtherOptimizer
from nanotron.parallel.parameters import NanotronParameter
class OptimizerFromGradientAccumulator(InheritFromOtherOptimizer):
def __init__(
self,
gradient_accumulator_builder: Callable[[Iterable[Tuple[str, NanotronParameter]]], GradientAccumulator],
named_params_or_groups: Iterable[Union[Tuple[str, torch.Tensor], Dict[str, Any]]],
optimizer_builder: Callable[[Iterable[Dict[str, Any]]], BaseOptimizer],
):
named_param_groups = list(named_params_or_groups)
if len(named_param_groups) == 0 or not isinstance(named_param_groups[0], dict):
named_param_groups = [{"named_params": named_param_groups}]
name_to_param = {}
for named_param_group in named_param_groups:
for name, param in named_param_group["named_params"]:
if name in name_to_param:
raise ValueError(f"Duplicate key. {name} is already in `name_to_param`")
else:
name_to_param[name] = param
# Build gradient accumulator
gradient_accumulator = gradient_accumulator_builder(name_to_param.items())
self.gradient_accumulator = gradient_accumulator
# Obtained new params depending on the gradient accumulator
converted_named_param_group = [
{
**{k: v for k, v in named_param_group.items() if k != "named_params"},
"named_params": [
(name, gradient_accumulator.get_parameter_for_optimizer(name))
for name, _ in named_param_group["named_params"]
],
}
for named_param_group in named_param_groups
]
optimizer = optimizer_builder(converted_named_param_group)
super().__init__(optimizer=optimizer, id_to_name=optimizer.id_to_name)
def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
loss = super().step(closure)
self.gradient_accumulator.step()
return loss
def zero_grad(self):
super().zero_grad()
return self.gradient_accumulator.zero_grad()
@cache
def state_dict_additional_keys(self) -> Set[str]:
return super().state_dict_additional_keys() | {"gradient_accumulator"}
def state_dict(self) -> dict:
state_dict = super().state_dict()
assert "gradient_accumulator" not in state_dict
state_dict["gradient_accumulator"] = self.gradient_accumulator.state_dict()
return state_dict
def load_state_dict(self, state_dict: dict, map_location: Optional[Union[str, torch.device]] = None) -> None:
gradient_accumulator_state_dict = state_dict.pop("gradient_accumulator")
super().load_state_dict(state_dict, map_location=map_location)
self.gradient_accumulator.load_state_dict(gradient_accumulator_state_dict)
optim/zero.py 0 → 100644
View file @ 61e92904
import itertools
import re
from collections import defaultdict
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
import numpy as np
import torch.optim
from functorch.dim import tree_map
from torch import nn
from tqdm import tqdm
from nanotron import distributed as dist
from nanotron import logging
from nanotron.distributed import ProcessGroup
from nanotron.logging import human_format, log_rank, warn_once
from nanotron.optim.base import BaseOptimizer
from nanotron.optim.inherit_from_other_optimizer import InheritFromOtherOptimizer
from nanotron.parallel import ParallelContext
from nanotron.parallel.parameters import NanotronParameter
logger = logging.get_logger(__name__)
class ZeroDistributedOptimizer(InheritFromOtherOptimizer):
"""Optimizer that handles partitioning of optimizer's states across DP ranks. See ZeRO Stage 1 in the paper https://arxiv.org/abs/1910.02054v3 for more details."""
def __init__(
self,
named_params_or_groups: Iterable[Union[Tuple[str, NanotronParameter], Dict[str, Any]]],
optimizer_builder: Callable[[Iterable[Dict[str, Any]]], BaseOptimizer],
dp_pg: ProcessGroup,
):
named_params_or_groups = list(named_params_or_groups)
if len(named_params_or_groups) == 0 or isinstance(named_params_or_groups[0], dict):
# case where named_params_or_groups is Iterable[Dict[str, Any]]
for d in named_params_or_groups:
assert (
"named_params" in d
), f"param_groups must contain a 'named_params' key, got a dict with keys {d.keys()}"
# keep only named_params_or_groups that require grads
named_params_or_groups = [
{
"named_params": [
(name, param) for name, param in named_param_group["named_params"] if param.requires_grad
],
**{k: v for k, v in named_param_group.items() if k != "named_params"},
}
for named_param_group in named_params_or_groups
]
self.zero_named_param_groups = named_params_or_groups
else:
# case where named_params_or_groups is Iterable[Tuple[str, NanotronParameter]]
# keep only named_params_or_groups that require grads
named_params_or_groups = [(name, param) for name, param in named_params_or_groups if param.requires_grad]
self.zero_named_param_groups = [{"named_params": named_params_or_groups}]
self.dp_pg = dp_pg # DP process group
# partition model's params across DP ranks.
# `self.param_name_to_dp_rank_offsets` sets mapping between each param inside self.named_params and its rank
# NOTE: some param_groups may have no params in the current rank. we still keep them in self.optimizer.param_groups
self.param_name_to_dp_rank_offsets = self._partition_parameters()
current_dp_rank = dist.get_rank(self.dp_pg)
param_groups_in_rank = [
{
"named_params": [
(
name,
get_sliced_tensor(
param=param,
start_offset=self.param_name_to_dp_rank_offsets[name][current_dp_rank][0],
end_offset=self.param_name_to_dp_rank_offsets[name][current_dp_rank][1],
),
)
for name, param in param_group["named_params"]
if current_dp_rank in self.param_name_to_dp_rank_offsets[name]
],
**{k: v for k, v in param_group.items() if k != "named_params"},
}
for param_group in self.zero_named_param_groups
]
# initialize rank's optimizer which is responsible for updating the rank's parameters
# NOTE: In case of ZeRO, `self.id_to_name` stores only names of parameters that are going to be updated by this DP rank's optimizer.
# NOTE: In case of ZeRO, `self.optimizer` will only get the parameters that are going to be updated by this DP's optimizer. Which
# means that `self.optimizer.param_groups` is only a subset of `self.param_groups`.
optimizer = optimizer_builder(param_groups_in_rank)
super().__init__(optimizer=optimizer, id_to_name=optimizer.id_to_name)
@torch.no_grad()
def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
"""Performs a single optimization step (parameter update)."""
# TODO: @nouamanetazi: handle syncing param groups attrs (e.g. if we update lr)
loss = super().step(closure=closure)
# calculate param_size (model) + param_size (grads) + 2*param_size/DP_if_zero1 (optim_states)
expected_allocated = sum(
param.numel() * param.element_size() * 2 + param.numel() * param.element_size() * 2 / self.dp_pg.size()
for named_param_group in self.zero_named_param_groups
for _, param in named_param_group["named_params"]
)
log_rank(
f"[After optim states allocation] Memory usage: {torch.cuda.memory_allocated() / 1024**2:.2f}MB "
f"(Expected 2*param_size + 2*param_size/DP_if_zero1={expected_allocated / 1024**2:.2f}MB). "
f"Peak reserved memory: {torch.cuda.max_memory_reserved() / 1024**2:.2f}MB",
logger=logger,
level=logging.DEBUG,
group=self.dp_pg,
rank=0,
)
# All gather updated params
self._all_gather_params()
return loss
def zero_grad(self):
"""Copied from `torch.optim.optimizer.zero_grad` with the only change of using `self.param_groups` instead of `self.optimizer.param_groups`
because we want to zero out the gradients of all model params (not just the params in the current rank)"""
super().zero_grad()
# TODO @thomasw21: This is a call to torch internal API, we need to fix this
foreach = False # self.optimizer.defaults.get("foreach", False)
# TODO @thomasw21: This is a call to torch internal API, we need to fix this
# if not hasattr(self.optimizer, "_zero_grad_profile_name"):
# self.optimizer._hook_for_profile()
if foreach:
per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))
# TODO @thomasw21: This is a call to torch internal API, we need to fix this
# with torch.autograd.profiler.record_function(self.optimizer._zero_grad_profile_name):
# zero out the gradients of all model params (not just the params in the current rank)
for named_param_group in self.zero_named_param_groups:
for _, p in named_param_group["named_params"]:
if p.grad is not None:
p.grad = None
if foreach:
for _, per_dtype_grads in per_device_and_dtype_grads.items():
for grads in per_dtype_grads.values():
torch._foreach_zero_(grads)
def _partition_parameters(self) -> Dict[str, Dict[int, Tuple[int, int]]]:
named_params = [
(name, param)
for named_param_group in self.zero_named_param_groups
for name, param in named_param_group["named_params"]
if param.requires_grad
]
# maps each model's param to the optimizer's dp rank that is responsible for updating it
# We assume that parameters can be sharded across DP, ie we can "split" a parameter in different DP. This does break some optimizers, like Adafactor and such.
# `param_name_to_dp_rank_offsets[name]` is a `Dict[int, Tuple[int, int]]` keys are dp_rank, and `Tuple[int, int]` are the offsets of the param belonging to this DP
param_name_to_dp_rank_offsets = {}
# NOTE: save the original shapes before flattening the params
# so that later on, we can reshape the params to their original shapes
# for topology-agnostic optimizer states loading
self._orig_param_shapes = {}
for name, param in named_params:
self._orig_param_shapes[name] = param.shape
for name, param in named_params:
# We assume parameter to be contiguous in order to have an easy way of sharding it.
assert param.is_contiguous(), f"Parameter {name} is not contiguous"
numel = param.numel()
padded_numel_per_dp = (numel - 1) // self.dp_pg.size() + 1
sizes = np.full(shape=(self.dp_pg.size()), fill_value=padded_numel_per_dp)
remainder = padded_numel_per_dp * self.dp_pg.size() - numel
# Last `remainder` indices has one less element
if remainder > 0:
# It's weird that `size[-0:]` returns the entire list instead of nothing
sizes[-remainder:] -= 1
end_offsets = np.cumsum(sizes)
assert len(end_offsets) == self.dp_pg.size()
assert end_offsets[-1] == numel, f"Somehow {end_offsets[-1]} != {numel}"
# We want start indices,
start_offsets = np.concatenate([[0], end_offsets[:-1]])
param_name_to_dp_rank_offsets[name] = {
dp_rank: (start_offsets[dp_rank], end_offsets[dp_rank])
for dp_rank in range(self.dp_pg.size())
if start_offsets[dp_rank] < end_offsets[dp_rank] # Only if the slice is not empty.
}
log_rank("[ZeRO sharding] Size of optimizer params per rank:", logger=logger, level=logging.INFO, rank=0)
all_numel = sum(
param_name_to_dp_rank_offsets[name][dp_rank][1] - param_name_to_dp_rank_offsets[name][dp_rank][0]
for name, param in named_params
for dp_rank in range(self.dp_pg.size())
if dp_rank in param_name_to_dp_rank_offsets[name]
)
for dp_rank in range(self.dp_pg.size()):
acc_numel = sum(
value[dp_rank][1] - value[dp_rank][0]
for value in param_name_to_dp_rank_offsets.values()
if dp_rank in value
)
log_rank(
f"[ZeRO sharding] DP Rank {dp_rank} has {human_format(acc_numel)} out of {human_format(all_numel)} ({0 if all_numel == 0 else acc_numel / all_numel * 100:.2f}%) params' optimizer states",
logger=logger,
level=logging.INFO,
rank=0,
)
return param_name_to_dp_rank_offsets
def _all_gather_params(self):
"""All gather updated params"""
all_named_tensors_to_gather = [
(name, param.view(-1))
for named_param_groups in self.zero_named_param_groups
for name, param in named_param_groups["named_params"]
]
if len(all_named_tensors_to_gather) == 0:
# No need to broadcast if there's nothing
return
if self.dp_pg.size() == 1:
# They should already be updated
return
current_dp_rank = dist.get_rank(self.dp_pg)
dist.all_gather_coalesced(
output_tensor_lists=[
[
tensor[slice(*self.param_name_to_dp_rank_offsets[name][dp_rank])]
if dp_rank in self.param_name_to_dp_rank_offsets[name]
else torch.empty(0, dtype=tensor.dtype, device=tensor.device)
for dp_rank in range(self.dp_pg.size())
]
for name, tensor in all_named_tensors_to_gather
],
input_tensor_list=[
tensor[slice(*self.param_name_to_dp_rank_offsets[name][current_dp_rank])]
if current_dp_rank in self.param_name_to_dp_rank_offsets[name]
else torch.empty(0, dtype=tensor.dtype, device=tensor.device)
for name, tensor in all_named_tensors_to_gather
],
group=self.dp_pg,
)
# Helpers
class SlicedFlatTensor(torch.Tensor):
"""Subclass of `torch.Tensor` that unable to define `grad` getter on a slice of a flattened tensor."""
# Based on torch/testing/_internal/logging_tensor.py
# https://github.com/pytorch/pytorch/issues/102337#issuecomment-1579673041
__torch_function__ = torch._C._disabled_torch_function_impl
@staticmethod
def get_sliced_flat_tensor(data, start_offset, end_offset):
with torch.no_grad():
return data.view(-1)[start_offset:end_offset]
@staticmethod
def __new__(cls, data, start_offset, end_offset):
sliced_tensor = cls.get_sliced_flat_tensor(data=data, start_offset=start_offset, end_offset=end_offset)
result = torch.Tensor._make_wrapper_subclass( # type: ignore[attr-defined]
cls,
sliced_tensor.size(),
strides=sliced_tensor.stride(),
storage_offset=sliced_tensor.storage_offset(),
# TODO: clone storage aliasing
dtype=sliced_tensor.dtype,
layout=sliced_tensor.layout,
device=sliced_tensor.device,
requires_grad=sliced_tensor.requires_grad,
)
return result
def __init__(self, data, start_offset, end_offset):
super().__init__()
# TODO @thomasw21: Make is so that you can never update this value
self.sliced_flat_tensor = self.get_sliced_flat_tensor(
data=data, start_offset=start_offset, end_offset=end_offset
)
self.orig_data = data
self.start_offset = start_offset
self.end_offset = end_offset
@classmethod
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
def unwrap(e):
return e.sliced_flat_tensor if isinstance(e, cls) else e
def never_wrap(e):
# Never re-wrap
return e
return tree_map(never_wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs)))
def _get_grad(self):
if self.orig_data.grad is None:
return None
with torch.no_grad():
return self.orig_data.grad.view(-1)[self.start_offset : self.end_offset]
def _set_grad(self, grad):
if grad is not None:
orig_grad = self._get_grad()
if orig_grad is None:
raise NotImplementedError(
"Trying to set gradient on a sliced tensor when the original tensor hasn't allocated the buffer for the gradient"
)
orig_grad.copy_(grad)
return
# TODO @thomasw21: This is unfortunately necessary since we might pass `SliceTensor` to the optimizer.
warn_once(
logger=logger,
msg="You're setting a `SlicedTensor` gradient to None. We're going to assume you meant to set the original tensor gradient to None.",
rank=0,
)
self.orig_data.grad = None
def _del_grad(self):
raise NotImplementedError
# TODO @thomasw21: Figure out why this function doesn't get inherited. https://github.com/pytorch/pytorch/issues/102337#issuecomment-1634363356
def data_ptr(self):
return self.sliced_flat_tensor.data_ptr()
grad = property(_get_grad, _set_grad, _del_grad)
def get_sliced_tensor(param: NanotronParameter, start_offset: int, end_offset: int):
# This allows us to create a leaf tensor, despite sharing the underlying storage
result = SlicedFlatTensor(data=param, start_offset=start_offset, end_offset=end_offset)
return result
def find_optim_index_from_param_name(
param_name: str,
# NOTE: (pp_rank, dp_rank, tp_rank) or (pp_rank, tp_rank)
ckp_sharded_optim_states: Union[Tuple[Tuple[int, int, int], torch.Tensor], Tuple[Tuple[int, int], torch.Tensor]],
is_zero1: bool,
pp_rank=0,
) -> int:
param_name = param_name.replace("module.", "")
# NOTE: since all shards have the same optim state names
# so we take the first shard (except optionally the pp dimension)
if is_zero1 is True:
# NOTE: (pp_rank, dp_rank, tp_rank)
OPTIM_STATE_INDEX_TO_PARAM_NAME = ckp_sharded_optim_states[(pp_rank, 0, 0)]["names"]
else:
# NOTE: (pp_rank, tp_rank)
OPTIM_STATE_INDEX_TO_PARAM_NAME = ckp_sharded_optim_states[(pp_rank, 0)]["names"]
return next((k for k, v in OPTIM_STATE_INDEX_TO_PARAM_NAME.items() if v == param_name), None)
def extract_parallel_ranks_from_shard_path(
shard_path: Path, is_zero1: bool
) -> Union[Tuple[int, int, int], Tuple[int, int]]:
"""Extract parallel ranks from shard path
For example, if the shard path is:
+ For ZeRO-1: /path/to/optimizer_pp-0-of-1_dp-0-of-2_tp-0-of-1.pt
then the function will return (0, 0, 0) (pp_rank, dp_rank, tp_rank)
For ZeRO-0: /path/to/optimizer_pp-0-of-1_tp-0-of-1.pt
then the function will return (0, 0) (pp_rank, tp_rank)
"""
if is_zero1 is True:
# TODO(xrsrke): use the same pattern as weight checkpoints
# in weight checkpoints, we do pp-rank-.... but here we only do pp-...
# TODO(xrsrke): don't hardcode this
pattern = r"optimizer_pp-(\d+)-of-\d+_dp-(\d+)-of-\d+_tp-(\d+)-of-\d+\.pt"
match = re.search(pattern, str(shard_path))
pp_rank, dp_rank, tp_rank = match.groups()
return int(pp_rank), int(dp_rank), int(tp_rank)
else:
# NOTE: this is zero0 checkpoint
pattern = r"pp-(\d+)-of-\d+_tp-(\d+)-of-\d+"
match = re.search(pattern, str(shard_path))
pp_rank, tp_rank = match.groups()
return int(pp_rank), int(tp_rank)
def merge_dp_shard_in_zero1_optimizer(
model: nn.Module,
optimizer_config,
shard_paths: List[Path],
parallel_context: ParallelContext,
map_location: Optional[str] = None,
) -> Dict[Tuple[int, int], Dict[str, torch.Tensor]]: # (pp_rank, tp_rank): param_name -> optim_state
assert (
optimizer_config["configs"]["param_name_to_dp_rank_offsets"] is not None
), "param_name_to_dp_rank_offsets is required"
checkpoint_pp_size = optimizer_config["parallelism"]["pp_size"]
checkpoint_tp_size = optimizer_config["parallelism"]["tp_size"]
ckp_sharded_optim_states = {}
for shard_path in shard_paths:
pp_rank, dp_rank, tp_rank = extract_parallel_ranks_from_shard_path(shard_path, is_zero1=True)
ckp_sharded_optim_states[(pp_rank, dp_rank, tp_rank)] = torch.load(shard_path, map_location=map_location)
param_name_to_dp_rank_offsets = optimizer_config["configs"]["param_name_to_dp_rank_offsets"]
optimizer_state_names = ckp_sharded_optim_states[(0, 0, 0)]["state"][0].keys()
def get_numel_of_unsharded_dp_param(param_name):
dp_offsets = param_name_to_dp_rank_offsets[param_name]
return max(int(value) for values in dp_offsets.values() for value in values)
def assign_shard_to_buffer(buffer, offset, value):
offset_start, offset_end = map(int, offset)
buffer[offset_start:offset_end] = value
param_names = sorted(model.state_dict().keys(), key=lambda x: x)
ckp_merged_dp_shards_optim_states = {}
for pp_rank, tp_rank in tqdm(
list(itertools.product(range(int(checkpoint_pp_size)), range(int(checkpoint_tp_size)))),
disable=dist.get_rank(parallel_context.world_pg) != 0,
desc="Merging ZeRO-1's shards across data parallel dimension",
):
# NOTE: filter only the shards that correspond to the current pp_rank and tp_rank
filtered_ckp_sharded_optim_states = {}
for (pp, dp, tp), ckp_optim_state in ckp_sharded_optim_states.items():
if pp == pp_rank and tp == tp_rank:
filtered_ckp_sharded_optim_states[dp] = ckp_optim_state
# NOTE: now merge the shards across data parallel dimension
# for each parameter, we need to merge all shards across data parallel dimension
merged_dp_shards_optim_states = {}
merged_dp_shards_optim_states["state"] = {}
for param_name in param_names:
unshard_dp_size = get_numel_of_unsharded_dp_param(param_name)
optim_state_index = find_optim_index_from_param_name(
param_name=param_name,
ckp_sharded_optim_states=ckp_sharded_optim_states,
is_zero1=True,
)
merged_dp_shards_optim_states["state"][optim_state_index] = {}
for state_name in optimizer_state_names:
unsharded_dp_buffer = torch.zeros(unshard_dp_size, device="cuda")
# NOTE: now merge all the params across data parallel dimension
for dp_rank, ckp_optim_state in filtered_ckp_sharded_optim_states.items():
# NOTE: extract the optimizer state of the current parameter
ckp_optim_state = ckp_optim_state["state"][optim_state_index]
ckp_offset = param_name_to_dp_rank_offsets[param_name][str(dp_rank)]
assign_shard_to_buffer(unsharded_dp_buffer, ckp_offset, ckp_optim_state[state_name])
# NOTE: in optimizer states, the "state" use an index to represent the parameter
# not the parameter name
merged_dp_shards_optim_states["state"][optim_state_index][state_name] = unsharded_dp_buffer
# NOTE: each dp shard has the same step
merged_dp_shards_optim_states["state"][optim_state_index]["step"] = ckp_optim_state["step"]
ckp_merged_dp_shards_optim_states[(pp_rank, tp_rank)] = merged_dp_shards_optim_states
# NOTE: each dp shard has the same names, and param_groups since it's the same tp shard
# the 0 in (pp_rank, 0, tp_rank) is the dp_rank
ckp_merged_dp_shards_optim_states[(pp_rank, tp_rank)]["names"] = ckp_sharded_optim_states[
(pp_rank, 0, tp_rank)
]["names"]
ckp_merged_dp_shards_optim_states[(pp_rank, tp_rank)]["param_groups"] = ckp_sharded_optim_states[
(pp_rank, 0, tp_rank)
]["param_groups"]
assert len(ckp_merged_dp_shards_optim_states) == int(checkpoint_pp_size) * int(
checkpoint_tp_size
), f"Expect {int(checkpoint_pp_size) * int(checkpoint_tp_size)} merged dp shards, got {len(ckp_merged_dp_shards_optim_states)}"
# NOTE: sanity check, make sure each merged checkpoint
# has the same dict key as the original checkpoint
for (pp_rank, tp_rank), ckp_optim_state in ckp_merged_dp_shards_optim_states.items():
# NOTE: we remove the gradient_accumulator key from sanity check
# because we don't merge gradient_accumulator states
missing_keys = set(ckp_optim_state.keys()) - set(ckp_sharded_optim_states[(pp_rank, 0, tp_rank)].keys())
assert (
len(missing_keys - {"gradient_accumulator"}) == 0
), "Expected the merged dp shards to have the same keys as the original dp shards, but merged dp shard misses: {}".format(
missing_keys
)
return ckp_merged_dp_shards_optim_states
parallel/__init__.py 0 → 100644
View file @ 61e92904
# flake8: noqa
from nanotron.parallel.context import ParallelContext
parallel/context.py 0 → 100644
View file @ 61e92904
import os
from typing import Literal, Tuple, Annotated
import numpy as np
import torch
import nanotron.distributed as dist
DistributedBackend = Literal["gloo", "mpi", "nccl"]
class ParallelContext:
def __init__(
self,
tensor_parallel_size: int,
pipeline_parallel_size: int,
data_parallel_size: int,
expert_parallel_size: int = 1,
backend: DistributedBackend = "nccl",
):
"""Initialize parallel context."""
num_gpus_per_model = tensor_parallel_size * pipeline_parallel_size * expert_parallel_size
world_size = int(os.environ["WORLD_SIZE"])
assert (
world_size % data_parallel_size == 0
), "The total number of processes must be divisible by the data parallel size."
assert world_size % num_gpus_per_model == 0, (
"The total number of processes must be divisible by"
"the number of GPUs per model (tensor_parallel_size * pipeline_parallel_size)."
)
if num_gpus_per_model * data_parallel_size != world_size:
raise ValueError(
f"The number of process requires to run all replicas ({num_gpus_per_model * data_parallel_size})",
f"must be equal to the world size ({world_size}).",
)
if not dist.is_available():
raise ValueError("torch.distributed is not available as a package, please install it.")
self.tensor_parallel_size = tensor_parallel_size
self.pipeline_parallel_size = pipeline_parallel_size
self.data_parallel_size = data_parallel_size
self.expert_parallel_size = expert_parallel_size
self._groups = {}
self.set_device()
assert backend == "nccl", "Only nccl backend is supported for now."
if not dist.is_initialized():
dist.initialize_torch_distributed()
world_size = int(os.getenv("WORLD_SIZE", "1"))
ranks = list(range(world_size))
process_group = dist.new_group(
ranks=ranks,
backend=dist.get_backend(),
)
self.world_pg = process_group
self._init_parallel_groups()
def _init_parallel_groups(self):
"""Initialize 3D parallelism's all process groups."""
dist.barrier()
world_size = int(os.environ["WORLD_SIZE"])
ranks = np.arange(0, world_size).reshape(
(
self.expert_parallel_size,
self.pipeline_parallel_size,
self.data_parallel_size,
self.tensor_parallel_size,
)
)
self.world_ranks_to_pg = {}
# Relevant process groups containing the current rank
self.tp_pg = self.create_new_group(ranks.transpose((0, 1, 2, 3)).reshape((-1, self.tensor_parallel_size)))
self.dp_pg = self.create_new_group(ranks.transpose((3, 0, 1, 2)).reshape((-1, self.data_parallel_size)))
self.pp_pg = self.create_new_group(ranks.transpose((2, 3, 0, 1)).reshape((-1, self.pipeline_parallel_size)))
self.expert_pg = self.create_new_group(ranks.transpose((1, 2, 3, 0)).reshape((-1, self.expert_parallel_size)))
# model parallel group = combination of tp and pp and exp for a given dp rank
self.mp_pg = self.create_new_group(
[ranks[:, :, dp_rank, :].reshape(-1) for dp_rank in range(self.data_parallel_size)]
)
self.tp_and_expert_pg = self.create_new_group(
[
ranks[:, pp_rank, dp_rank, :].reshape(-1)
for pp_rank in range(self.pipeline_parallel_size)
for dp_rank in range(self.data_parallel_size)
]
)
self.world_rank_matrix: np.ndarray = ranks
def create_new_group(self, all_groups_ranks: np.ndarray) -> dist.ProcessGroup:
dist.barrier()
rank = int(os.environ["RANK"])
new_group_containing_rank = None
for group_ranks in all_groups_ranks:
sorted_ranks = tuple(sorted(group_ranks))
# add new group to `world_ranks_to_pg`
if sorted_ranks not in self.world_ranks_to_pg:
new_group = dist.new_group(ranks=group_ranks)
self.world_ranks_to_pg[sorted_ranks] = new_group
else:
new_group = self.world_ranks_to_pg[sorted_ranks]
if rank in sorted_ranks:
new_group_containing_rank = new_group
dist.barrier()
return new_group_containing_rank
def set_device(self):
local_rank = int(os.getenv("LOCAL_RANK", "0"))
# NOTE: Set the device id.
# `torch.cuda.device_count` should return the number of device on a single node.
# We assume the nodes to be homogeneous (same number of gpus per node)
device_id = local_rank
torch.cuda.set_device(torch.cuda.device(device_id))
def get_local_ranks(self, world_rank: int) -> Tuple[int, int, int]:
return tuple(i.item() for i in np.where(self.world_rank_matrix == world_rank))
def destroy(self):
if not dist.is_initialized():
return
dist.barrier()
dist.destroy_process_group()
def get_global_rank(
self,
ep_rank: int,
pp_rank: int,
dp_rank: int,
tp_rank: int,
) -> np.int64:
"""
Get the global rank based on the specified ranks in different parallel groups.
:param ep_rank: int, Rank in the expert parallel group.
:param pp_rank: int, Rank in the pipeline parallel group.
:param dp_rank: int, Rank in the data parallel group.
:param tp_rank: int, Rank in the tensor parallel group.
:return: numpy.int64, The global rank.
"""
return self.world_rank_matrix[ep_rank, pp_rank, dp_rank, tp_rank]
parallel/data_parallel/utils.py 0 → 100644
View file @ 61e92904
from contextlib import contextmanager
from typing import Optional
import torch
from nanotron import distributed as dist
from nanotron.optim.gradient_accumulator import GradientAccumulator
from torch import nn
@contextmanager
def ddp_trigger_sync_in_bwd(model_ddp):
"""Trigger the sync of the gradients in the next backward pass of DDP model."""
assert isinstance(model_ddp, torch.nn.parallel.DistributedDataParallel)
old_require_backward_grad_sync = model_ddp.require_backward_grad_sync
old_require_forward_param_sync = model_ddp.require_forward_param_sync
model_ddp.require_backward_grad_sync = True
model_ddp.require_forward_param_sync = True
# https://github.com/pytorch/pytorch/blob/master/torch/csrc/distributed/c10d/reducer.cpp#L1325-L1356
model_ddp.reducer.prepare_for_backward([])
try:
yield
finally:
model_ddp.require_backward_grad_sync = old_require_backward_grad_sync
model_ddp.require_forward_param_sync = old_require_forward_param_sync
def sync_gradients_across_dp(
module: nn.Module,
dp_pg: dist.ProcessGroup,
reduce_op: dist.ReduceOp,
grad_accumulator: Optional[GradientAccumulator],
**sync_options,
):
"""Sync gradients across data parallelism.
Args:
module: The module to sync gradients for.
dp_pg: The data parallelism process group.
reduce_op: The reduce operation to use.
grad_accumulator: The gradient accumulator to use.
sync_options: Additional options given when using `grad_accumulator`. Please look at `GradientAccumulator.sync_gradients_across_dp` for documentation
"""
if grad_accumulator is not None:
# This is an optimized path that
grad_accumulator.sync_gradients_across_dp(dp_pg=dp_pg, reduce_op=reduce_op, **sync_options)
return
# Sync gradients
for name, param in module.named_parameters():
dist.all_reduce(param.grad, op=reduce_op, group=dp_pg)
parallel/parameters.py 0 → 100644
View file @ 61e92904
import dataclasses
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
import torch
from torch import nn
from nanotron import distributed as dist
from nanotron import logging
if TYPE_CHECKING:
from nanotron.models import NanotronModel
logger = logging.get_logger(__name__)
@dataclasses.dataclass
class SlicesPair:
local_slices: Tuple[slice, ...]
global_slices: Tuple[slice, ...]
@staticmethod
def slice_to_str(s: slice):
# e.g. slice(0, 10, 2) -> "0,10,2"
# e.g. slice(None, None, None) -> "None,None,None"
return ",".join(str(x) if x is not None else "None" for x in (s.start, s.stop, s.step))
@staticmethod
def str_to_slice(s: str):
return slice(*(int(x) if x != "None" else None for x in s.split(",")))
def __str__(self):
# e.g. local_slices (slice(0, 10, 2), slice(None, None, None)) -> "0,10,2|None,None,None"
local_slices_str = "|".join(map(self.slice_to_str, self.local_slices))
# e.g. global_slices (slice(0, 20, 4), slice(None, None, None)) -> "0,20,4|None,None,None"
global_slices_str = "|".join(map(self.slice_to_str, self.global_slices))
# e.g. "0,10,2|None,None,None#0,20,4|None,None,None"
return f"{local_slices_str}#{global_slices_str}"
@classmethod
def from_str(cls, string: str):
local_slices_str, global_slices_str = string.split("#")
local_slices = tuple(map(cls.str_to_slice, local_slices_str.split("|")))
global_slices = tuple(map(cls.str_to_slice, global_slices_str.split("|")))
return cls(local_slices, global_slices)
@classmethod
def tuple_to_str(cls, pairs):
# e.g. 2 SlicesPair, 1st SlicesPair local_slices "0,10,2|None,None,None" and global_slices "0,10,2|None,None,None"
# 2nd SlicesPair local_slices "0,20,4|None,None,None" and global_slices "0,40,8|None,None,None"
# -> "0,10,2|None,None,None#0,10,2|None,None,None;0,20,4|None,None,None#0,40,8|None,None,None"
return ";".join(map(str, pairs))
@classmethod
def tuple_from_str(cls, string: str):
return tuple(map(cls.from_str, string.split(";")))
@dataclasses.dataclass
class TiedInfo:
name: str
# name must be defined starting from `root_module` (e.g. root_module.dense0.dense1.weight)
root_module: nn.Module
global_ranks: Tuple[int, ...]
# None signifies that we do not reduce
reduce_op: Optional[dist.ReduceOp]
def get_full_name_from_model(self, model: nn.Module) -> str:
module_id_to_prefix = {id(module): f"{module_name}." for module_name, module in model.named_modules()}
# Fix the root_model
module_id_to_prefix[id(model)] = ""
return self.get_full_name_from_module_id_to_prefix(module_id_to_prefix)
def get_full_name_from_module_id_to_prefix(self, module_id_to_prefix: Dict[int, str]) -> str:
return f"{module_id_to_prefix[id(self.root_module)]}{self.name}" # this assumes root_module is part of module_id_to_prefix
@dataclasses.dataclass
class ShardedInfo:
global_ranks: Tuple[int, ...]
# Info of to what slice of the unsharded tensor (global_slices) the current sharded tensor corresponds (local_slices)
local_global_slices_pairs: Tuple[SlicesPair, ...]
# The shape of the unsharded tensor
unsharded_shape: Tuple[int, ...]
def is_tp_sharded(self, parallel_context) -> bool:
return set(dist.get_global_ranks(parallel_context.tp_pg)).issubset(set(self.global_ranks))
def is_expert_sharded(self, parallel_context) -> bool:
return set(dist.get_global_ranks(parallel_context.expert_pg)).issubset(set(self.global_ranks))
def is_dp_sharded(self, parallel_context):
return set(dist.get_global_ranks(parallel_context.dp_pg)).issubset(set(self.global_ranks))
class NanotronParameter(nn.Parameter):
"""Base class for all parameters in Nanotronmodels
A NanotronParameter can have specific properties:
- sharded: the parameter is considered to be `sharded` across multiple devices
- tied: the parameter is considered to be `tied` with other parameters. We sum gradients over those.
.. note::
Notes about tied weights:
- Tied weights means weights that need to be synced only within the same DP rank, regardless if they are part of TP strategy or just shared weights between two layers.
- Syncing tied weights usually require to sum gradients.
- Some weights are synced without needing to reduce grads over ranks. They can be in the same device (ex: enc/dec embeds in the same PP stage) or they can be duplicated across TP and duplicate the workload across TP ranks (ex: LN using traditional TP)
- Even if some weights don't need their grads to be reduced, it's still useful for them to be marked as tied. For example, current serialization format requires to mark them correctly.
"""
NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME = "__nanotron_metadata__"
NANOTRON_PARAMETER_METADATA_TIED_KEY = "tied"
NANOTRON_PARAMETER_METADATA_SHARDED_KEY = "sharded"
def __new__(cls, tensor: torch.Tensor, requires_grad: bool = True):
param = nn.Parameter.__new__(cls, data=tensor.data.detach(), requires_grad=requires_grad)
if isinstance(tensor, NanotronParameter):
# Check that we don't inherit a weird class
# We copy in order not to make in-place operation
assert type(tensor) == NanotronParameter
setattr(
param,
cls.NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME,
getattr(tensor, cls.NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME).copy(),
)
else:
setattr(param, cls.NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME, {})
return param
def _set_metadata(self, key: str, value: Any):
metadata = getattr(self, self.NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME)
if key in metadata:
raise ValueError(
f"We shouldn't override previous metadata. Key to be overridden: {key}, current metadata: {metadata}"
)
else:
metadata[key] = value
def mark_as_tied(
self,
name: str,
global_ranks: Tuple[int, ...],
reduce_op: Optional[dist.ReduceOp],
root_module: "NanotronModel",
):
self._set_metadata(
self.NANOTRON_PARAMETER_METADATA_TIED_KEY,
TiedInfo(name=name, global_ranks=global_ranks, reduce_op=reduce_op, root_module=root_module),
)
def get_tied_info(self) -> TiedInfo:
return getattr(self, self.NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME)[
self.NANOTRON_PARAMETER_METADATA_TIED_KEY
]
@property
def is_tied(self) -> bool:
return self.NANOTRON_PARAMETER_METADATA_TIED_KEY in getattr(
self, self.NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME
)
def mark_as_sharded(
self,
global_ranks: Tuple[int, ...],
local_global_slices_pairs: Tuple[SlicesPair, ...],
unsharded_shape: Tuple[int, ...],
):
self._set_metadata(
self.NANOTRON_PARAMETER_METADATA_SHARDED_KEY,
ShardedInfo(
global_ranks=global_ranks,
local_global_slices_pairs=local_global_slices_pairs,
unsharded_shape=unsharded_shape,
),
)
def get_sharded_info(self) -> ShardedInfo:
return getattr(self, self.NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME)[
self.NANOTRON_PARAMETER_METADATA_SHARDED_KEY
]
@property
def is_sharded(self) -> bool:
return self.NANOTRON_PARAMETER_METADATA_SHARDED_KEY in getattr(
self, self.NANOTRON_PARAMETER_METADATA_ATTRIBUTE_NAME
)
def sanity_check(root_module: nn.Module):
"""Makes sure that the module is in Nanotronformat
Format:
- all parameters are `NanotronParameter`, this allows us to add metadata to a parameter.
"""
for name, param in root_module.named_parameters():
if not isinstance(param, NanotronParameter):
raise ValueError(
f"Nanotronrequires model to be in Nanotronformat, ie all parameters are required to be a NanotronParameter. {name} isn't."
)
parallel/pipeline_parallel/README.md 0 → 100644
View file @ 61e92904
## Pipeline parallelism
We choose to mimic the "torch" eager semantics:
- Declare module/blocks at init time
- Declare edges during forward
# Scheduling
# All forward, all backward (easy, but memory expensive)
# 1f1b (much nicer)
We're going to assume that all Pipeline blocks are assigned to a rank in a contiguous manner.
Warmup:
```
Rank 1: [forward(), forward(), forward(), forward(), backward()]
Rank 2: [forward(), forward(), forward(), forward(), backward(), backward()]
Rank 3: [forward(), forward(), forward(), forward(), backward(), backward(), backward()]
Rank 4: [forward(), backward(), forward(), backward(), forward(), backward()]
```
// TODO @thomasw21: How do we extrapolate this notion to a tree. Not sure exactly, but topological ordering should be fine
# TODOs:
- [ ] passing activation that don't require backward screws me as 1f1b works because you have the same number of forward and the same number of backward (in the stage sense)
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