mirror of https://github.com/hpcaitech/ColossalAI
441 lines
16 KiB
Python
441 lines
16 KiB
Python
import contextlib
|
|
import copy
|
|
import gc
|
|
import pprint
|
|
from typing import Callable, List, Optional, Union
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils._pytree import tree_map
|
|
|
|
from colossalai.device.device_mesh import DeviceMesh
|
|
from colossalai.fx.profiler import MetaTensor
|
|
from colossalai.tensor.shape_consistency import ShapeConsistencyManager
|
|
from colossalai.tensor.sharding_spec import ShardingSpec
|
|
|
|
# reference: https://pytorch.org/cppdocs/notes/tensor_creation.html
|
|
_TorchFactoryMethod = [
|
|
"arange",
|
|
"empty",
|
|
"eye",
|
|
"full",
|
|
"linspace",
|
|
"logspace",
|
|
"ones",
|
|
"rand",
|
|
"randn",
|
|
"randint",
|
|
"randperm",
|
|
"zeros",
|
|
"tensor",
|
|
]
|
|
|
|
orig_empty = torch.empty # avoid override
|
|
|
|
scm = ShapeConsistencyManager()
|
|
|
|
|
|
class LazyTensor(torch.Tensor):
|
|
"""A naive implementation of LazyTensor (https://arxiv.org/pdf/2102.13267.pdf).
|
|
|
|
Usage:
|
|
1. Use ``LazyTensor`` instead of ``torch.Tensor``.
|
|
>>> x = LazyTensor(torch.zeros, 2, 3)
|
|
>>> x += 1
|
|
>>> y = x * x
|
|
>>> y = y.cuda().half()
|
|
>>> y[0, 0] = 0
|
|
>>> y = y.materialize() # materialize the tensor
|
|
>>> print(y)
|
|
tensor([[0., 1., 1.],
|
|
[1., 1., 1.]], device='cuda:0', dtype=torch.float16)
|
|
|
|
2. Generate ``MetaTensor`` from ``LazyTensor``
|
|
>>> x = LazyTensor(torch.zeros, 2, 3)
|
|
>>> x.reshape(3, 2)
|
|
>>> x = x.traceable() # generate ``MetaTensor``
|
|
>>> print(x)
|
|
MetaTensor(..., size=(3, 2), device=cpu, dtype=torch.float32)
|
|
|
|
3. Use ``LazyTensor`` to generate sharded ``nn.Parameter``.
|
|
>>> x = LazyTensor(torch.zeros, 2, 3)
|
|
>>> x.spec = ... # some ``ShardingSpec``
|
|
>>> x.distribute() # distribute the tensor according to the ``ShardingSpec``
|
|
|
|
Warnings:
|
|
1. Cases that ``LazyTensor`` can't deal with.
|
|
>>> x = LazyTensor(torch.ones, 2, 3)
|
|
>>> x[0, 0] = -x[0, 0] # this will cause infinite recursion
|
|
|
|
2. ``LazyTensor.materialize()`` can't be called multiple times.
|
|
>>> x = LazyTensor(torch.ones, 2, 3)
|
|
>>> x.materialize()
|
|
>>> x.materialize() # this is disallowed
|
|
"""
|
|
|
|
_repr = True
|
|
_meta_data: Optional[MetaTensor] = None # shape, dtype, device
|
|
_cached_data: Optional[torch.Tensor] = None # materialized data
|
|
|
|
@staticmethod
|
|
def __new__(cls, func, *args, dtype=None, device=None, **kwargs):
|
|
elem = func(*args, dtype=dtype, device='meta', **kwargs)
|
|
r = torch.Tensor._make_wrapper_subclass(cls,
|
|
elem.size(),
|
|
strides=elem.stride(),
|
|
storage_offset=elem.storage_offset(),
|
|
dtype=elem.dtype,
|
|
layout=elem.layout,
|
|
device=device if device is not None else torch.device('cpu'),
|
|
requires_grad=elem.requires_grad)
|
|
r._meta_data = MetaTensor(elem, fake_device=device)
|
|
return r
|
|
|
|
def __init__(self, func, *args, dtype=None, device=None, **kwargs):
|
|
self._factory_method = (func, args, {'dtype': dtype, 'device': device, **kwargs}) # (func, args, kwargs)
|
|
self._cached_buffer = list() # (func, args, kwargs)
|
|
self._spec = None
|
|
self._data = self
|
|
|
|
def __repr__(self):
|
|
if self._repr:
|
|
# avoid recursive representation
|
|
self.__class__._repr = False
|
|
s = f'LazyTensor(..., size={tuple(self._meta_data.shape)}, device={self._meta_data.device}, dtype={self._meta_data.dtype})\n'\
|
|
f'factory method: {self._factory_method}\n'\
|
|
f'cached: {pprint.pformat(self._cached_buffer) if self._cached_data is None else self._cached_data}\n'\
|
|
f'spec: {self._spec}'
|
|
self.__class__._repr = True
|
|
return s
|
|
else:
|
|
return 'LazyTensor(...)'
|
|
|
|
def materialize(self) -> torch.Tensor:
|
|
"""Materialize the ``LazyTensor`` to ``torch.Tensor``.
|
|
|
|
Warnings:
|
|
Calling ``self.materialize()`` will clear all cached sequence and factory method,
|
|
because we don't allow materialize the same ``LazyTensor`` twice.
|
|
This is mentioned in the paper: https://arxiv.org/pdf/2102.13267.pdf (Part 4.3).
|
|
|
|
Returns:
|
|
torch.Tensor: The materialized tensor.
|
|
"""
|
|
target = self._data._realize_cached_data()
|
|
if isinstance(self, nn.Parameter):
|
|
target = nn.Parameter(target, requires_grad=self.requires_grad)
|
|
self._clear_all()
|
|
return target
|
|
|
|
def traceable(self) -> MetaTensor:
|
|
"""Generate ``MetaTensor`` from ``LazyTensor``. (Mostly for tracing)
|
|
|
|
Returns:
|
|
MetaTensor: The generated ``MetaTensor``.
|
|
"""
|
|
if isinstance(self, nn.Parameter):
|
|
return nn.Parameter(self._meta_data, requires_grad=self.requires_grad)
|
|
else:
|
|
return self._meta_data
|
|
|
|
def distribute(self) -> torch.Tensor:
|
|
"""Distribute the ``LazyTensor`` according to the ``ShardingSpec``.
|
|
|
|
Returns:
|
|
torch.Tensor: The sharded tensor.
|
|
"""
|
|
if self._spec is None:
|
|
raise RuntimeError('ShardingSpec is not set for\n{self}')
|
|
spec, device_mesh = self._spec, self._spec.device_mesh
|
|
target = self.materialize()
|
|
|
|
# TODO(some man): better not be coupled with auto-parallel
|
|
target.data = scm.apply_for_autoparallel_runtime(target.data, ShardingSpec(device_mesh, target.shape, {}),
|
|
spec).detach().clone()
|
|
return target
|
|
|
|
def _realize_cached_data(self) -> torch.Tensor:
|
|
# self._cached_data should be generated after the first call of this function
|
|
if self._cached_data is None:
|
|
if self._factory_method is not None:
|
|
# apply factory method
|
|
func, args, kwargs = self._factory_method
|
|
|
|
# apply cached sequence
|
|
self._cached_data = self._apply_cache_buffer(func(*args, **kwargs))
|
|
else:
|
|
# apply cached sequence only
|
|
self._cached_data = self._apply_cache_buffer()
|
|
return self._cached_data
|
|
|
|
def _apply_cache_buffer(self, target=None) -> torch.Tensor:
|
|
# dump all cached sequence
|
|
# super-dainiu: support methods for single Tensor only
|
|
def replace(x):
|
|
if x is self:
|
|
return target
|
|
elif isinstance(x, LazyTensor):
|
|
return x._realize_cached_data()
|
|
return x
|
|
|
|
packed = None
|
|
|
|
for (func, args, kwargs) in self._cached_buffer:
|
|
if func == torch.Tensor.requires_grad_:
|
|
packed = func, args, kwargs # requires grad should be set at last
|
|
else:
|
|
o = func(*tree_map(replace, args), **tree_map(replace, kwargs))
|
|
target = o if isinstance(o, torch.Tensor) else target # if func returns non-Tensor, discard the value
|
|
|
|
# super-dainiu: set requires_grad after all inplace-ops are done
|
|
if packed is not None:
|
|
func, args, kwargs = packed
|
|
func(*tree_map(replace, args), **tree_map(replace, kwargs))
|
|
|
|
return target
|
|
|
|
# clear all means:
|
|
# 1. clear factory method
|
|
# 2. clear cached sequence
|
|
# 3. clear cached data
|
|
def _clear_all(self):
|
|
self._cached_data = None
|
|
self._cached_buffer = None
|
|
self._data = None
|
|
gc.collect() # avoid memory leak
|
|
|
|
# cache everything with __torch_function__
|
|
@classmethod
|
|
def __torch_function__(cls, func, types, args=(), kwargs=None):
|
|
if kwargs is None:
|
|
kwargs = {}
|
|
target = None
|
|
|
|
if isinstance(func, torch._C.ScriptMethod):
|
|
|
|
def unwrap(x):
|
|
if isinstance(x, LazyTensor):
|
|
return x._meta_data
|
|
return x
|
|
|
|
target: LazyTensor = args[0].clone()
|
|
target._cached_buffer.append((func, args, kwargs))
|
|
target._meta_data = getattr(target._meta_data, func.name)(*tree_map(unwrap, args[1:]),
|
|
**tree_map(unwrap, kwargs))
|
|
|
|
else:
|
|
|
|
def unwrap(x):
|
|
nonlocal target
|
|
if isinstance(x, LazyTensor):
|
|
target = x if (func.__name__.endswith('_') and not (func.__name__.endswith('__'))
|
|
or func.__name__ == "__setitem__") else x.clone()
|
|
target._cached_buffer.append((func, args, kwargs))
|
|
return x._meta_data
|
|
return x
|
|
|
|
args = tree_map(unwrap, args)
|
|
kwargs = tree_map(unwrap, kwargs)
|
|
o = func(*args, **kwargs)
|
|
|
|
if isinstance(o, MetaTensor):
|
|
target._meta_data = o
|
|
return target
|
|
else:
|
|
return o
|
|
|
|
@classmethod
|
|
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
|
|
pass # skip
|
|
|
|
def clone(self) -> "LazyTensor":
|
|
"""Create a new ``LazyTensor`` with same cached sequence and factory method.
|
|
|
|
Returns:
|
|
LazyTensor: the new ``LazyTensor``
|
|
"""
|
|
target = LazyTensor(orig_empty, 0, dtype=self._meta_data.dtype, device=self._meta_data.device)
|
|
target._factory_method = None
|
|
target._cached_buffer = list()
|
|
target._meta_data = self._meta_data.clone()
|
|
target._cached_data = self._cached_data.clone() if self._cached_data is not None else None
|
|
target._spec = copy.deepcopy(self._spec)
|
|
return target
|
|
|
|
def detach(self) -> "LazyTensor":
|
|
target = self.clone()
|
|
target._cached_buffer.append((torch.Tensor.detach_, (self,), {}))
|
|
return target
|
|
|
|
@property
|
|
def spec(self) -> ShardingSpec:
|
|
return self._spec
|
|
|
|
@spec.setter
|
|
def spec(self, other: ShardingSpec):
|
|
self._spec = other
|
|
|
|
@property
|
|
def data(self) -> "LazyTensor":
|
|
return self._data.detach()
|
|
|
|
@data.setter
|
|
def data(self, other: "LazyTensor") -> "LazyTensor":
|
|
"""This avoid the following infinite recursion, which is very common in ``nn.Module`` initialization.
|
|
|
|
Usage:
|
|
>>> a = LazyTensor(torch.empty, 0, dtype=torch.float32, device='cpu')
|
|
>>> b = a.cuda()
|
|
>>> a.data = b
|
|
"""
|
|
self._data = other
|
|
|
|
|
|
class LazyInitContext():
|
|
"""Context manager for lazy initialization. Enables initializing the model without allocating real memory.
|
|
|
|
Usage:
|
|
1. The model is initialized, but no real memory is allocated.
|
|
>>> ctx = LazyInitContext()
|
|
>>> with ctx:
|
|
>>> model = MyModel().cuda()
|
|
|
|
2. The model is initialized with ``MetaTensor`` as weights, but still no real memory is allocated.
|
|
>>> with ctx.traceable(model):
|
|
>>> gm = symbolic_trace(model, meta_args=meta_args)
|
|
>>> # Solve the execution strategy and apply the strategy to the model
|
|
>>> strategy = StrategyAndSpec()
|
|
|
|
3. The model is initialized with ``torch.Tensor`` as weights, and real memory is allocated. (single device)
|
|
>>> model = ctx.materialize(model)
|
|
|
|
3. The model is initialized with sharded ``torch.Tensor`` as weights, and real memory is allocated. (distributed scenario)
|
|
>>> model = apply_strategy_to_all_params(model, strategy)
|
|
>>> model = ctx.distribute(model)
|
|
|
|
Warnings:
|
|
This API is still experimental and further modifications can be made to it.
|
|
For example:
|
|
1. Quantization strategies can be applied before allocating real memory.
|
|
2. Lazy initialization seems slower than normal initialization.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.overrides = {}
|
|
|
|
def __enter__(self):
|
|
|
|
def wrap_factory_method(target):
|
|
# factory functions (eg. torch.empty())
|
|
def wrapper(*args, **kwargs):
|
|
return LazyTensor(target, *args, **kwargs)
|
|
|
|
return wrapper, target
|
|
|
|
def wrap_factory_like_method(orig_target, target):
|
|
# factory_like functions (eg. torch.empty_like())
|
|
def wrapper(*args, **kwargs):
|
|
orig_t = args[0]
|
|
return LazyTensor(orig_target, *args[1:], device=orig_t.device, dtype=orig_t.dtype, **kwargs)
|
|
|
|
return wrapper, target
|
|
|
|
self.overrides = {
|
|
target: wrap_factory_method(getattr(torch, target))
|
|
for target in _TorchFactoryMethod
|
|
if callable(getattr(torch, target, None))
|
|
}
|
|
|
|
self.overrides.update({
|
|
target + '_like': wrap_factory_like_method(getattr(torch, target), getattr(torch, target + '_like'))
|
|
for target in _TorchFactoryMethod
|
|
if callable(getattr(torch, target + '_like', None))
|
|
})
|
|
|
|
for name, (wrapper, orig) in self.overrides.items():
|
|
setattr(torch, name, wrapper)
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
for name, (wrapper, orig) in self.overrides.items():
|
|
setattr(torch, name, orig)
|
|
|
|
@staticmethod
|
|
def materialize(module: torch.nn.Module):
|
|
"""Initialize all ``nn.Parameter`` from ``LazyTensor``.
|
|
|
|
Args:
|
|
module (torch.nn.Module): Target ``nn.Module``
|
|
"""
|
|
|
|
@torch.no_grad()
|
|
def init_recursively(module: nn.Module):
|
|
# recursively initialize the module
|
|
for mod in module.children():
|
|
init_recursively(mod)
|
|
|
|
# initialize tensors directly attached to the current module
|
|
for name, param in module.named_parameters(recurse=False):
|
|
setattr(module, name, param.materialize())
|
|
|
|
for name, buf in module.named_buffers(recurse=False):
|
|
setattr(module, name, buf.materialize())
|
|
|
|
init_recursively(module)
|
|
return module
|
|
|
|
@staticmethod
|
|
def distribute(module: torch.nn.Module):
|
|
"""Initialize and shard all ``nn.Parameter`` from ``LazyTensor``.
|
|
|
|
Args:
|
|
module (torch.nn.Module): Sharded target ``nn.Module``
|
|
"""
|
|
|
|
@torch.no_grad()
|
|
def init_recursively(module: nn.Module):
|
|
# recursively initialize the module
|
|
for mod in module.children():
|
|
init_recursively(mod)
|
|
|
|
# initialize tensors directly attached to the current module
|
|
for name, param in module.named_parameters(recurse=False):
|
|
setattr(module, name, param.distribute())
|
|
|
|
for name, buf in module.named_buffers(recurse=False):
|
|
setattr(module, name, buf.distribute())
|
|
|
|
init_recursively(module)
|
|
return module
|
|
|
|
@staticmethod
|
|
@contextlib.contextmanager
|
|
def traceable(module: torch.nn.Module):
|
|
"""Initialize all ``nn.Parameters`` as ``MetaTensor``. This enables ``ColoTracer`` with control flow.
|
|
|
|
Args:
|
|
module (torch.nn.Module): Traceable ``nn.Module`` with ``MetaTensor`` as parameters.
|
|
"""
|
|
orig_val = dict()
|
|
|
|
def init_recursively(module: nn.Module):
|
|
# recursively initialize the module
|
|
for mod in module.children():
|
|
init_recursively(mod)
|
|
|
|
# initialize tensors directly attached to the current module
|
|
for name, param in module.named_parameters(recurse=False):
|
|
setattr(module, name, param.traceable())
|
|
orig_val[(module, name)] = param
|
|
|
|
for name, buf in module.named_buffers(recurse=False):
|
|
setattr(module, name, buf.traceable())
|
|
orig_val[(module, name)] = buf
|
|
|
|
init_recursively(module)
|
|
|
|
yield
|
|
|
|
# restore original values
|
|
for (module, name), val in orig_val.items():
|
|
setattr(module, name, val)
|