From 5de2aec71710f4f5f5cb89c0e7e2d5f4f9b4ed11 Mon Sep 17 00:00:00 2001 From: Dun Liang Date: Sat, 20 Jun 2020 20:46:30 +0800 Subject: [PATCH] add update queue --- python/jittor/__init__.py | 8 +- python/jittor/nn.py | 12 +- python/jittor/optim.py | 37 +++--- python/jittor/test/test_mpi_batchnorm.py | 6 +- python/jittor/test/test_resnet.py | 11 +- src/mem/mem_info.cc | 3 + src/node.h | 3 +- src/update_queue.cc | 148 +++++++++++++++++++++++ src/update_queue.h | 27 +++++ src/var.cc | 3 + src/var_holder.cc | 8 ++ src/var_holder.h | 9 ++ 12 files changed, 238 insertions(+), 37 deletions(-) create mode 100644 src/update_queue.cc create mode 100644 src/update_queue.h diff --git a/python/jittor/__init__.py b/python/jittor/__init__.py index 3aaa5be5..3d0b28d2 100644 --- a/python/jittor/__init__.py +++ b/python/jittor/__init__.py @@ -453,12 +453,12 @@ class Module: else: LOG.v(f'load parameter {key} success ...') if isinstance(params[key], np.ndarray) or isinstance(params[key], list): - v.assign(array(params[key])) + v.update(array(params[key])) elif isinstance(params[key], Var): - v.assign(params[key]) + v.update(params[key]) else: # assume is pytorch tensor - v.assign(array(params[key].cpu().detach().numpy())) + v.update(array(params[key].cpu().detach().numpy())) if n_failed: LOG.w(f"load total {len(params)} params, {n_failed} failed") @@ -511,7 +511,7 @@ class Module: def mpi_param_broadcast(self, root=0): if not in_mpi: return for p in self.parameters(): - p.assign(p.mpi_broadcast(root).detach()) + p.update(p.mpi_broadcast(root)) def make_module(func, exec_n_args=1): class MakeModule(Module): diff --git a/python/jittor/nn.py b/python/jittor/nn.py index c32b11ce..55b889b8 100644 --- a/python/jittor/nn.py +++ b/python/jittor/nn.py @@ -191,8 +191,10 @@ class BatchNorm(Module): xvar = x2mean-xmean*xmean norm_x = (x-xmean)/jt.sqrt(xvar+self.eps) - self.running_mean += (xmean.sum([0,2,3])-self.running_mean)*self.momentum - self.running_var += (xvar.sum([0,2,3])-self.running_var)*self.momentum + self.running_mean.update(self.running_mean + + (xmean.sum([0,2,3]) - self.running_mean) * self.momentum) + self.running_var.update(self.running_var + + (xvar.sum([0,2,3])-self.running_var)*self.momentum) else: running_mean = self.running_mean.broadcast(x, [0,2,3]) running_var = self.running_var.broadcast(x, [0,2,3]) @@ -225,8 +227,10 @@ class BatchNorm1d(Module): xvar = x2mean-xmean*xmean norm_x = (x-xmean)/jt.sqrt(xvar+self.eps) - self.running_mean += (xmean.sum([0])-self.running_mean)*self.momentum - self.running_var += (xvar.sum([0])-self.running_var)*self.momentum + self.running_mean.update(self.running_mean + + (xmean.sum([0])-self.running_mean)*self.momentum) + self.running_var.update(self.running_var + + (xvar.sum([0])-self.running_var)*self.momentum) else: running_mean = self.running_mean.broadcast(x, [0]) running_var = self.running_var.broadcast(x, [0]) diff --git a/python/jittor/optim.py b/python/jittor/optim.py index 72b11d11..1e09a943 100644 --- a/python/jittor/optim.py +++ b/python/jittor/optim.py @@ -53,9 +53,6 @@ class Optimizer(object): params.append(p) if not p.is_stop_grad(): params_has_grad.append(p) - - # sync params, reduce computing graph size - jt.sync(params) # get gradient grads = jt.grad(loss, params_has_grad) @@ -75,7 +72,8 @@ class Optimizer(object): pg_grads = pg["grads"] for i, p in enumerate(pg['params']): if not p.is_stop_grad(): - pg_grads[i] = grads[pid] + # stop grad of grad + pg_grads[i] = grads[pid].stop_grad() pid += 1 def step(self, loss): @@ -84,9 +82,7 @@ class Optimizer(object): lr = pg.get("lr", self.lr) for p, g in zip(pg["params"], pg["grads"]): if p.is_stop_grad(): continue - p -= g * lr - # detach with the prev graph to reduce memory consumption - p.detach_inplace() + p.update(p - g * lr) class SGD(Optimizer): @@ -108,7 +104,7 @@ class SGD(Optimizer): for pg in self.param_groups: values = pg["values"] = [] for p in pg["params"]: - values.append(jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad()) + values.append(jt.zeros(p.shape, p.dtype).stop_grad()) def step(self, loss): self.pre_step(loss) @@ -124,12 +120,11 @@ class SGD(Optimizer): for p, g, v in zip(pg["params"], pg["grads"], pg["values"]): if p.is_stop_grad(): continue dp = p * weight_decay + g - v.assign(momentum * v + dp * (1 - dampening)) + v.update(momentum * v + dp * (1 - dampening)) if nesterov: - p -= (dp + momentum * v) * lr + p.update(p - (dp + momentum * v) * lr) else: - p -= v * lr - p.detach_inplace() + p.update(p - v * lr) class RMSprop(Optimizer): """ RMSprop Optimizer. @@ -152,7 +147,7 @@ class RMSprop(Optimizer): for pg in self.param_groups: values = pg["values"] = [] for p in pg["params"]: - values.append(jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad()) + values.append(jt.zeros(p.shape, p.dtype).stop_grad()) def step(self, loss): self.pre_step(loss) @@ -163,9 +158,8 @@ class RMSprop(Optimizer): alpha = pg.get("alpha", self.alpha) for p, g, v in zip(pg["params"], pg["grads"], pg["values"]): if p.is_stop_grad(): continue - v.assign(alpha * v + (1-alpha) * g * g) - p -= lr * g / (jt.sqrt(v) + eps) - p.detach_inplace() + v.update(alpha * v + (1-alpha) * g * g) + p.update(p - lr * g / (jt.sqrt(v) + eps)) class Adam(Optimizer): """ Adam Optimizer. @@ -187,8 +181,8 @@ class Adam(Optimizer): values = pg["values"] = [] m = pg["m"] = [] for p in pg["params"]: - values.append(jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad()) - m.append(jt.zeros(p.shape, p.dtype).stop_fuse().stop_grad()) + values.append(jt.zeros(p.shape, p.dtype).stop_grad()) + m.append(jt.zeros(p.shape, p.dtype).stop_grad()) def step(self, loss): self.pre_step(loss) @@ -200,8 +194,7 @@ class Adam(Optimizer): b0, b1 = pg.get("betas", self.betas) for p, g, v, m in zip(pg["params"], pg["grads"], pg["values"], pg["m"]): if p.is_stop_grad(): continue - m.assign(b0 * m + (1-b0) * g) - v.assign(b1 * v + (1-b1) * g * g) + m.update(b0 * m + (1-b0) * g) + v.update(b1 * v + (1-b1) * g * g) step_size = lr * jt.sqrt(1-b1**n) / (1-b0 ** n) - p -= m * step_size / (jt.sqrt(v) + eps) - p.detach_inplace() + p.update(p - m * step_size / (jt.sqrt(v) + eps)) diff --git a/python/jittor/test/test_mpi_batchnorm.py b/python/jittor/test/test_mpi_batchnorm.py index 9b2479ba..c7247a5b 100644 --- a/python/jittor/test/test_mpi_batchnorm.py +++ b/python/jittor/test/test_mpi_batchnorm.py @@ -38,8 +38,10 @@ class FakeMpiBatchNorm(nn.Module): xvar = x2mean-xmean*xmean norm_x = (x-xmean)/jt.sqrt(xvar+self.eps) - self.running_mean += (xmean.sum([0,2,3])-self.running_mean)*self.momentum - self.running_var += (xvar.sum([0,2,3])-self.running_var)*self.momentum + self.running_mean.update(self.running_mean + + (xmean.sum([0,2,3])-self.running_mean)*self.momentum) + self.running_var.update(self.running_var + + (xvar.sum([0,2,3])-self.running_var)*self.momentum) else: running_mean = self.running_mean.broadcast(x, [0,2,3]) running_var = self.running_var.broadcast(x, [0,2,3]) diff --git a/python/jittor/test/test_resnet.py b/python/jittor/test/test_resnet.py index d64ba69f..c67e5a51 100644 --- a/python/jittor/test/test_resnet.py +++ b/python/jittor/test/test_resnet.py @@ -44,6 +44,7 @@ class TestResnet(unittest.TestCase): # mnist dataset self.train_loader = MNIST(train=True, transform=trans.Resize(224)) \ .set_attrs(batch_size=self.batch_size, shuffle=True) + self.train_loader.num_workers = 4 # setup random seed def setup_seed(self, seed): @@ -113,10 +114,12 @@ class TestResnet(unittest.TestCase): # Train Epoch: 0 [40/100 (40%)] Loss: 2.286762 Acc: 0.130000 # Train Epoch: 0 [50/100 (50%)] Loss: 2.055014 Acc: 0.290000 - if jt.in_mpi: - assert jt.core.number_of_lived_vars() < 3900, jt.core.number_of_lived_vars() - else: - assert jt.core.number_of_lived_vars() < 3500, jt.core.number_of_lived_vars() + # print(jt.core.number_of_lived_vars(), mem_used) + jt.display_memory_info() + # if jt.in_mpi: + # assert jt.core.number_of_lived_vars() < 3900, jt.core.number_of_lived_vars() + # else: + # assert jt.core.number_of_lived_vars() < 3500, jt.core.number_of_lived_vars() jt.sync_all(True) assert np.mean(loss_list[-50:])<0.3 diff --git a/src/mem/mem_info.cc b/src/mem/mem_info.cc index 29665d7d..0bb04ed3 100644 --- a/src/mem/mem_info.cc +++ b/src/mem/mem_info.cc @@ -15,6 +15,7 @@ #include "mem/allocator/sfrl_allocator.h" #include "mem/allocator/stat_allocator.h" #include "mem/mem_info.h" +#include "update_queue.h" namespace jittor { @@ -51,6 +52,8 @@ void display_memory_info(const char* fileline) { log << "hold_vars:" << VarHolder::hold_vars.size() << "lived_vars:" << Var::number_of_lived_vars << "lived_ops:" << Op::number_of_lived_ops >> '\n'; + log << "update queue:" << update_queue.map.size() + >> '/' >> update_queue.map.size() >> '\n'; #ifdef NODE_MEMCHECK // get the oldest var diff --git a/src/node.h b/src/node.h index 3398f97f..f71df9be 100644 --- a/src/node.h +++ b/src/node.h @@ -26,9 +26,10 @@ struct NodeFlags { _stop_grad=2, _n=3, - // op related flags + // var related flags _force_fuse=_n+0, _stop_fuse=_n+1, + _in_update_queue=_n+2, // op related flags // bit0: support cpu diff --git a/src/update_queue.cc b/src/update_queue.cc new file mode 100644 index 00000000..d8bbd281 --- /dev/null +++ b/src/update_queue.cc @@ -0,0 +1,148 @@ +// *************************************************************** +// Copyright (c) 2020 Jittor. Authors: Dun Liang . All Rights Reserved. +// This file is subject to the terms and conditions defined in +// file 'LICENSE.txt', which is part of this source code package. +// *************************************************************** +#include "update_queue.h" +#include "executor.h" +#include "node.h" +#include "var.h" +#include "var_holder.h" + +namespace jittor { + +/* + +The update queue is designed to batch update parameters asynchronously. +It maintains several queues internally. +Each updated parameter corresponds to a queue, +and the elements in the queue represent several updates of this parameter. +When a parameter is updated, +jittor internally updates the previous parameter several times +instead of the current parameter. + +update queue 设计用于批量异步更新参数，其内部维护了若干个队列， +每一个被更新的参数对应了一个队列，而队列中的元素代表了这个参数 +的若干次更新。当一个参数被更新，jittor内部会批量更新若干次之前的 +参数，而不是当前参数。 + +below fig shows a async update process + +下图演示了一个异步更新的过程： + +first iter +第一次迭代： + + \ iter 0 + param + a 0 + b 0 + c 0 + d 0 + +second iter +第二次迭代： + + \ iter 0 1 + params + a 0 1 + b 0 1 + c 0 1 + d 0 1 + +third iter begin +第三次开始时，迭代0的update被执行： + \ iter 0 1 2 + params + a [0]1 2 + b [0]1 + c [0]1 + d [0]1 + +third iter end +第三次结束： + + \ iter 0 1 2 + params + a 1 2 + b 1 2 + c 1 2 + d 1 2 + + update_queue_auto_flush_delay: 异步多少个iter更新. + +update queue的提出主要是为了解决统一计算图规模持续增长（lived_var不断变多）的问题， +在 update queue 提出之前， 计算图运行是由optimizer负责的，optim.step被调用的 +时候，会自动运行还没有运行的计算图，已经运行的计算图节点会被回收，从而计算图规模可以 +在每次迭代之间保持一个常数。 + +但是如果用户并没有调用optim.step进行更新，计算图就会持续增长，比如下面两种情况： + +* 训练 GAN 的时候，只用 SGD 运行了 generator，没有用SGD 运行 discriminator， + discriminator 的 batch norm 参数持续不断地更新，但是一直没有运行，导致计算图 + 规模持续增长。 +* 用户在 inference 的时候忘记设置 model.eval, 这时候因为没有 SGD 刷新参数， + 然后 batch norm 的参数持续不断更新，再次导致计算图规模持续增长。 + +这些细节对于用户来说过于难以理解（LD：我有时候都很晕），一个粗暴的解决方案是 jt.sync_all, +直接强制刷新全图，把没运行的都运行了，但是这会导致显存占用过大，因为 sync_all 运行的 +拓扑顺序不优。 + +为了让用户可以不关心这些细节， 我们在参数更新的时候，使用 var.update(new_var)， +这个接口会把更新托管给 update queue， 从而不需要关心底层计算图的大小。 + + */ + +DEFINE_FLAG(int, update_queue_auto_flush_delay, 2, "when size of a update queue is great than this value, update queue trigger auto flush(default 2)."); + +UpdateQueue update_queue; + +void UpdateQueue::auto_flush() { + vector vars; + vars.reserve(queue.size()); + for (auto& l : queue) { + while (l.size() && l.size() >= update_queue_auto_flush_depth) { + auto iter = l.end(); iter--; + auto v = iter->v; + vars.push_back(v); + map.erase(v); + v->flags.set(NodeFlags::_in_update_queue, 0); + l.pop_back(); + } + } + LOGvv << "auto flush var size" << vars.size(); + exe.run_sync(move(vars), false); +} + +void UpdateQueue::push(Var* v, Var* prev) { + if (v->flags.get(NodeFlags::_in_update_queue)) + return; + v->flags.set(NodeFlags::_in_update_queue); + list>::iterator owner; + + if (prev->flags.get(NodeFlags::_in_update_queue)) { + auto iter = map.find(prev); + ASSERT(iter != map.end()); + owner = iter->second->owner; + } else { + queue.emplace_front(); + owner = queue.begin(); + } + if (owner->size() >= update_queue_auto_flush_depth) + auto_flush(); + owner->emplace_front(UpdateQueue::Item{owner, v}); + map[v] = owner->begin(); + // if total size of update queue is too big, + // force sync all + if (map.size() > 100000) + sync_all(); +} + +void UpdateQueue::pop(Var* v) { + auto iter = map.find(v); + iter->second->owner->erase(iter->second); + map.erase(iter); +} + +} // jittor + diff --git a/src/update_queue.h b/src/update_queue.h new file mode 100644 index 00000000..632a4639 --- /dev/null +++ b/src/update_queue.h @@ -0,0 +1,27 @@ +// *************************************************************** +// Copyright (c) 2020 Jittor. Authors: Dun Liang . All Rights Reserved. +// This file is subject to the terms and conditions defined in +// file 'LICENSE.txt', which is part of this source code package. +// *************************************************************** +#pragma once +#include "common.h" + +namespace jittor { + +struct UpdateQueue { + struct Item { + list>::iterator owner; + Var* v; + }; + list> queue; + unordered_map::iterator> map; + + void push(Var* v, Var* prev); + void pop(Var* v); + void auto_flush(); +}; + +extern UpdateQueue update_queue; + +} // jittor + diff --git a/src/var.cc b/src/var.cc index fb53f1e4..7203bcf2 100644 --- a/src/var.cc +++ b/src/var.cc @@ -9,6 +9,7 @@ #include "op.h" #include "mem/allocator.h" #include "pybind/py_var_tracer.h" +#include "update_queue.h" namespace jittor { @@ -30,6 +31,8 @@ Var::~Var() { if (mem_ptr != nullptr) allocator->free(mem_ptr, size, allocation); number_of_lived_vars--; + if (flags.get(NodeFlags::_in_update_queue)) + update_queue.pop(this); } string Var::to_string() { diff --git a/src/var_holder.cc b/src/var_holder.cc index d53cfea8..3cb80a38 100644 --- a/src/var_holder.cc +++ b/src/var_holder.cc @@ -12,6 +12,7 @@ #include "var.h" #include "executor.h" #include "graph.h" +#include "update_queue.h" namespace jittor { @@ -75,6 +76,13 @@ VarHolder* VarHolder::assign(VarHolder* v) { return this; } +VarHolder* VarHolder::update(VarHolder* v) { + auto dv = jittor::detach(v->var); + update_queue.push(dv.ptr, var); + *this = move(dv); + return this; +} + extern Executor exe; void VarHolder::sync(bool device_sync) { diff --git a/src/var_holder.h b/src/var_holder.h index a543b5c2..c103378c 100644 --- a/src/var_holder.h +++ b/src/var_holder.h @@ -43,6 +43,15 @@ struct VarHolder { // @attrs(return_self) VarHolder* assign(VarHolder* v); + /* update parameter and global variable, + different from assign, it will + stop grad between origin var and assigned var, and + will update in the background + */ + // @pyjt(update) + // @attrs(return_self) + VarHolder* update(VarHolder* v); + // @pyjt(swap) // @attrs(return_self) inline VarHolder* swap(VarHolder* v) { std::swap(var, v->var); return this; };