wehub-resource-sync
812 lines
34 KiB
Python
812 lines
34 KiB
Python
# Copyright (c) 2019 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 logging
|
|
|
|
import paddle
|
|
from paddle.base import core, framework, unique_name
|
|
from paddle.base.backward import append_backward
|
|
from paddle.base.framework import Variable, in_dygraph_mode, program_guard
|
|
from paddle.optimizer import Optimizer
|
|
|
|
|
|
class RecomputeOptimizer(Optimizer):
|
|
"""
|
|
:api_attr: Static Graph
|
|
|
|
Recompute Optimizer Wrapper
|
|
|
|
Normally, a training step contains three sub-steps: first, run forward
|
|
Operators to calculate the loss; second, run backward Operators to
|
|
calculate gradient of the parameters; third, apply optimization method
|
|
to update the value of the parameters.
|
|
|
|
In the forward computation process, all variables that are needed by
|
|
backward computation process will be kept in memory, which occupy a great
|
|
amount of memory when the network becomes very deep.
|
|
|
|
Recompute split the network to k segments. In each segment, It will
|
|
recompute the forward Operators, before running backward operators. It is
|
|
very helpful for saving memory.
|
|
|
|
The Variables that separate a network to segments are called as checkpoints,
|
|
and users should set it manually. The usage is very simple:
|
|
|
|
Args:
|
|
optimizer (Optimizer): The optimizer that is applied to parameters.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> import numpy as np
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> def gen_data():
|
|
... return {
|
|
... "x": np.random.random(size=(32, 32)).astype('float32'),
|
|
... "y": np.random.randint(2, size=(32, 1)).astype('int64'),
|
|
... }
|
|
>>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
|
|
... print(input_x)
|
|
... fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
|
|
... prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
|
|
... cost = paddle.nn.functional.cross_entropy(
|
|
... input=prediction,
|
|
... label=input_y,
|
|
... reduction='none',
|
|
... use_softmax=False,
|
|
... )
|
|
... sum_cost = paddle.mean(cost)
|
|
... return sum_cost, fc_1, prediction
|
|
>>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
|
|
>>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
|
|
>>> cost, fc_1, pred = mlp(input_x, input_y)
|
|
|
|
>>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
|
|
>>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
|
|
>>> sgd._set_checkpoints([fc_1, pred])
|
|
>>> sgd.minimize(cost)
|
|
|
|
>>> print("Finished optimize")
|
|
Finished optimize
|
|
>>> place = paddle.CPUPlace()
|
|
>>> exe = paddle.static.Executor(place)
|
|
>>> exe.run(paddle.static.default_startup_program())
|
|
>>> step = 10
|
|
|
|
>>> for i in range(step):
|
|
... cost_val = exe.run(
|
|
... feed=gen_data(),
|
|
... program=paddle.static.default_main_program(),
|
|
... fetch_list=[cost.name],
|
|
... )
|
|
... print("step=%d cost=%f" % (i, cost_val[0]))
|
|
var x : DENSE_TENSOR.shape(-1, 32).dtype(float32).stop_gradient(True)
|
|
Finished optimize
|
|
step=0 cost=0.737203
|
|
step=1 cost=1.308077
|
|
step=2 cost=0.768422
|
|
step=3 cost=1.239475
|
|
step=4 cost=0.882643
|
|
step=5 cost=0.738027
|
|
step=6 cost=0.819374
|
|
step=7 cost=0.818534
|
|
step=8 cost=0.753692
|
|
step=9 cost=0.787448
|
|
|
|
"""
|
|
|
|
def __init__(self, optimizer):
|
|
if in_dygraph_mode():
|
|
raise Exception("In dygraph, don't support RecomputeOptimizer.")
|
|
self._optimizer = optimizer
|
|
self._checkpoints = None
|
|
self._learning_rate = self._optimizer._learning_rate
|
|
self._learning_rate_map = self._optimizer._learning_rate_map
|
|
self.enable_offload = False
|
|
|
|
def _set_checkpoints(self, checkpoints):
|
|
"""
|
|
Args:
|
|
checkpoints (list): List of Variable or string
|
|
"""
|
|
assert isinstance(checkpoints, list), (
|
|
"_checkpoints should be a list of Variable or a list of String"
|
|
)
|
|
for ckpt in checkpoints:
|
|
assert isinstance(ckpt, (Variable, str)), (
|
|
"_checkpoints should be a list of Variable or a list of String"
|
|
)
|
|
self._checkpoints = checkpoints
|
|
|
|
# should enable offload before calling backward
|
|
def _enable_offload(self):
|
|
self.enable_offload = True
|
|
|
|
@framework.deprecate_stat_dict
|
|
def load(self, state_dict):
|
|
"""
|
|
:api_attr: Static Graph
|
|
|
|
load function is not supported by Recompute Optimizer for now.
|
|
:return: None
|
|
|
|
Args:
|
|
state_dict: the dict load by load_persistable method
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> paddle.enable_static()
|
|
>>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
|
|
... fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
|
|
... prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
|
|
... cost = paddle.nn.functional.cross_entropy(
|
|
... input=prediction,
|
|
... label=input_y,
|
|
... reduction='none',
|
|
... use_softmax=False,
|
|
... )
|
|
... sum_cost = paddle.mean(cost)
|
|
... return sum_cost, fc_1, prediction
|
|
|
|
>>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
|
|
>>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
|
|
>>> cost, fc_1, pred = mlp(input_x, input_y)
|
|
>>> print("Finished FF")
|
|
Finished FF
|
|
|
|
>>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
|
|
>>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
|
|
>>> sgd._set_checkpoints([fc_1, pred])
|
|
>>> try:
|
|
... state_dict = {}
|
|
... sgd.load(state_dict)
|
|
>>> except NotImplementedError as e:
|
|
... print(e)
|
|
load function is not supported by Recompute Optimizer for now
|
|
"""
|
|
raise NotImplementedError(
|
|
"load function is not supported by Recompute Optimizer for now"
|
|
)
|
|
|
|
def apply_gradients(self, params_grads):
|
|
"""
|
|
call apply_gradients function of self._optimizer.
|
|
|
|
Args:
|
|
params_grads (list): list of (param, grad) pair to do optimization.
|
|
|
|
Returns:
|
|
list: A list of operators appended to the current program.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> import paddle.base.framework as framework
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
|
|
... fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
|
|
... prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
|
|
... cost = paddle.nn.functional.cross_entropy(
|
|
... input=prediction,
|
|
... label=input_y,
|
|
... reduction='none',
|
|
... use_softmax=False,
|
|
... )
|
|
... sum_cost = paddle.mean(cost)
|
|
... return sum_cost, fc_1, prediction
|
|
|
|
>>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
|
|
>>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
|
|
>>> cost, fc_1, pred = mlp(input_x, input_y)
|
|
>>> print("Finished FF")
|
|
Finished FF
|
|
|
|
>>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
|
|
>>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
|
|
>>> sgd._set_checkpoints([fc_1, pred])
|
|
>>> params_grads = sgd.backward(
|
|
... cost,
|
|
... startup_program=None,
|
|
... parameter_list=None,
|
|
... no_grad_set=None,
|
|
... )
|
|
|
|
>>> program = cost.block.program
|
|
>>> with framework.program_guard(program, None):
|
|
... optimize_ops = sgd.apply_gradients(params_grads)
|
|
|
|
>>> print("Finished apply gradients")
|
|
Finished apply gradients
|
|
"""
|
|
|
|
return self._optimizer.apply_gradients(params_grads=params_grads)
|
|
|
|
def _create_vars(self, varname):
|
|
pinned_var_name = unique_name.generate(varname + "@Pinned")
|
|
fetched_var_name = unique_name.generate(varname + "@Fetch")
|
|
|
|
pinned_var = self._main_program.global_block().create_var(
|
|
name=pinned_var_name,
|
|
shape=self.checkpoint_shape,
|
|
dtype=self._main_program.global_block().var(varname).dtype,
|
|
persistable=False,
|
|
stop_gradient=True,
|
|
)
|
|
|
|
fetch_var = self._main_program.global_block().create_var(
|
|
name=fetched_var_name,
|
|
shape=self.checkpoint_shape,
|
|
dtype=self._main_program.global_block().var(varname).dtype,
|
|
persistable=False,
|
|
stop_gradient=False,
|
|
)
|
|
|
|
return pinned_var_name, fetched_var_name
|
|
|
|
def _append_fill_constant_ops(self, startup_program):
|
|
"""
|
|
add fill_constant_ops to the end of the prog
|
|
|
|
we should fill the pinned vars before running the main_prog
|
|
to instantiate their tensor hold_, which could tell us whether
|
|
the host memory could hold all the checkpoints from all the
|
|
GPU devices in this node.
|
|
"""
|
|
op_role = 0
|
|
block = startup_program.global_block()
|
|
fill_constant_vars = self.checkpoint_name2pinned_name.values()
|
|
OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
|
|
for varname in fill_constant_vars:
|
|
var = self._main_program.global_block().var(varname)
|
|
# NOTE (JZ-LIANG) to pre-allocate the CUDAPinned MEM
|
|
pinned_var = block.create_var(
|
|
name=varname,
|
|
shape=self.checkpoint_shape,
|
|
dtype=self._main_program.global_block().var(var.name).dtype,
|
|
persistable=False,
|
|
stop_gradient=True,
|
|
)
|
|
block.append_op(
|
|
type='fill_constant',
|
|
outputs={'Out': varname},
|
|
attrs={
|
|
"shape": var.shape,
|
|
"dtype": var.dtype,
|
|
"value": 0.0,
|
|
"place_type": 2,
|
|
OP_ROLE_KEY: op_role,
|
|
},
|
|
)
|
|
|
|
def _insert_async_memcpy_op(
|
|
self, insert_idx, src_varname, dst_varname, op_role, dst_place_type
|
|
):
|
|
OP_ROLE_KEY = core.op_proto_and_checker_maker.kOpRoleAttrName()
|
|
self.block._insert_op_without_sync(
|
|
insert_idx,
|
|
type='memcpy',
|
|
inputs={'X': [self._main_program.global_block().var(src_varname)]},
|
|
outputs={
|
|
'Out': [self._main_program.global_block().var(dst_varname)]
|
|
},
|
|
attrs={"dst_place_type": int(dst_place_type), OP_ROLE_KEY: op_role},
|
|
)
|
|
|
|
def _insert_fetch_op(self, idx, varname):
|
|
assert varname in self.checkpoint_name2pinned_name, (
|
|
f"Try to fetch {varname} from Pinned Memory, but it is NOT a checkpoint"
|
|
)
|
|
|
|
pinned_varname = self.checkpoint_name2pinned_name[varname]
|
|
fetch_varname = self.checkpoint_name2fetch_name[varname]
|
|
self._insert_async_memcpy_op(idx, pinned_varname, fetch_varname, 1, 1)
|
|
|
|
def _insert_offload_op(self, idx, varname):
|
|
assert varname in self.checkpoint_name2pinned_name, (
|
|
f"Try to offload {varname} to Pinned Memory, but it is NOT a checkpoint"
|
|
)
|
|
pinned_varname = self.checkpoint_name2pinned_name[varname]
|
|
self._insert_async_memcpy_op(idx, varname, pinned_varname, 0, 2)
|
|
|
|
def _insert_sync_op(self, op_idx, checkpoint_name):
|
|
# single stream offload no need sync
|
|
pass
|
|
|
|
def _record_fetch_op(self, idx):
|
|
assert len(self.un_fetch_checkpoint_names) > 0, (
|
|
"Could NOT found checkpoint to fetch"
|
|
)
|
|
checkpoint_name = self.un_fetch_checkpoint_names.pop(-1)
|
|
logging.debug(f"Record fetch [{checkpoint_name}]")
|
|
self.idx2insertions[idx] = ("fetch", checkpoint_name)
|
|
|
|
return checkpoint_name
|
|
|
|
def _record_offload_op(self, idx, checkpoint_name):
|
|
expected_checkpoint_name = self.un_offload_checkpoint_names.pop(0)
|
|
assert checkpoint_name == expected_checkpoint_name, (
|
|
f"expected to offload [{expected_checkpoint_name}] but got [{checkpoint_name}]"
|
|
)
|
|
logging.debug(f"Record offload [{checkpoint_name}]")
|
|
self.idx2insertions[idx] = ("offload", checkpoint_name)
|
|
|
|
def _record_sync_op(self, idx, checkpoint_name):
|
|
assert checkpoint_name not in self.synced_checkpoints, (
|
|
f"Try to sync the checkpoint [{checkpoint_name}] twice"
|
|
)
|
|
self.synced_checkpoints.add(checkpoint_name)
|
|
logging.debug(f"Record offload sync [{checkpoint_name}]")
|
|
self.idx2insertions[idx] = ("sync", checkpoint_name)
|
|
|
|
def _parse_backward(self):
|
|
self.idx2insertions = {}
|
|
# don't offload the last checkpoints, to favor throughput
|
|
self.un_fetch_checkpoint_names = self.sorted_checkpoint_names[:]
|
|
self.un_fetch_checkpoint_names.pop(-1)
|
|
need_fetch_checkpoint_names = self.un_fetch_checkpoint_names[:]
|
|
self.checkpoint_usage_count = {}
|
|
for checkpoint_name in self.un_fetch_checkpoint_names:
|
|
self.checkpoint_usage_count[checkpoint_name] = 0
|
|
|
|
self.bw_start_op_idx = len(self.block.ops)
|
|
for idx, op in enumerate(self.block.ops):
|
|
if int(op.desc.attr("op_role")) == 1:
|
|
self.bw_start_op_idx = idx
|
|
break
|
|
|
|
assert self.bw_start_op_idx < len(self.block.ops), (
|
|
"Could NOT found backward op in prog"
|
|
)
|
|
|
|
# fetch second to last checkpoint at the beginning of BW
|
|
fetched_checkpoint_varname = self._record_fetch_op(self.bw_start_op_idx)
|
|
last_last_fetch_checkpoint = None
|
|
|
|
for i, op in enumerate(self.block.ops[self.bw_start_op_idx :]):
|
|
idx = self.bw_start_op_idx + i
|
|
input_vars = op.desc.input_arg_names()
|
|
|
|
for input_var in input_vars:
|
|
if input_var in need_fetch_checkpoint_names:
|
|
if input_var not in self.un_fetch_checkpoint_names:
|
|
# fetch the offload checkpoint when the first usage of its previous one
|
|
if self.checkpoint_usage_count[input_var] == 0:
|
|
# TODO (JZ-LIANG) sync memcpy_stream if extra stream for memcpy
|
|
second_to_last_fetch_checkpoint = (
|
|
fetched_checkpoint_varname
|
|
)
|
|
# there is NO fetch ahead the first checkpoint
|
|
if input_var != self.sorted_checkpoint_names[0]:
|
|
fetched_checkpoint_varname = (
|
|
self._record_fetch_op(idx)
|
|
)
|
|
|
|
# should check the current used checkpoint is the last fetch one
|
|
assert second_to_last_fetch_checkpoint == input_var, (
|
|
f"Current recompute segment should use [{second_to_last_fetch_checkpoint}] BUT got [{input_var}]"
|
|
)
|
|
# rename
|
|
self.block.ops[idx]._rename_input(
|
|
input_var,
|
|
self.checkpoint_name2fetch_name[input_var],
|
|
)
|
|
self.checkpoint_usage_count[input_var] += 1
|
|
else:
|
|
raise ValueError(
|
|
f"use checkpoint [{input_var}] before fetch in BW"
|
|
)
|
|
|
|
assert len(self.un_fetch_checkpoint_names) == 0, (
|
|
f"{self.un_fetch_checkpoint_names} checkpoints have NOT been Recorded"
|
|
)
|
|
|
|
def _update_backward(self):
|
|
if len(self.idx2insertions) == 0:
|
|
return
|
|
total_op = len(self.block.ops)
|
|
for op_idx in reversed(range(self.bw_start_op_idx, total_op)):
|
|
if op_idx in self.idx2insertions:
|
|
operation, checkpoint_name = self.idx2insertions[op_idx]
|
|
if operation == "fetch":
|
|
self._insert_fetch_op(op_idx, checkpoint_name)
|
|
logging.debug(f"Insert [{checkpoint_name}] fetch op.")
|
|
del self.idx2insertions[op_idx]
|
|
elif operation == "sync":
|
|
self._insert_sync_op(op_idx, checkpoint_name)
|
|
logging.debug(f"Sync [{checkpoint_name}] fetch op.")
|
|
self.block._sync_with_cpp()
|
|
assert len(self.idx2insertions) == 0, (
|
|
f"{[ele[1] for ele in self.idx2insertions.values()]} checkpoints left un-Fetched"
|
|
)
|
|
|
|
def _parse_forward(self):
|
|
self.idx2insertions = {}
|
|
# don't offload the last checkpoints, faster, less memory saving
|
|
self.un_offload_checkpoint_names = self.sorted_checkpoint_names[:]
|
|
last_checkpoint = self.un_offload_checkpoint_names.pop(-1)
|
|
need_offload_checkpoint_names = self.un_offload_checkpoint_names[:]
|
|
self.checkpoint_usage_count_and_idx = {}
|
|
for checkpoint_name in self.un_offload_checkpoint_names:
|
|
self.checkpoint_usage_count_and_idx[checkpoint_name] = {
|
|
'count': 0,
|
|
'idx': -1,
|
|
}
|
|
self.synced_checkpoints = set()
|
|
self.fw_start_op_idx = len(self.block.ops)
|
|
for idx, op in enumerate(self.block.ops):
|
|
if int(op.desc.attr("op_role")) == 0:
|
|
self.fw_start_op_idx = idx
|
|
break
|
|
|
|
assert self.fw_start_op_idx < len(self.block.ops), (
|
|
"Could NOT found Forward op in prog"
|
|
)
|
|
last_offload_checkpoint = None
|
|
|
|
for i, op in enumerate(
|
|
self.block.ops[self.fw_start_op_idx : self.bw_start_op_idx]
|
|
):
|
|
idx = self.fw_start_op_idx + i
|
|
output_vars = op.desc.output_arg_names()
|
|
input_vars = op.desc.input_arg_names()
|
|
|
|
for output_var in output_vars:
|
|
if output_var in need_offload_checkpoint_names:
|
|
assert len(output_vars) == 1, (
|
|
f"checkpoint should be the only Output of a certain op, but [{output_var}] is from [{op}]"
|
|
)
|
|
|
|
if output_var in self.un_offload_checkpoint_names:
|
|
# insert sync op if last checkpoint has not been sync
|
|
if last_offload_checkpoint is not None:
|
|
if (
|
|
self.checkpoint_usage_count_and_idx[
|
|
last_offload_checkpoint
|
|
]['count']
|
|
== 0
|
|
):
|
|
self._record_sync_op(
|
|
idx, last_offload_checkpoint
|
|
)
|
|
else:
|
|
last_usage_idx = (
|
|
self.checkpoint_usage_count_and_idx[
|
|
last_offload_checkpoint
|
|
]['idx']
|
|
)
|
|
assert last_usage_idx > 0, (
|
|
f"last_usage_idx of checkpoint [{last_offload_checkpoint}] should large than 0"
|
|
)
|
|
self._record_sync_op(
|
|
last_usage_idx + 1, last_offload_checkpoint
|
|
)
|
|
# insert offload op after the checkpoint's generation op
|
|
self._record_offload_op(idx + 1, output_var)
|
|
last_offload_checkpoint = output_var
|
|
else:
|
|
raise ValueError(
|
|
f"There should be just ONE op that output checkpoint [{output_var}]"
|
|
)
|
|
# need to sync the last need to offload checkpoint before the last checkpoint as output op
|
|
if output_var == last_checkpoint:
|
|
assert len(output_vars) == 1, (
|
|
f"checkpoint should be the only Output of a certain op, but [{output_var}] is from [{op}]"
|
|
)
|
|
assert (
|
|
last_offload_checkpoint
|
|
== self.sorted_checkpoint_names[-2]
|
|
), (
|
|
f"the last offload checkpoint before [{last_checkpoint}] is suppose to be [{self.sorted_checkpoint_names[-2]}], but got [{last_offload_checkpoint}]"
|
|
)
|
|
# sync if last checkpoint has not been sync
|
|
if (
|
|
self.checkpoint_usage_count_and_idx[
|
|
last_offload_checkpoint
|
|
]['idx']
|
|
== 0
|
|
):
|
|
self._record_sync_op(idx, last_offload_checkpoint)
|
|
else:
|
|
last_usage_idx = self.checkpoint_usage_count_and_idx[
|
|
last_offload_checkpoint
|
|
]['idx']
|
|
assert last_usage_idx > 0, (
|
|
f"last_usage_idx of checkpoint [{last_offload_checkpoint}] should large than 0"
|
|
)
|
|
self._record_sync_op(
|
|
last_usage_idx + 1, last_offload_checkpoint
|
|
)
|
|
# record checkpoint usage
|
|
for input_var in input_vars:
|
|
if input_var in need_offload_checkpoint_names:
|
|
assert input_var not in self.synced_checkpoints, (
|
|
f"checkpoint [{input_var}] used after sync"
|
|
)
|
|
self.checkpoint_usage_count_and_idx[input_var]['count'] += 1
|
|
self.checkpoint_usage_count_and_idx[input_var]['idx'] = idx
|
|
|
|
assert len(self.un_offload_checkpoint_names) == 0, (
|
|
f"{self.un_fetch_checkpoint_names} checkpoints have NOT been Recorded"
|
|
)
|
|
assert len(self.synced_checkpoints) == len(
|
|
need_offload_checkpoint_names
|
|
), (
|
|
f"{set(need_offload_checkpoint_names) - set(self.synced_checkpoints)} checkpoints have NOT been Recorded"
|
|
)
|
|
|
|
def _update_forward(self):
|
|
if len(self.idx2insertions) == 0:
|
|
return
|
|
for op_idx in reversed(
|
|
range(self.fw_start_op_idx, self.bw_start_op_idx)
|
|
):
|
|
if op_idx in self.idx2insertions:
|
|
operation, checkpoint_name = self.idx2insertions[op_idx]
|
|
if operation == "offload":
|
|
self._insert_offload_op(op_idx, checkpoint_name)
|
|
logging.debug(f"Insert [{checkpoint_name}] offload op.")
|
|
del self.idx2insertions[op_idx]
|
|
elif operation == "sync":
|
|
self._insert_sync_op(op_idx, checkpoint_name)
|
|
logging.debug(
|
|
f"Insert [{checkpoint_name}] offload_sync op."
|
|
)
|
|
del self.idx2insertions[op_idx]
|
|
|
|
self.block._sync_with_cpp()
|
|
assert len(self.idx2insertions) == 0, (
|
|
f"{[ele[1] for ele in self.idx2insertions.values()]} checkpoints left un-Offloaded"
|
|
)
|
|
|
|
def _check_offload_fetch(self):
|
|
# TODO(JZ-LIANG) the single stream offload need no sync
|
|
pass
|
|
|
|
def _offload(self, loss, startup_program=None):
|
|
"""
|
|
core steps for recompute offload
|
|
1. create pinned vars and temp vars
|
|
2. parse & update Forward pass: offload, sync
|
|
3. parse & update Backward pass: rename, fetch, sync
|
|
4. verify the correctness
|
|
"""
|
|
self._main_program = loss.block.program
|
|
self.block = loss.block
|
|
if startup_program is None:
|
|
startup_program = paddle.static.default_startup_program()
|
|
|
|
with program_guard(self._main_program, startup_program):
|
|
assert len(self.checkpoint_shape) > 0, (
|
|
f"checkpoints shape {self.checkpoint_shape} should be an non empty list like: [12, 512, 1024]"
|
|
)
|
|
assert all(ele > 0 for ele in self.checkpoint_shape), (
|
|
f"all ele in checkpoints shape {self.checkpoint_shape} should be a determined integer larger than 0"
|
|
)
|
|
self.checkpoint_name2pinned_name = {}
|
|
self.checkpoint_name2fetch_name = {}
|
|
for checkpoint_varname in self.sorted_checkpoint_names:
|
|
pinned_var_name, fetch_var_name = self._create_vars(
|
|
checkpoint_varname
|
|
)
|
|
self.checkpoint_name2pinned_name[checkpoint_varname] = (
|
|
pinned_var_name
|
|
)
|
|
self.checkpoint_name2fetch_name[checkpoint_varname] = (
|
|
fetch_var_name
|
|
)
|
|
self._append_fill_constant_ops(startup_program)
|
|
# TODO (JZ-LIANG) to provide two offload strategy in future
|
|
# step 2. parse & update FW: rename, offload, sync
|
|
self._parse_backward()
|
|
self._update_backward()
|
|
# step 3. parse & update BW: rename, offload, sync
|
|
self._parse_forward()
|
|
self._update_forward()
|
|
# step 4. verify the correctness
|
|
self._check_offload_fetch()
|
|
|
|
def backward(
|
|
self,
|
|
loss,
|
|
startup_program=None,
|
|
parameter_list=None,
|
|
no_grad_set=None,
|
|
callbacks=None,
|
|
):
|
|
"""
|
|
call append_backward with checkpoints.
|
|
|
|
Args:
|
|
loss (Variable): loss variable to run optimizations.
|
|
startup_program (Program): startup_program for initializing parameters
|
|
in `parameter_list`.
|
|
parameter_list (list): list of Variables or Variable.names to update.
|
|
no_grad_set (set|None): set of Variables or Variables.names should be ignored.
|
|
callbacks (list|None): list of callables to run when appending backward
|
|
operator for one parameter.
|
|
checkpoints (list): list of Variables as checkpoints
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
|
|
... fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
|
|
... prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
|
|
... cost = paddle.nn.functional.cross_entropy(
|
|
... input=prediction,
|
|
... label=input_y,
|
|
... reduction='none',
|
|
... use_softmax=False,
|
|
... )
|
|
... sum_cost = paddle.mean(cost)
|
|
... return sum_cost, fc_1, prediction
|
|
|
|
>>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
|
|
>>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
|
|
>>> cost, fc_1, pred = mlp(input_x, input_y)
|
|
>>> print("Finished FF")
|
|
Finished FF
|
|
|
|
>>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
|
|
>>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
|
|
>>> sgd._set_checkpoints([fc_1, pred])
|
|
>>> params_grads = sgd.backward(
|
|
... cost,
|
|
... startup_program=None,
|
|
... parameter_list=None,
|
|
... no_grad_set=None,
|
|
... )
|
|
>>> print("Finished backward")
|
|
Finished backward
|
|
"""
|
|
assert self._checkpoints is not None, (
|
|
"You should call _set_checkpoints first"
|
|
)
|
|
|
|
if in_dygraph_mode():
|
|
raise NotImplementedError(
|
|
"DyGraph current does not support recompute"
|
|
)
|
|
|
|
self._dtype = loss.dtype
|
|
program = loss.block.program
|
|
with program_guard(program, startup_program):
|
|
checkpoint_vars = []
|
|
for ckpt in self._checkpoints:
|
|
if isinstance(ckpt, Variable):
|
|
checkpoint_vars.append(ckpt)
|
|
else:
|
|
checkpoint_vars.append(loss.block.var(ckpt))
|
|
|
|
# allow return to non-recompute when checkpoints is empty
|
|
if len(checkpoint_vars) > 0:
|
|
params_grads, sorted_checkpoint_names = append_backward(
|
|
loss,
|
|
parameter_list,
|
|
no_grad_set,
|
|
checkpoints=checkpoint_vars,
|
|
)
|
|
else:
|
|
params_grads = append_backward(
|
|
loss,
|
|
parameter_list,
|
|
no_grad_set,
|
|
checkpoints=checkpoint_vars,
|
|
)
|
|
|
|
if self.enable_offload:
|
|
self.sorted_checkpoint_names = sorted_checkpoint_names
|
|
self._offload(loss, startup_program=startup_program)
|
|
|
|
return params_grads
|
|
|
|
def apply_optimize(self, loss, startup_program, params_grads):
|
|
"""
|
|
call the apply_optimize function of self._optimizer
|
|
Args:
|
|
loss (Variable): loss variable to run optimizations.
|
|
startup_program (Program): startup_program for initializing parameters
|
|
in `parameter_list`.
|
|
params_grads (list): list of (param, grad) pair to do optimization.
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> paddle.enable_static()
|
|
|
|
>>> def mlp(input_x, input_y, hid_dim=128, label_dim=2):
|
|
... fc_1 = paddle.static.nn.fc(x=input_x, size=hid_dim)
|
|
... prediction = paddle.static.nn.fc(x=[fc_1], size=label_dim, activation='softmax')
|
|
... cost = paddle.nn.functional.cross_entropy(
|
|
... input=prediction,
|
|
... label=input_y,
|
|
... reduction='none',
|
|
... use_softmax=False,
|
|
... )
|
|
... sum_cost = paddle.mean(cost)
|
|
... return sum_cost, fc_1, prediction
|
|
|
|
>>> input_x = paddle.static.data(name="x", shape=[-1, 32], dtype='float32')
|
|
>>> input_y = paddle.static.data(name="y", shape=[-1, 1], dtype='int64')
|
|
>>> cost, fc_1, pred = mlp(input_x, input_y)
|
|
>>> print("Finished FF")
|
|
Finished FF
|
|
|
|
>>> sgd = paddle.optimizer.Adam(learning_rate=0.01)
|
|
>>> sgd = paddle.incubate.optimizer.RecomputeOptimizer(sgd)
|
|
>>> sgd._set_checkpoints([fc_1, pred])
|
|
>>> params_grads = sgd.backward(
|
|
... cost,
|
|
... startup_program=None,
|
|
... parameter_list=None,
|
|
... no_grad_set=None,
|
|
... )
|
|
|
|
>>> optimize_ops = sgd.apply_optimize(
|
|
... cost,
|
|
... startup_program=None,
|
|
... params_grads=params_grads,
|
|
... )
|
|
|
|
>>> print("Finished apply_optimize")
|
|
Finished apply_optimize
|
|
"""
|
|
|
|
func = (
|
|
self._optimizer.apply_optimize
|
|
if hasattr(self._optimizer, 'apply_optimize')
|
|
else self._optimizer._apply_optimize
|
|
)
|
|
return func(
|
|
loss, startup_program=startup_program, params_grads=params_grads
|
|
)
|
|
|
|
def minimize(
|
|
self, loss, startup_program=None, parameter_list=None, no_grad_set=None
|
|
):
|
|
assert isinstance(loss, Variable), "The loss should be an Variable."
|
|
assert self._checkpoints is not None, (
|
|
"You should call _set_checkpoints first"
|
|
)
|
|
if in_dygraph_mode():
|
|
raise NotImplementedError(
|
|
"DyGraph current does not support recompute"
|
|
)
|
|
params_grads = self.backward(
|
|
loss,
|
|
startup_program=startup_program,
|
|
parameter_list=parameter_list,
|
|
no_grad_set=no_grad_set,
|
|
)
|
|
|
|
optimize_ops = self.apply_optimize(
|
|
loss, startup_program=startup_program, params_grads=params_grads
|
|
)
|
|
|
|
return optimize_ops, params_grads
|