add update queue

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):

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_mean.update(self.running_mean +
-            self.running_var += (xvar.sum([0,2,3])-self.running_var)*self.momentum
+                (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_mean.update(self.running_mean + 
-            self.running_var += (xvar.sum([0])-self.running_var)*self.momentum
+                (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])

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
+                p.update(p - g * lr)
-                # detach with the prev graph to reduce memory consumption
-                p.detach_inplace()
 
 
 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.update(p - v * lr)
-                p.detach_inplace()
 
 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)
+                v.update(alpha * v + (1-alpha) * g * g)
-                p -= lr * g / (jt.sqrt(v) + eps)
+                p.update(p - lr * g / (jt.sqrt(v) + eps))
-                p.detach_inplace()
 
 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())
+                values.append(jt.zeros(p.shape, p.dtype).stop_grad())
-                m.append(jt.zeros(p.shape, p.dtype).stop_fuse().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)
+                m.update(b0 * m + (1-b0) * g)
-                v.assign(b1 * v + (1-b1) * g * 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.update(p - m * step_size / (jt.sqrt(v) + eps))
-                p.detach_inplace()

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_mean.update(self.running_mean + 
-            self.running_var += (xvar.sum([0,2,3])-self.running_var)*self.momentum
+                (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])

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:
+            # print(jt.core.number_of_lived_vars(), mem_used)
-                assert jt.core.number_of_lived_vars() < 3900, jt.core.number_of_lived_vars()
+            jt.display_memory_info()
-            else:
+            # if jt.in_mpi:
-                assert jt.core.number_of_lived_vars() < 3500, jt.core.number_of_lived_vars()
+            #     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

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

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

src/update_queue.cc

@@ -0,0 +1,148 @@
+// ***************************************************************
+// Copyright (c) 2020 Jittor. Authors: Dun Liang <randonlang@gmail.com>. 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<Var*> 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<list<Item>>::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
+

src/update_queue.h

@@ -0,0 +1,27 @@
+// ***************************************************************
+// Copyright (c) 2020 Jittor. Authors: Dun Liang <randonlang@gmail.com>. 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<list<Item>>::iterator owner;
+        Var* v;
+    };
+    list<list<Item>> queue;
+    unordered_map<Var*, list<Item>::iterator> map;
+
+    void push(Var* v, Var* prev);
+    void pop(Var* v);
+    void auto_flush();
+};
+
+extern UpdateQueue update_queue;
+
+} // jittor
+

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() {

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) {

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; };