wehub-resource-sync
322 lines
10 KiB
Python
322 lines
10 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
import paddle.autograd as imperative_base
|
|
from paddle.framework import (
|
|
_current_expected_place_,
|
|
)
|
|
from paddle.incubate.tensor.manipulation import (
|
|
async_offload_with_offset,
|
|
create_async_load,
|
|
)
|
|
|
|
alignment = {
|
|
"gpu": 256,
|
|
"npu": 256,
|
|
"xpu": 256,
|
|
}
|
|
|
|
align = {
|
|
paddle.float16: 2,
|
|
paddle.bfloat16: 2,
|
|
paddle.float32: 4,
|
|
}
|
|
|
|
__current_device_type__ = None
|
|
|
|
|
|
def _share_tensor_ipc_meta(tensor):
|
|
if tensor is None:
|
|
return None
|
|
|
|
if paddle.is_compiled_with_xpu():
|
|
return tensor.value().get_tensor()._share_xpu()
|
|
|
|
if paddle.is_compiled_with_cuda() and not paddle.is_compiled_with_rocm():
|
|
return tensor.value().get_tensor()._share_cuda()
|
|
return None
|
|
|
|
|
|
def get_current_device_type():
|
|
global __current_device_type__
|
|
if __current_device_type__ is None:
|
|
if paddle.is_compiled_with_cuda():
|
|
device_type = "gpu"
|
|
elif paddle.is_compiled_with_xpu():
|
|
device_type = "xpu"
|
|
else:
|
|
current_device = _current_expected_place_()
|
|
try:
|
|
device_type = current_device.get_device_type()
|
|
except:
|
|
device_type = "unknown"
|
|
assert device_type in alignment.keys(), (
|
|
f"tensor fusion helper now only support {alignment.keys()}, but got device {device_type} instead."
|
|
)
|
|
__current_device_type__ = device_type
|
|
return __current_device_type__
|
|
|
|
|
|
def get_align(t):
|
|
size = np.prod(t.shape) * align[t.dtype]
|
|
remaining = size % alignment[get_current_device_type()]
|
|
ali = (
|
|
0
|
|
if remaining == 0
|
|
else alignment[get_current_device_type()] - remaining
|
|
)
|
|
align_ = ali // align[t.dtype]
|
|
return align_
|
|
|
|
|
|
class FusionStorage:
|
|
def __init__(
|
|
self,
|
|
accumulators,
|
|
master_weights,
|
|
merged_model_params=None,
|
|
dtype=paddle.float32,
|
|
):
|
|
assert isinstance(accumulators, dict), "accumulators must be a dict"
|
|
assert isinstance(master_weights, dict), "master_weights must be a dict"
|
|
assert (
|
|
isinstance(merged_model_params, dict) or merged_model_params is None
|
|
), "merged_model_params must be a dict or None"
|
|
self.accumulators = accumulators
|
|
self.master_weights = master_weights
|
|
self.merged_model_params = merged_model_params
|
|
self.accumulators_meta = {}
|
|
self.master_weights_meta = {}
|
|
self.merged_model_params_meta = {}
|
|
self.dtype = dtype
|
|
self.buffer = None
|
|
self.offset = 0
|
|
self.build_buffer()
|
|
self.mapping_tensor()
|
|
|
|
@imperative_base.no_grad()
|
|
def build_buffer(self):
|
|
self.offset = 0
|
|
|
|
for k, v in self.accumulators.items():
|
|
if k not in self.accumulators_meta:
|
|
self.accumulators_meta[k] = {}
|
|
for para_name, var_tmp in v.items():
|
|
assert var_tmp.dtype == self.dtype
|
|
src_len = var_tmp._numel() + get_align(var_tmp)
|
|
self.accumulators_meta[k][para_name] = {
|
|
"start": self.offset,
|
|
"end": self.offset + src_len,
|
|
"name": var_tmp.name,
|
|
"shape": var_tmp.shape,
|
|
}
|
|
self.offset += src_len
|
|
|
|
for k, v in self.master_weights.items():
|
|
assert v.dtype == self.dtype
|
|
src_len = v._numel() + get_align(v)
|
|
self.master_weights_meta[k] = {
|
|
"start": self.offset,
|
|
"end": self.offset + src_len,
|
|
"name": v.name,
|
|
"shape": v.shape,
|
|
}
|
|
self.offset += src_len
|
|
|
|
if self.merged_model_params is not None:
|
|
for k, v in self.merged_model_params.items():
|
|
assert v.dtype == self.dtype
|
|
src_len = v._numel() + get_align(v)
|
|
self.merged_model_params_meta[k] = {
|
|
"start": self.offset,
|
|
"end": self.offset + src_len,
|
|
"name": v.name,
|
|
"shape": v.shape,
|
|
}
|
|
self.offset += src_len
|
|
|
|
self.buffer = paddle.zeros((self.offset,), dtype=self.dtype)
|
|
|
|
@imperative_base.no_grad()
|
|
def mapping_tensor(self):
|
|
for k, v in self.accumulators_meta.items():
|
|
for para_name, meta in v.items():
|
|
self.mapping_tensor_impl(
|
|
src=self.accumulators[k][para_name],
|
|
start=meta["start"],
|
|
end=meta["end"],
|
|
)
|
|
|
|
for k, v in self.master_weights_meta.items():
|
|
self.mapping_tensor_impl(
|
|
src=self.master_weights[k], start=v["start"], end=v["end"]
|
|
)
|
|
|
|
for k, v in self.merged_model_params_meta.items():
|
|
self.mapping_tensor_impl(
|
|
src=self.merged_model_params[k],
|
|
start=v["start"],
|
|
end=v["end"],
|
|
)
|
|
|
|
@imperative_base.no_grad()
|
|
def mapping_tensor_impl(self, src, start, end):
|
|
tensor_shape = src.shape
|
|
stop_gradient = src.stop_gradient
|
|
src.stop_gradient = True
|
|
src.flatten_()
|
|
paddle.assign(
|
|
src,
|
|
self.buffer._slice(start, end),
|
|
)
|
|
src.get_tensor()._set_dims(tensor_shape)
|
|
src.stop_gradient = stop_gradient
|
|
self.buffer._slice(start, end)._share_buffer_to(src)
|
|
|
|
def _refresh_buffer_ipc_meta(self):
|
|
return _share_tensor_ipc_meta(self.buffer)
|
|
|
|
@property
|
|
def buffer_ipc_meta(self):
|
|
return self._refresh_buffer_ipc_meta()
|
|
|
|
|
|
class FusionStorageHelper:
|
|
def __init__(
|
|
self,
|
|
accumulators_meta,
|
|
master_weights_meta,
|
|
merged_model_params_meta,
|
|
buffer_ipc_meta,
|
|
):
|
|
self.async_loader = create_async_load()
|
|
self.accumulators_meta = None
|
|
self.master_weights_meta = None
|
|
self.merged_model_params_meta = None
|
|
self.buffer = None
|
|
self.cpu_buffer = None
|
|
self.buffer_length = None
|
|
self.tasks = []
|
|
self.reset_meta(
|
|
accumulators_meta,
|
|
master_weights_meta,
|
|
merged_model_params_meta,
|
|
buffer_ipc_meta,
|
|
)
|
|
|
|
@imperative_base.no_grad()
|
|
def reset_meta(
|
|
self,
|
|
accumulators_meta,
|
|
master_weights_meta,
|
|
merged_model_params_meta,
|
|
buffer_ipc_meta,
|
|
):
|
|
assert isinstance(accumulators_meta, dict), (
|
|
"accumulators_meta must be a dict"
|
|
)
|
|
self.accumulators_meta = accumulators_meta
|
|
assert isinstance(master_weights_meta, dict), (
|
|
"master_weights_meta must be a dict"
|
|
)
|
|
self.master_weights_meta = master_weights_meta
|
|
assert (
|
|
isinstance(merged_model_params_meta, dict)
|
|
or merged_model_params_meta is None
|
|
), "merged_model_params_meta must be a dict or None"
|
|
self.merged_model_params_meta = merged_model_params_meta
|
|
|
|
assert isinstance(buffer_ipc_meta, tuple), (
|
|
"buffer_ipc_meta must be a tuple"
|
|
)
|
|
assert len(buffer_ipc_meta) in (5, 7), (
|
|
"buffer_ipc_meta must be a tuple with length 5 when FLAGS_use_virtual_memory_auto_growth is True or 7 when FLAGS_use_virtual_memory_auto_growth is False."
|
|
)
|
|
|
|
if paddle.is_compiled_with_xpu():
|
|
new_tensor = paddle.base.core.DenseTensor._new_shared_xpu(
|
|
buffer_ipc_meta
|
|
)
|
|
else:
|
|
new_tensor = paddle.base.core.DenseTensor._new_shared_cuda(
|
|
buffer_ipc_meta
|
|
)
|
|
|
|
self.buffer = paddle.to_tensor(new_tensor)
|
|
self.cpu_buffer = self.buffer.pin_memory()
|
|
self.buffer_length = self.buffer._numel()
|
|
|
|
def sync_param(self):
|
|
self.sync_partial_param(0, self.buffer_length)
|
|
|
|
@imperative_base.no_grad()
|
|
def sync_partial_param(self, start, end):
|
|
assert isinstance(start, int), "start must be an integer"
|
|
assert isinstance(end, int), "end must be an integer"
|
|
assert start >= 0, "start must be non-negative"
|
|
assert end <= self.buffer_length, (
|
|
"end must be less than or equal to the total buffer length"
|
|
)
|
|
task = async_offload_with_offset(
|
|
src_tensor=self.buffer,
|
|
dst_tensor=self.cpu_buffer,
|
|
src_offset=start,
|
|
dst_offset=start,
|
|
offload_size=(end - start),
|
|
async_loader=self.async_loader,
|
|
)
|
|
self.tasks.append(task)
|
|
|
|
def wait_all(self):
|
|
if len(self.tasks) == 0:
|
|
return
|
|
last_task = self.tasks.pop(-1)
|
|
while len(self.tasks) > 0:
|
|
task = self.tasks.pop(0)
|
|
if paddle.is_compiled_with_xpu():
|
|
task.xpu_wait()
|
|
else:
|
|
task.cuda_wait()
|
|
last_task.cpu_wait()
|
|
|
|
def state_dict(self):
|
|
state_dict = {"master_weights": {}}
|
|
for k, v in self.accumulators_meta.items():
|
|
for para_name, tensor_meta in v.items():
|
|
var_tmp = self.restore_tensor_from_meta(tensor_meta)
|
|
state_dict[var_tmp.name] = var_tmp
|
|
for k, v in self.master_weights_meta.items():
|
|
var_tmp = self.restore_tensor_from_meta(v)
|
|
state_dict["master_weights"][k] = var_tmp
|
|
if self.merged_model_params_meta:
|
|
state_dict["merged_model_params"] = {}
|
|
for k, v in self.merged_model_params_meta.items():
|
|
var_tmp = self.restore_tensor_from_meta(v)
|
|
state_dict["merged_model_params"][k] = var_tmp
|
|
return state_dict
|
|
|
|
@imperative_base.no_grad()
|
|
def restore_tensor_from_meta(self, tensor_meta):
|
|
shape = tensor_meta["shape"]
|
|
name = tensor_meta["name"]
|
|
start = tensor_meta["start"]
|
|
end = tensor_meta["end"]
|
|
tensor = self.cpu_buffer._slice(start, end)
|
|
tensor.get_tensor()._set_dims(shape)
|
|
tensor.name = name
|
|
return tensor
|