From 82f7ce5c6cd9841ded9a6b05bc0eebfa3e2fdafd Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Thu, 1 Sep 2022 06:21:46 +0000
Subject: [PATCH 01/10] init

---
 .gitignore                                    |  3 +-
 tests/test_all.py                             | 17 ++++
 tests/test_backward.py                        | 84 -------------------
 tests/test_dropout.py                         | 24 ++++--
 tests/{test_loss.py => test_loss_func.py}     | 37 ++------
 tests/{test_dist.py => test_nccl_backward.py} | 16 ++--
 tests/utils.py                                | 11 +++
 7 files changed, 63 insertions(+), 129 deletions(-)
 create mode 100644 tests/test_all.py
 delete mode 100644 tests/test_backward.py
 rename tests/{test_loss.py => test_loss_func.py} (74%)
 rename tests/{test_dist.py => test_nccl_backward.py} (63%)
 create mode 100644 tests/utils.py

diff --git a/.gitignore b/.gitignore
index 03e62300..0222862f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -149,4 +149,5 @@ log
 *.qdrep
 .vscode
 
-!bmtrain/dist
\ No newline at end of file
+!bmtrain/dist
+tests/test_log.txt
\ No newline at end of file
diff --git a/tests/test_all.py b/tests/test_all.py
new file mode 100644
index 00000000..c88520c0
--- /dev/null
+++ b/tests/test_all.py
@@ -0,0 +1,17 @@
+import subprocess
+from tqdm import tqdm
+
+
+tq = tqdm([
+    ("dropout", 1),
+    # ("nccl_backward", 4),
+    # ("loss_func", 1)
+])
+
+for t, num_gpu in tq:
+    PREFIX = f"python3 -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node={num_gpu} --master_addr=localhost --master_port=32123"
+    SUFFIX = f"> test_log.txt 2>&1"
+    command = f"{PREFIX} test_{t}.py {SUFFIX}"
+    completedProc = subprocess.run(command, shell=True)
+    assert completedProc.returncode == 0, f"test {t} failed, see test_log.txt for more details."
+    print(f"finish testing {t}")       
diff --git a/tests/test_backward.py b/tests/test_backward.py
deleted file mode 100644
index 6bda8070..00000000
--- a/tests/test_backward.py
+++ /dev/null
@@ -1,84 +0,0 @@
-import torch
-import bmtrain
-from tqdm import tqdm
-
-from bmtrain.utils import print_rank
-
-class DistributedLinear(bmtrain.DistributedModule):
-    def __init__(self, dim_in, dim_out):
-        super().__init__()
-        self.dim_in = dim_in
-        self.dim_out = dim_out
-
-        self.weight = bmtrain.DistributedParameter(torch.empty(dim_in, dim_out, dtype=torch.float))
-    
-    def forward(self, x : torch.Tensor):
-        last_dim = x.size(-1)
-        x_viewd = x.view(-1, last_dim)
-        out = torch.matmul(x_viewd, self.weight).view(x.size()[:-1] + (self.dim_out,))
-        var_out = out.var(dim=-1, keepdim=True)
-        return out / var_out + x
-
-class TestMLP(bmtrain.DistributedModule):
-    def __init__(self):
-        super().__init__()
-
-        self.mlp = torch.nn.ModuleList([
-            DistributedLinear(8, 8) for _ in range(16)
-        ])
-    
-    @bmtrain.checkpoint
-    def forward_segments(self, layer_ids, hidden_state : torch.Tensor):
-        for idx in layer_ids:
-            hidden_state = self.mlp[idx](hidden_state)
-        return hidden_state
-    
-    def forward(self, x):
-        for segments in [
-            [0, 1], [2, 3],
-            [4, 5], [6, 7],
-            [8, 9], [10, 11],
-            [12, 13], [14, 15]
-        ]:
-            x = self.forward_segments(segments, x)
-            bmtrain.wait_loader()
-        return x
-
-def main():
-    bmtrain.init_distributed()
-
-    model = TestMLP()
-
-    print_rank("Model mem\n", torch.cuda.memory_summary())
-    state = model.state_dict()
-    for kw, param in state.items():
-        torch.nn.init.normal_(param, std=0.02)
-    model.load_state_dict(state)    # split parameters automatically
-    print_rank("Model after loading state dict\n", torch.cuda.memory_summary())
-    
-    for i in range(bmtrain.world_size()):
-        v = torch.randn(4, 8, dtype=torch.float, device="cuda")
-        if i == bmtrain.rank():
-            raw_data = v
-
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-    for iter in tqdm(range(100)):
-        # load data
-        data = raw_data.clone().requires_grad_()
-        optimizer.zero_grad()
-        output = model(data)
-        loss = ((output - bmtrain.rank()) ** 2).mean()
-        print_rank("Iter %d, loss: " % iter, bmtrain.sum_loss(loss).item())
-        
-        loss.backward()
-        optimizer.step()
-        bmtrain.wait_optimizer() # loader stream wait for optimizer
-
-    print_rank(torch.cuda.memory_summary())
-    bmtrain.synchronize()
-    print_rank(torch.cuda.memory_summary(), rank=1)
-
-    
-
-if __name__ == "__main__":
-    main()
\ No newline at end of file
diff --git a/tests/test_dropout.py b/tests/test_dropout.py
index d72ee75f..a237857c 100644
--- a/tests/test_dropout.py
+++ b/tests/test_dropout.py
@@ -1,3 +1,5 @@
+from utils import *
+
 import torch
 import bmtrain as bmt
 
@@ -8,7 +10,6 @@ def __init__(self):
         self.drop = torch.nn.Dropout(p=0.5)
     
     def forward(self, x):
-        bmt.print_rank(x)
         return self.drop(x)
 
 class OutterModule(bmt.DistributedModule):
@@ -21,17 +22,22 @@ def __init__(self) -> None:
         ])
     
     def forward(self, x):
-        bmt.print_rank(self.training)
         return self.blk(x)
 
-def main():
-    bmt.init_distributed()
-
+def test_main():
     model = OutterModule()
 
-    x = torch.ones(32, device="cuda")
-    y = model(x)
-    bmt.print_rank(y)
+    for _ in range(5):
+        model.train()
+        x = torch.ones(32, device="cuda")
+        y = model(x)
+        assert_neq(x.numel()-y.nonzero().size(0), 0)
+
+        model.eval()
+        x = torch.ones(32, device="cuda")
+        y = model(x)
+        assert_eq(x.numel()-y.nonzero().size(0), 0)
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    bmt.init_distributed()
+    test_main()
\ No newline at end of file
diff --git a/tests/test_loss.py b/tests/test_loss_func.py
similarity index 74%
rename from tests/test_loss.py
rename to tests/test_loss_func.py
index ac7692de..a76be5f4 100644
--- a/tests/test_loss.py
+++ b/tests/test_loss_func.py
@@ -1,8 +1,9 @@
+from utils import *
+
 import torch
 import bmtrain as bmt
 import torch
 import random
-import time
 import copy
 
 def run(x, tgt, loss_func, bigmodel=None, scale=32768, use_float=False):
@@ -13,28 +14,23 @@ def run(x, tgt, loss_func, bigmodel=None, scale=32768, use_float=False):
         if bigmodel is not None:
             bigmodel = bigmodel.float()
     x = x.requires_grad_()
-    # print(torch.cuda.memory_summary())
     if bigmodel is None:
         loss = loss_func(x, tgt)
     else:
         t = bigmodel(x)
         loss = loss_func(t, tgt)
-        # loss = loss_func(bigmodel(x), tgt)
-    # print(torch.cuda.memory_summary())
     (loss * scale).backward()
-    # print(torch.cuda.memory_summary())
     return loss, x.grad
 
 def check(x, tgt, loss_func1, loss_func2, bigmodel=None):
     loss_1, grad_1 = run(x, tgt, loss_func1, bigmodel=bigmodel)
     loss_2, grad_2 = run(x, tgt, loss_func2, bigmodel=bigmodel, use_float=True)
-    assert(grad_1.isnan().sum()==0)
-    assert(grad_2.isnan().sum()==0)
-    print((loss_1 - loss_2).abs().item())
-    print((grad_1 - grad_2).abs().max().item())
-    print()
+    assert_eq(grad_1.isnan().sum(), 0)
+    assert_eq(grad_2.isnan().sum(), 0)
+    assert_lt((loss_1 - loss_2).abs().item(), 1e-5)
+    assert_lt((grad_1 - grad_2).abs().max().item(), 1e-2)
 
-def simple_test():
+def test_simple():
     loss_func1 = bmt.loss.FusedCrossEntropy()
     loss_func2 = torch.nn.CrossEntropyLoss()
 
@@ -83,24 +79,7 @@ def test_inplace():
         tgt = torch.randint(0, C, (N,)).cuda().long()
         check(x, tgt, loss_func1, loss_func2, bigmodel=bigmodel)
 
-
-def benchmark_memory():
-    N = 32 * 512
-    C = 50000
-
-    bigmodel = torch.nn.Linear(1, C).cuda().half()
-    x = torch.randn(N, 1).cuda().half()
-    tgt = torch.randint(0, C, (N,)).cuda().long()
-
-    # loss_func = bmt.loss.FusedCrossEntropy(inplace=True)
-    loss_func = bmt.loss.FusedCrossEntropy(inplace=False)
-    # loss_func = torch.nn.CrossEntropyLoss()
-
-    loss, grad = run(x, tgt, loss_func, bigmodel=bigmodel)
-    print(loss, grad)
-
 if __name__ == "__main__":
     test_other()
     test_inplace()
-    simple_test()
-    # benchmark_memory()
\ No newline at end of file
+    test_simple()
\ No newline at end of file
diff --git a/tests/test_dist.py b/tests/test_nccl_backward.py
similarity index 63%
rename from tests/test_dist.py
rename to tests/test_nccl_backward.py
index c6844ee3..e9ce87eb 100644
--- a/tests/test_dist.py
+++ b/tests/test_nccl_backward.py
@@ -1,15 +1,19 @@
+from utils import *
+
 import bmtrain as bmt
 import torch
 
-def main():
-    bmt.init_distributed()
+def test_main():
     x = torch.full((1,), bmt.rank() + 1, dtype=torch.half, device="cuda").requires_grad_(True)
     y = bmt.distributed.all_reduce(x, "prod").view(-1)
-    bmt.print_rank(y)
     loss = (y * y).sum() / 2
     loss.backward()
-    print(x.grad)
-
+    ref = y
+    for i in range(bmt.world_size()):
+        if i != bmt.rank(): ref *= i+1
+    assert_eq(x.grad, ref)
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    bmt.init_distributed()
+
+    test_main()
\ No newline at end of file
diff --git a/tests/utils.py b/tests/utils.py
new file mode 100644
index 00000000..1e60064a
--- /dev/null
+++ b/tests/utils.py
@@ -0,0 +1,11 @@
+def assert_eq(a, b):
+    assert a == b, f"{a} != {b}"
+
+def assert_neq(a, b):
+    assert a != b, f"{a} == {b}"
+
+def assert_lt(a, b):
+    assert a < b, f"{a} >= {b}"
+
+def assert_gt(a, b):
+    assert a > b, f"{a} <= {b}"
\ No newline at end of file

From 3b06d68c062ef9eed2cc2a13a23125d137076dfb Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Thu, 1 Sep 2022 06:57:34 +0000
Subject: [PATCH 02/10] automize more test

---
 tests/test_all.py                             |   5 +-
 ...{test_hasinfnan.py => test_has_inf_nan.py} |   4 +-
 ...{test_wrapper.py => test_model_wrapper.py} | 101 +++---------------
 tests/test_send_recv.py                       |  36 +++----
 4 files changed, 34 insertions(+), 112 deletions(-)
 rename tests/{test_hasinfnan.py => test_has_inf_nan.py} (93%)
 rename tests/{test_wrapper.py => test_model_wrapper.py} (72%)

diff --git a/tests/test_all.py b/tests/test_all.py
index c88520c0..4f34c604 100644
--- a/tests/test_all.py
+++ b/tests/test_all.py
@@ -3,7 +3,10 @@
 
 
 tq = tqdm([
-    ("dropout", 1),
+    ("model_wrapper", 4),
+    # ("send_recv", 4),
+    # ("has_inf_nan", 1),
+    # ("dropout", 1),
     # ("nccl_backward", 4),
     # ("loss_func", 1)
 ])
diff --git a/tests/test_hasinfnan.py b/tests/test_has_inf_nan.py
similarity index 93%
rename from tests/test_hasinfnan.py
rename to tests/test_has_inf_nan.py
index b310e546..b1b9b4a9 100644
--- a/tests/test_hasinfnan.py
+++ b/tests/test_has_inf_nan.py
@@ -1,3 +1,5 @@
+from utils import *
+
 import torch
 import bmtrain.optim._cuda as G
 import random
@@ -5,7 +7,7 @@
 def check(x, v):
     out = torch.zeros(1, dtype=torch.uint8, device="cuda")[0]
     G.f_has_inf_nan(x, out)
-    assert(out.item() == v)
+    assert_eq(out.item(), v)
 
 def test_main():
     for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
diff --git a/tests/test_wrapper.py b/tests/test_model_wrapper.py
similarity index 72%
rename from tests/test_wrapper.py
rename to tests/test_model_wrapper.py
index e2b0ea11..cb7eddfa 100644
--- a/tests/test_wrapper.py
+++ b/tests/test_model_wrapper.py
@@ -1,3 +1,5 @@
+from utils import *
+
 from typing import Optional
 import torch
 import math
@@ -169,11 +171,7 @@ def forward(self,
 
         return logits
 
-def main():
-    bmt.init_distributed(
-        seed=0,
-    )
-
+def test_main():
     model = GPT(
         num_layers=8,
         vocab_size=10240, 
@@ -186,16 +184,10 @@ def main():
         dtype=torch.half
     )
 
-    model.load_state_dict(torch.load("../example/ckpt-0.pt"))
-
-    model = bmt.BMTrainModelWrapper(model)
+    bmt_model = bmt.BMTrainModelWrapper(model)
+    model = model.cuda()
 
-    bmt.print_rank("Model memory")
-    bmt.print_rank(torch.cuda.memory_summary())
-    bmt.synchronize()
-
-    # data
-    # generate dummy data for each rank
+    # use the break if i == bmt.rank() to generate different data on different rank
     torch.manual_seed(1234)
     batch_size = 2
     seq_len = 512
@@ -214,80 +206,13 @@ def main():
 
         if i == bmt.rank():
             break
-    
-    loss_func = torch.nn.CrossEntropyLoss(ignore_index=-100)
-    optimizer = bmt.optim.AdamOffloadOptimizer(model.parameters(), weight_decay=1e-2)
-    lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0)
 
-    loss_scaler = bmt.optim.LossScaler(loss_scale=2**20)
-    loss_scaler.add_optimizer(optimizer, lr_scheduler)
-
-    bmt.synchronize()
-    
-    avg_time_recorder = bmt.utils.AverageRecorder()
-    avg_loss_recorder = bmt.utils.AverageRecorder()
-
-    for iteration in range(1000):
-        # load data
-        st = time.time()
-        loss_scaler.zero_grad()
-
-        with bmt.inspect.inspect_tensor() as inspector:
-            pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
-            logits = model(
-                enc_input,
-                pos,
-                pos < enc_length[:, None]
-            )
-            batch, seq_len, vocab_out_size = logits.size()
-
-            loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
-        
-            global_loss = bmt.sum_loss(loss).item()
-
-            loss = loss_scaler(loss)
-            loss.backward()
-        
-        # print inspected tensors in the forward & backward pass
-        # print parameters of the model
-        if iteration % 100 == 0:
-            bmt.print_rank(
-                bmt.inspect.format_summary(
-                    inspector.get_summary()
-                )
-            )
-            bmt.print_rank(
-                bmt.inspect.format_summary(
-                    bmt.inspect.inspect_model(model, "*")
-                )
-            )
-        
-
-        loss_scaler.step()
-
-        # record time and loss
-        iteration_time = time.time() - st
-
-        avg_time_recorder.record(iteration_time)
-        avg_loss_recorder.record(global_loss)
-
-        # print time and loss
-        bmt.print_rank(
-            "| Iter: {:6d} | loss: {:.4f} average_loss: {:.4f} | lr: {:.4e} scale: {:10.4f} | time: {:.4f}".format(
-                iteration,
-                global_loss,
-                avg_loss_recorder.value,
-                lr_scheduler.current_lr,
-                loss_scaler.loss_scale,
-                avg_time_recorder.value
-            )
-        )
-
-        # save model
-        if iteration % 1000 == 0:
-            bmt.save(model, "ckpt-%d.pt" % iteration)
-    
-    bmt.save(model, "checkpoint.pt")
+    pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
+    logits = model(enc_input, pos, pos < enc_length[:, None])
+    bmt_logits = bmt_model(enc_input, pos, pos < enc_length[:, None])
+    assert_eq((logits==bmt_logits).all(), True)
 
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    bmt.init_distributed(seed=0)
+
+    test_main()
\ No newline at end of file
diff --git a/tests/test_send_recv.py b/tests/test_send_recv.py
index f969d661..87a9fc96 100644
--- a/tests/test_send_recv.py
+++ b/tests/test_send_recv.py
@@ -1,27 +1,19 @@
-import bmtrain as bmt
+from utils import *
+
 import torch
+import bmtrain as bmt
 from bmtrain.global_var import config
-from bmtrain.pipe_comm import send_activations, recv_activations, gather_input
-from bmtrain import nccl
-from torch.distributed.distributed_c10d import recv
-from time import sleep
-def test_send_tensor():
-    bmt.init_distributed()
-    current_stream = torch.cuda.current_stream()
-    groups = [0,2]
-    rank = config['rank']
-def test_gather_input():
-    bmt.init_distributed(pipe_size=2)
-    if config['topology'].get_group_id("pipe") == 0:
-        a = torch.ones((2,1))
+
+def test_send_recv():
+    if config["topology"].stage_id == 0:
+        a = torch.ones((2,1)) * (config["zero_rank"]+1)
+        bmt.distributed.send_activations(a.cuda(), 1, config["pipe_comm"])
     else:
-        a = torch.zeros((2,1))
-    res = gather_input(a.cuda(),config['pipe_comm'])
-    if config['topology'].get_group_rank("pipe") == 0:
-        print(res)
-        
-def main():
-    test_send_tensor()
+        ref = torch.ones((2,1)) * (config["zero_rank"]+1)
+        a = bmt.distributed.recv_activations(0, config["pipe_comm"])
+        assert_eq((a == ref.cuda()).all(), True)
 
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    bmt.init_distributed(pipe_size=2)
+
+    test_send_recv()
\ No newline at end of file

From 4c44fb8f456ef13e8744bce09247f9e47d3e98c6 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Thu, 1 Sep 2022 07:26:12 +0000
Subject: [PATCH 03/10] automize more test

---
 tests/test_all.py           | 15 +++++----
 tests/test_dropout.py       |  2 ++
 tests/test_middle_hidden.py | 62 +++++++++++++++++++++----------------
 tests/test_model_wrapper.py |  3 ++
 tests/test_nccl_backward.py |  1 +
 tests/test_requires_grad.py | 59 ++++++++++++++++++++++-------------
 tests/test_send_recv.py     |  5 ++-
 7 files changed, 91 insertions(+), 56 deletions(-)

diff --git a/tests/test_all.py b/tests/test_all.py
index 4f34c604..5dad9b4e 100644
--- a/tests/test_all.py
+++ b/tests/test_all.py
@@ -3,12 +3,15 @@
 
 
 tq = tqdm([
+    ("requires_grad", 1),
+    ("has_inf_nan", 1),
+    ("dropout", 1),
+    ("loss_func", 1),
+
+    ("middle_hidden", 4),
     ("model_wrapper", 4),
-    # ("send_recv", 4),
-    # ("has_inf_nan", 1),
-    # ("dropout", 1),
-    # ("nccl_backward", 4),
-    # ("loss_func", 1)
+    ("send_recv", 4),
+    ("nccl_backward", 4),
 ])
 
 for t, num_gpu in tq:
@@ -17,4 +20,4 @@
     command = f"{PREFIX} test_{t}.py {SUFFIX}"
     completedProc = subprocess.run(command, shell=True)
     assert completedProc.returncode == 0, f"test {t} failed, see test_log.txt for more details."
-    print(f"finish testing {t}")       
+    print(f"finish testing {t}")
diff --git a/tests/test_dropout.py b/tests/test_dropout.py
index a237857c..f02da4e3 100644
--- a/tests/test_dropout.py
+++ b/tests/test_dropout.py
@@ -31,11 +31,13 @@ def test_main():
         model.train()
         x = torch.ones(32, device="cuda")
         y = model(x)
+        print(y)
         assert_neq(x.numel()-y.nonzero().size(0), 0)
 
         model.eval()
         x = torch.ones(32, device="cuda")
         y = model(x)
+        print(y)
         assert_eq(x.numel()-y.nonzero().size(0), 0)
 
 if __name__ == "__main__":
diff --git a/tests/test_middle_hidden.py b/tests/test_middle_hidden.py
index ca0dc9b3..2073588d 100644
--- a/tests/test_middle_hidden.py
+++ b/tests/test_middle_hidden.py
@@ -1,3 +1,5 @@
+from utils import *
+
 import bmtrain as bmt
 import random
 import torch
@@ -115,7 +117,7 @@ def manual_seed(seed=33):
     except ModuleNotFoundError:
         pass
 
-def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_middle=False, mix_test=False):
+def sub_run(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_middle=False, mix_test=False):
     manual_seed()
 
     pre = Linear(dim, dim)
@@ -135,15 +137,14 @@ def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_mi
         bmt.init_parameters(m)
     m = cls(pre, [m for m in ms], post)
 
+    ret = ""
     if only_last:
         logits = m(inp)
         loss = (logits * last_weight).sum()
         loss.backward()
-        bmt.print_rank(f"========================{name}:only last========================")
-        bmt.print_rank(
-            bmt.inspect.format_summary(
-                bmt.inspect.inspect_model(m, '*')
-            )
+        ret += f"========================only last========================\n"
+        ret += bmt.inspect.format_summary(
+            bmt.inspect.inspect_model(m, '*')
         )
     if only_middle:
         logits, hidden_states = m(inp, return_hidden_states=True)
@@ -152,11 +153,9 @@ def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_mi
             for i, hidden_state in enumerate(hidden_states)
         ])
         loss.backward()
-        bmt.print_rank(f"========================{name}:only middle========================")
-        bmt.print_rank(
-            bmt.inspect.format_summary(
-                bmt.inspect.inspect_model(m, '*')
-            )
+        ret += f"========================only middle========================\n"
+        ret += bmt.inspect.format_summary(
+            bmt.inspect.inspect_model(m, '*')
         )
     if mix_test:
         logits, hidden_states = m(inp, return_hidden_states=True)
@@ -165,24 +164,33 @@ def sub_test(name, cls, num_layer, dim, batch, seq_len, only_last=False, only_mi
             for i, hidden_state in enumerate(hidden_states)
         ]) + (logits * last_weight).sum()
         loss.backward()
-        bmt.print_rank(f"========================{name}:mix========================")
-        bmt.print_rank(
-            bmt.inspect.format_summary(
-                bmt.inspect.inspect_model(m, '*')
-            )
+        ret += f"========================mix========================\n"
+        ret += bmt.inspect.format_summary(
+            bmt.inspect.inspect_model(m, '*')
         )
+    return ret.replace("None  ", "0.0000") + "\n" # replace for matching None grad with zero_grad
 
-def test(name, cls, num_layer=4, dim=4096, batch=32, seq_len=256):
-    sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_last=True)
+def run(name, cls, num_layer=4, dim=4096, batch=32, seq_len=256):
+    ret = ""
+    ret += sub_run(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_last=True)
     bmt.synchronize()
-    sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_middle=True)
+    ret += sub_run(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, only_middle=True)
     bmt.synchronize()
-    sub_test(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, mix_test=True)
+    ret += sub_run(name, cls, num_layer=num_layer, dim=dim, batch=batch, seq_len=seq_len, mix_test=True)
     bmt.synchronize()
-
-bmt.init_distributed(pipe_size=4)
-
-test("normal", Model_NORMAL)
-test("block", Model_BLOCK)
-test("zero", Model_ZERO)
-test("pipe", Model_PIPE)
\ No newline at end of file
+    return ret
+
+def test_main():
+    ret = []
+    ret.append( run("normal", Model_NORMAL) )
+    ret.append( run("block", Model_BLOCK) )
+    ret.append( run("zero", Model_ZERO) )
+    ret.append( run("pipe", Model_PIPE) )
+    for r in ret:
+        bmt.prnit_rank(r)
+    for r in ret:
+        for r2 in ret:
+            assert_eq(r, r2)
+
+if __name__ == "__main__":
+    bmt.init_distributed(pipe_size=4)
diff --git a/tests/test_model_wrapper.py b/tests/test_model_wrapper.py
index cb7eddfa..0835a4da 100644
--- a/tests/test_model_wrapper.py
+++ b/tests/test_model_wrapper.py
@@ -210,6 +210,9 @@ def test_main():
     pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
     logits = model(enc_input, pos, pos < enc_length[:, None])
     bmt_logits = bmt_model(enc_input, pos, pos < enc_length[:, None])
+    print(logits)
+    bmt.synchronize()
+    print(bmt_logits)
     assert_eq((logits==bmt_logits).all(), True)
 
 if __name__ == '__main__':
diff --git a/tests/test_nccl_backward.py b/tests/test_nccl_backward.py
index e9ce87eb..3dcd0560 100644
--- a/tests/test_nccl_backward.py
+++ b/tests/test_nccl_backward.py
@@ -11,6 +11,7 @@ def test_main():
     ref = y
     for i in range(bmt.world_size()):
         if i != bmt.rank(): ref *= i+1
+    print(x.grad)
     assert_eq(x.grad, ref)
 
 if __name__ == "__main__":
diff --git a/tests/test_requires_grad.py b/tests/test_requires_grad.py
index 2a9bab43..707bd793 100644
--- a/tests/test_requires_grad.py
+++ b/tests/test_requires_grad.py
@@ -1,3 +1,5 @@
+from utils import *
+
 import bmtrain as bmt
 import torch
 from bmtrain import config
@@ -32,27 +34,40 @@ def run(m, a, b):
     loss = logits.sum()
     loss.backward()
 
-    bmt.print_rank(
-        bmt.inspect.format_summary(
+    sm = bmt.inspect.format_summary(
             bmt.inspect.inspect_model(m, '*')
         )
-    )
-    print(a.weight.grad is None)
-    print(a.bias.grad is None)
-
-bmt.init_distributed()
-a = Linear(256, 256)
-b = Linear(256, 256)
-m = TransformerBlockList([CheckpointBlock(a), CheckpointBlock(b)])
-bmt.init_parameters(m)
-
-a.bias.requires_grad_(False)
-run(m, a, b)
-
-a.weight.requires_grad_(False)
-a.bias.requires_grad_(True)
-run(m, a, b)
-
-a.weight.requires_grad_(True)
-a.bias.requires_grad_(False)
-run(m, a, b)
\ No newline at end of file
+    return a.weight.grad is None, a.bias.grad is None, sm
+
+def test_main():
+    a = Linear(256, 256)
+    b = Linear(256, 256)
+    m = TransformerBlockList([CheckpointBlock(a), CheckpointBlock(b)])
+    bmt.init_parameters(m)
+
+    a.bias.requires_grad_(False)
+    awg, abg, sm1 = run(m, a, b)
+    print(awg, abg, sm1)
+    assert_eq((awg, abg), (False, True))
+    assert_eq(sm1.split('\n')[2].split()[-2:], ["0.0000", "0.0000"])
+
+    a.weight.requires_grad_(False)
+    a.bias.requires_grad_(True)
+    awg, abg, sm2 = run(m, a, b)
+    print(awg, abg, sm2)
+    assert_eq((awg, abg), (False, False))
+    assert_eq(sm1.split('\n')[1], sm2.split('\n')[1])
+    assert_neq(sm1.split('\n')[2], sm2.split('\n')[2])
+
+    a.weight.requires_grad_(True)
+    a.bias.requires_grad_(False)
+    awg, abg, sm3 = run(m, a, b)
+    print(awg, abg, sm3)
+    assert_eq((awg, abg), (False, False))
+    assert_neq(sm2.split('\n')[1], sm3.split('\n')[1])
+    assert_eq(sm2.split('\n')[2], sm3.split('\n')[2])
+
+if __name__ == "__main__":
+    bmt.init_distributed()
+
+    test_main()
\ No newline at end of file
diff --git a/tests/test_send_recv.py b/tests/test_send_recv.py
index 87a9fc96..2a176898 100644
--- a/tests/test_send_recv.py
+++ b/tests/test_send_recv.py
@@ -7,10 +7,13 @@
 def test_send_recv():
     if config["topology"].stage_id == 0:
         a = torch.ones((2,1)) * (config["zero_rank"]+1)
-        bmt.distributed.send_activations(a.cuda(), 1, config["pipe_comm"])
+        a = a.cuda()
+        print(f"send {a}")
+        bmt.distributed.send_activations(a, 1, config["pipe_comm"])
     else:
         ref = torch.ones((2,1)) * (config["zero_rank"]+1)
         a = bmt.distributed.recv_activations(0, config["pipe_comm"])
+        print(f"recv {a}")
         assert_eq((a == ref.cuda()).all(), True)
 
 if __name__ == '__main__':

From 909d39e1cfb5d886a1aec5dd4c1e0fd5ae989ddd Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 2 Sep 2022 01:13:14 +0000
Subject: [PATCH 04/10] add test for init parameters and save/load parameters

---
 bmtrain/param_init.py                   |   1 -
 tests/test_all.py                       |   4 +
 tests/test_init_parameters.py           | 219 ++++++++++++++++++++++++
 tests/test_init_parameters_multi_gpu.py | 146 ++++++++++++++++
 tests/test_load_ckpt.py                 |  66 +++++++
 tests/test_model_wrapper.py             |   2 +-
 tests/test_send_recv.py                 |   2 +-
 tests/utils.py                          |   5 +-
 8 files changed, 441 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_init_parameters.py
 create mode 100644 tests/test_init_parameters_multi_gpu.py
 create mode 100644 tests/test_load_ckpt.py

diff --git a/bmtrain/param_init.py b/bmtrain/param_init.py
index c95f90f8..125283e3 100644
--- a/bmtrain/param_init.py
+++ b/bmtrain/param_init.py
@@ -22,7 +22,6 @@ def init_distributed_parameter(params : Iterable[torch.nn.Parameter]):
             param._init_method(tmp_tensor)
 
             # Pytorch 1.11 changed the API of storage.__getitem__
-            # use zero_rank to support pipeline
             torch.tensor([], dtype=param.dtype, device=param.device).set_(param.storage())[:] = \
                 torch.tensor([], dtype=param.dtype, device=param.device).set_(tmp_storage)[partition_size * config['rank'] : partition_size * (config['rank'] + 1)]
             # param.storage().copy_(tmp_storage[partition_size * config['rank'] : partition_size * (config['rank'] + 1)])
diff --git a/tests/test_all.py b/tests/test_all.py
index 5dad9b4e..af5a0253 100644
--- a/tests/test_all.py
+++ b/tests/test_all.py
@@ -3,6 +3,10 @@
 
 
 tq = tqdm([
+    ("load_ckpt", 1),
+    ("init_parameters", 1),
+    ("init_parameters_multi_gpu", 4),
+
     ("requires_grad", 1),
     ("has_inf_nan", 1),
     ("dropout", 1),
diff --git a/tests/test_init_parameters.py b/tests/test_init_parameters.py
new file mode 100644
index 00000000..c23249ff
--- /dev/null
+++ b/tests/test_init_parameters.py
@@ -0,0 +1,219 @@
+from utils import *
+
+import torch
+import torch.nn.functional as F
+import bmtrain as bmt
+
+def manual_seed(seed=33):
+    torch.manual_seed(seed)
+    import random as random
+    random.seed(seed)
+    try:
+        import numpy as np
+        np.random.seed(seed)
+    except ModuleNotFoundError:
+        pass
+
+class Linear_NormalInitBefore(torch.nn.Module):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        w = torch.empty(out_features, in_features, dtype=dtype, device="cuda") # use cuda to match random algorithm
+        torch.nn.init.xavier_normal_(w)
+        self.weight = torch.nn.Parameter(w)
+        if bias:
+            b = torch.empty(out_features, dtype=dtype, device="cuda") # use cuda to match random algorithm
+            torch.nn.init.zeros_(b)
+            self.bias = torch.nn.Parameter(b)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+
+class Linear_NormalInitAfter(torch.nn.Module):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm
+        torch.nn.init.xavier_normal_(self.weight)
+        if bias:
+            self.bias = torch.nn.Parameter(torch.empty(out_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm
+            torch.nn.init.zeros_(self.bias)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+
+class Linear_BMTInitializer(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=dtype), init_method=torch.nn.init.xavier_normal_)
+        if bias:
+            self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=dtype), init_method=torch.nn.init.zeros_)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+
+class Linear_ManualInitBefore(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        w = torch.empty(out_features, in_features, dtype=dtype, device="cuda") # use cuda to match random algorithm
+        torch.nn.init.xavier_normal_(w)
+        self.weight = bmt.DistributedParameter(w)
+        if bias:
+            b = torch.empty(out_features, dtype=dtype, device="cuda") # use cuda to match random algorithm
+            torch.nn.init.zeros_(b)
+            self.bias = bmt.DistributedParameter(b)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+
+class Linear_ManualInitAfter(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm
+        # torch.nn.init.xavier_normal_(self.weight)
+        if bias:
+            self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm
+            # torch.nn.init.zeros_(self.bias)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+
+    def extra_repr(self) -> str:
+        return 'in_features={}, out_features={}, bias={}'.format(
+            self.in_features, self.out_features, self.bias is not None
+        )
+
+class Linear_Pipeline(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.l = bmt.PipelineTransformerBlockList([
+            Linear_BMTInitializer(in_features, out_features, bias, dtype),
+            Linear_BMTInitializer(in_features, out_features, bias, dtype),
+        ])
+    
+    def forward(self, input):
+        return self.l(input)
+
+class Linear_BlockList(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.l = bmt.TransformerBlockList([
+            Linear_BMTInitializer(in_features, out_features, bias, dtype),
+            Linear_BMTInitializer(in_features, out_features, bias, dtype),
+        ])
+    
+    def forward(self, input):
+        return self.l(input)
+
+class Linear_CheckpointList(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.l = torch.nn.ModuleList([
+            bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)),
+            bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)),
+        ])
+    
+    def forward(self, input):
+        return self.l(input)
+
+class Linear_Checkpoint(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.l = bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype))
+    
+    def forward(self, input):
+        return self.l(input)
+    
+def test_main():
+    shape = [3, 5]
+    # torch
+    m = [None] * 10
+    ret = [None] * 10
+    manual_seed(33)
+    m[0] = Linear_NormalInitBefore(*shape)
+    ret[0] = (m[0].weight.data, m[0].bias.data)
+
+    manual_seed(33)
+    m[1] = Linear_NormalInitAfter(*shape)
+    ret[1] = (m[1].weight.data, m[1].bias.data)
+
+    # bmtrain
+    manual_seed(33)
+    m[2] = Linear_BMTInitializer(*shape)
+    bmt.init_parameters(m[2])
+    ret[2] = (m[2].weight.data, m[2].bias.data)
+
+    manual_seed(33)
+    m[3] = Linear_ManualInitBefore(*shape)
+    ret[3] = (m[3].weight.data, m[3].bias.data)
+
+    # manual_seed(33)
+    # mw = Linear_ManualInitAfter(*shape) # not supported
+    # print(mw.weight.data, mw.bias.data)
+
+    manual_seed(33)
+    m[4] = bmt.BMTrainModelWrapper(m[0])
+    ret[4] = (m[4].weight.data, m[4].bias.data)
+
+    manual_seed(33)
+    m[5] = bmt.BMTrainModelWrapper(m[1])
+    ret[5] = (m[5].weight.data, m[5].bias.data)
+
+    manual_seed(33)
+    m[6] = Linear_Pipeline(*shape)
+    bmt.init_parameters(m[6])
+    ret[6] = (m[6].l[0].weight.data, m[6].l[0].bias.data)
+
+    manual_seed(33)
+    m[7] = Linear_BlockList(*shape)
+    bmt.init_parameters(m[7])
+    ret[7] = (m[7].l[0].weight.data, m[7].l[0].bias.data)
+
+    manual_seed(33)
+    m[8] = Linear_CheckpointList(*shape)
+    bmt.init_parameters(m[8])
+    ret[8] = (m[8].l[0].weight.data, m[8].l[0].bias.data)
+
+    manual_seed(33)
+    m[9] = Linear_Checkpoint(*shape)
+    bmt.init_parameters(m[9])
+    ret[9] = (m[9].l.weight.data, m[9].l.bias.data)
+
+    for i in range(10):
+        ret[i] = ( ret[i][0].view(-1), ret[i][1].view(-1) )
+    for i in range(10):
+        for j in range(10):
+            assert_all_eq(ret[i][0], ret[j][0])
+            assert_all_eq(ret[i][1], ret[j][1])
+
+if __name__ == "__main__":
+    bmt.init_distributed(pipe_size=1)
+
+    test_main()
\ No newline at end of file
diff --git a/tests/test_init_parameters_multi_gpu.py b/tests/test_init_parameters_multi_gpu.py
new file mode 100644
index 00000000..1e61568c
--- /dev/null
+++ b/tests/test_init_parameters_multi_gpu.py
@@ -0,0 +1,146 @@
+from utils import *
+
+import os
+import torch
+import torch.nn.functional as F
+import bmtrain as bmt
+
+def manual_seed(seed=33):
+    torch.manual_seed(seed)
+    import random as random
+    random.seed(seed)
+    try:
+        import numpy as np
+        np.random.seed(seed)
+    except ModuleNotFoundError:
+        pass
+
+class Linear_Normal(torch.nn.Module):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm
+        torch.nn.init.xavier_normal_(self.weight)
+        if bias:
+            self.bias = torch.nn.Parameter(torch.empty(out_features, dtype=dtype, device="cuda")) # use cuda to match random algorithm
+            torch.nn.init.zeros_(self.bias)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+
+class Linear_BMTInitializer(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=dtype), init_method=torch.nn.init.xavier_normal_)
+        if bias:
+            self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=dtype), init_method=torch.nn.init.zeros_)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+
+class Linear_NormalList(torch.nn.Module):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.l = torch.nn.ModuleList([
+            Linear_Normal(in_features, out_features, bias, dtype),
+            Linear_Normal(in_features, out_features, bias, dtype),
+        ])
+    
+    def forward(self, input):
+        return self.l(input)
+
+class Linear_Pipeline(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.l = bmt.PipelineTransformerBlockList([
+            Linear_BMTInitializer(in_features, out_features, bias, dtype),
+            Linear_BMTInitializer(in_features, out_features, bias, dtype),
+        ])
+    
+    def forward(self, input):
+        return self.l(input)
+
+class Linear_BlockList(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.l = bmt.TransformerBlockList([
+            Linear_BMTInitializer(in_features, out_features, bias, dtype),
+            Linear_BMTInitializer(in_features, out_features, bias, dtype),
+        ])
+    
+    def forward(self, input):
+        return self.l(input)
+
+class Linear_CheckpointList(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.l = torch.nn.ModuleList([
+            bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)),
+            bmt.CheckpointBlock(Linear_BMTInitializer(in_features, out_features, bias, dtype)),
+        ])
+    
+    def forward(self, input):
+        return self.l(input)
+
+def check(ckpt_path, ckpt_path_ref):
+    if bmt.rank() == 0:
+        ckpt1 = torch.load(ckpt_path)
+        ckpt2 = torch.load(ckpt_path_ref)
+        for (k1, v1), (k2, v2) in zip(ckpt1.items(), ckpt2.items()):
+            assert_eq(k1, k2)
+            print(v1, v2)
+            assert_all_eq(v1.cuda(), v2.cuda())
+    
+def test_main():
+    ckpt_path_ref = "test_ckpt_ref.pt"
+    ckpt_path = "test_ckpt.pt"
+    shape = [3, 5]
+    # torch
+    m = [None] * 4
+    ret = [None] * 4
+
+    manual_seed(33)
+    m[0] = Linear_NormalList(*shape)
+    if bmt.rank() == 0:
+        torch.save(m[0].state_dict(), ckpt_path_ref)
+
+    manual_seed(33)
+    m[1] = Linear_Pipeline(*shape)
+    bmt.init_parameters(m[1])
+    bmt.save(m[1], ckpt_path)
+    check(ckpt_path, ckpt_path_ref)
+
+    # bmtrain
+    manual_seed(33)
+    m[2] = Linear_BlockList(*shape)
+    bmt.init_parameters(m[2])
+    bmt.save(m[2], ckpt_path)
+    check(ckpt_path, ckpt_path_ref)
+
+    manual_seed(33)
+    m[3] = Linear_CheckpointList(*shape)
+    bmt.init_parameters(m[3])
+    bmt.save(m[3], ckpt_path)
+    check(ckpt_path, ckpt_path_ref)
+
+    if bmt.rank() == 0:
+        os.remove(ckpt_path)
+        os.remove(ckpt_path_ref)
+
+if __name__ == "__main__":
+    bmt.init_distributed(pipe_size=2)
+
+    test_main()
\ No newline at end of file
diff --git a/tests/test_load_ckpt.py b/tests/test_load_ckpt.py
new file mode 100644
index 00000000..56298940
--- /dev/null
+++ b/tests/test_load_ckpt.py
@@ -0,0 +1,66 @@
+from utils import *
+import torch
+import torch.nn.functional as F
+import bmtrain as bmt
+import os
+
+class Linear_Normal(torch.nn.Module):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
+        torch.nn.init.xavier_normal_(self.weight)
+        if bias:
+            self.bias = torch.nn.Parameter(torch.empty(out_features, dtype=dtype))
+            torch.nn.init.zeros_(self.bias)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+
+class Linear_BMT(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, bias: bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=dtype), init_method=torch.nn.init.xavier_normal_)
+        if bias:
+            self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=dtype), init_method=torch.nn.init.zeros_)
+        else:
+            self.register_parameter('bias', None)
+    
+    def forward(self, input):
+        return F.linear(input, self.weight, self.bias)
+    
+def test_main():
+    ckpt_path = "test_ckpt.pt"
+    m = Linear_Normal(7, 5).cuda()
+    m2 = Linear_BMT(7, 5)
+    torch.save(m.state_dict(), ckpt_path)
+    bmt.load(m2, ckpt_path)
+
+    print(m.weight.data)
+    print(m2.weight.data)
+    assert_all_eq(m.weight.data, m2.weight.data)
+    assert_all_eq(m.bias.data, m2.bias.data)
+
+    os.remove(ckpt_path)
+
+    bmt.save(m2, ckpt_path)
+    m.load_state_dict(torch.load(ckpt_path))
+
+    print(m.weight.data)
+    print(m2.weight.data)
+    assert_all_eq(m.weight.data, m2.weight.data)
+    assert_all_eq(m.bias.data, m2.bias.data)
+
+    os.remove(ckpt_path)
+
+if __name__ == "__main__":
+    bmt.init_distributed()
+
+    test_main()
\ No newline at end of file
diff --git a/tests/test_model_wrapper.py b/tests/test_model_wrapper.py
index 0835a4da..409107e3 100644
--- a/tests/test_model_wrapper.py
+++ b/tests/test_model_wrapper.py
@@ -213,7 +213,7 @@ def test_main():
     print(logits)
     bmt.synchronize()
     print(bmt_logits)
-    assert_eq((logits==bmt_logits).all(), True)
+    assert_all_eq(logits, bmt_logits)
 
 if __name__ == '__main__':
     bmt.init_distributed(seed=0)
diff --git a/tests/test_send_recv.py b/tests/test_send_recv.py
index 2a176898..95c9c1e5 100644
--- a/tests/test_send_recv.py
+++ b/tests/test_send_recv.py
@@ -14,7 +14,7 @@ def test_send_recv():
         ref = torch.ones((2,1)) * (config["zero_rank"]+1)
         a = bmt.distributed.recv_activations(0, config["pipe_comm"])
         print(f"recv {a}")
-        assert_eq((a == ref.cuda()).all(), True)
+        assert_all_eq(a, ref.cuda())
 
 if __name__ == '__main__':
     bmt.init_distributed(pipe_size=2)
diff --git a/tests/utils.py b/tests/utils.py
index 1e60064a..62831fea 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -8,4 +8,7 @@ def assert_lt(a, b):
     assert a < b, f"{a} >= {b}"
 
 def assert_gt(a, b):
-    assert a > b, f"{a} <= {b}"
\ No newline at end of file
+    assert a > b, f"{a} <= {b}"
+
+def assert_all_eq(a, b):
+    assert_eq((a==b).all(), True)
\ No newline at end of file

From 92cd2cbeb817a174f341e32a1eff5700ecd6d163 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 2 Sep 2022 10:08:47 +0000
Subject: [PATCH 05/10] FIX: adam eps scale

---
 csrc/cuda/adam.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/csrc/cuda/adam.cu b/csrc/cuda/adam.cu
index 4b3df9c7..c7bff5ca 100644
--- a/csrc/cuda/adam.cu
+++ b/csrc/cuda/adam.cu
@@ -26,7 +26,7 @@ __global__ void adam_fp32_accum(
         float local_m = beta1 * __half2float(m[col]) + (1 - beta1) * local_g;       // real_m * scale
         float local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g / scale;   // real_v * scale
         float local_p = param[col];
-        local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v * scale / bias_correction2) + eps) - lr * weight_decay * local_p;
+        local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v * scale / bias_correction2) + eps * scale) - lr * weight_decay * local_p;
 
         param_h[col] = __float2half(local_p);
         param[col] = local_p;

From 261146ca3fee446b5326420d18a33cbce2dd9b8f Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 2 Sep 2022 10:10:00 +0000
Subject: [PATCH 06/10] FIX: nccl reduce use no avg op for compatibility

---
 bmtrain/inspect/tensor.py | 8 ++++----
 bmtrain/synchronize.py    | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/bmtrain/inspect/tensor.py b/bmtrain/inspect/tensor.py
index 7d56ad7f..526ab8fc 100644
--- a/bmtrain/inspect/tensor.py
+++ b/bmtrain/inspect/tensor.py
@@ -129,9 +129,9 @@ def get_summary(self):
                     nccl.allReduce(
                         info.storage(),
                         info.storage(),
-                        "avg",
+                        "sum",
                         comm
-                    )
+                    ) / nccl.commCount(comm)
                     x_mean = info[0].cpu().item()
                     x_std = math.sqrt(info[1].cpu().item())
                     grad_mean = None
@@ -146,9 +146,9 @@ def get_summary(self):
                     nccl.allReduce(
                         info.storage(),
                         info.storage(),
-                        "avg",
+                        "sum",
                         comm
-                    )
+                    ) / nccl.commCount(comm)
                     x_mean = info[0].cpu().item()
                     x_std = math.sqrt(info[1].cpu().item())
                     grad_mean = info[2].cpu().item()
diff --git a/bmtrain/synchronize.py b/bmtrain/synchronize.py
index f704d159..80d39057 100644
--- a/bmtrain/synchronize.py
+++ b/bmtrain/synchronize.py
@@ -31,7 +31,7 @@ def sum_loss(loss : torch.Tensor):
     This is a helper function to reduce the loss across all workers.
     """
     warnings.warn("bmtrain.sum_loss is deprecated and will be removed in later version. Use bmtrain.distributed.all_reduce instead.", DeprecationWarning)
-    return distributed.all_reduce(loss, "avg")
+    return distributed.all_reduce(loss, "sum") / config['world_size']
 
 def gather_result(result: torch.Tensor):
     warnings.warn("bmtrain.gather_result is deprecated and will be removed in later version. Use bmtrain.distributed.all_gather instead.", DeprecationWarning)

From d3ae9d9b05845eb959f8ea1e38bb199369abcdf4 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 2 Sep 2022 10:10:29 +0000
Subject: [PATCH 07/10] FIX: pipeline loss_scale factor

---
 bmtrain/optim/loss_scaler.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/bmtrain/optim/loss_scaler.py b/bmtrain/optim/loss_scaler.py
index 2a243863..9fbaed37 100644
--- a/bmtrain/optim/loss_scaler.py
+++ b/bmtrain/optim/loss_scaler.py
@@ -70,7 +70,7 @@ def __call__(self, loss : torch.Tensor) -> torch.Tensor:
         Args:
             loss (torch.Tensor): loss
         """
-        return loss * (self.loss_scale * config["pipe_size"] / config['world_size'])
+        return loss * (self.loss_scale / config['world_size'])
 
     def zero_grad(self):
         """
@@ -146,7 +146,7 @@ def clip_grad_norm(self, param_groups, max_norm, norm_type=2, eps=1e-6):
             >>> loss_scaler.clip_grad_norm(optimizer.param_groups, max_norm=1.0, norm_type=2)
 
         """
-        scale = self.loss_scale * config['pipe_size'] / config['world_size']
+        scale = self.loss_scale / config['world_size']
         grads = []
         parameters = [p for group in param_groups for p in group['params']]
         for p in parameters:

From 42f52d11a7e19f6341ee3709d2218565148a9da7 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 2 Sep 2022 12:09:00 +0000
Subject: [PATCH 08/10] unify loss_scale and normal -> OptimManager

---
 README-ZH.md                                  | 61 +++++-----------
 README.md                                     | 33 +++++----
 bmtrain/__init__.py                           |  1 -
 bmtrain/backward.py                           | 28 --------
 bmtrain/optim/__init__.py                     |  2 +-
 .../{loss_scaler.py => optim_manager.py}      | 70 +++++++++++--------
 example/train.py                              | 15 ++--
 7 files changed, 86 insertions(+), 124 deletions(-)
 delete mode 100644 bmtrain/backward.py
 rename bmtrain/optim/{loss_scaler.py => optim_manager.py} (76%)

diff --git a/README-ZH.md b/README-ZH.md
index 264edbd5..598777b2 100644
--- a/README-ZH.md
+++ b/README-ZH.md
@@ -273,19 +273,16 @@ BMTrain 支持**所有** PyTorch 原生的优化器和损失函数，同时你
 ### 第 4 部分: 训练
 
 ```python
-# 新建损失缩放器实例
-loss_scaler = bmtrain.optim.LossScaler()
-# 将所有的 optimzer 及(可选)其对应的 lr_scheduler 收入损失缩放器管理。
-loss_scaler.add_optimizer(optimizer, lr_scheduler)
+# 新建优化器管理器实例
+optim_manager = bmtrain.optim.OptimManager(loss_scale=1024)
+# 将所有的 optimzer 及(可选)其对应的 lr_scheduler 收入优化器管理器管理。
+optim_manager.add_optimizer(optimizer, lr_scheduler)
 # 可以再次调用 add_optimizer 加入其他优化器
 
 for iteration in range(1000):
     # ... 为每个rank加载数据 ...
 
-    # 梯度清零
-    loss_scaler.zero_grad() # 为每个 optimizer 调用 zero_grad
-
-    # 前向传播
+    # 前向传播并计算梯度
     pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
     logits = model(
         enc_input,
@@ -296,14 +293,19 @@ for iteration in range(1000):
 
     loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
     
-    global_loss = bmtrain.sum_loss(loss).item() # 聚合所有rank上的损失
+    global_loss = bmtrain.sum_loss(loss).item() # 聚合所有rank上的损失, 仅用于输出训练日志
+
+    # 梯度清零
+    optim_manager.zero_grad() # 为每个 optimizer 调用 zero_grad
 
     # 损失缩放和反向传播
-    loss = loss_scaler(loss)
-    loss.backward()
+    optim_manager.backward(loss)
+
+    # 梯度裁剪
+    grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=1.0)
 
     # 更新参数
-    loss_scaler.step()
+    optim_manager.step()
 
     # ... 保存checkpoint、打印日志 ...
 ```
@@ -312,40 +314,11 @@ for iteration in range(1000):
 
 你可以根据代码中的注释来了解各部分代码的作用。
 
-唯一需要说明的是 `loss_scaler`，**损失缩放**是混合精度训练中的一项常用技术，需要在反向传播前通过 `loss_scaler(loss)` 对 `loss` 进行放缩，用于避免梯度下溢。在使用损失所放后，优化器的 `step()` 等步骤需要有一些细节上的调整。我们在 `loss_scaler` 帮你实现了这些细节, 你只需要通过 `add_optimizer` 将优化器和学习率调整策略收入 `loss_scaler` 管理，并由 `loss_scaler` 代为执行 `zero_grad()` 和 `step()` 操作。
-
-如果你没有使用 BMTrain 中的融合优化器，你可以不用 `loss_scaler`, 相应的代码修改为:
-
-```python
-for iteration in range(1000):
-    # ... 为每个rank加载数据 ...
-
-    # 梯度清零
-    optimizer.zero_grad()
-
-    # 前向传播
-    pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
-    logits = model(
-        enc_input,
-        pos,
-        pos < enc_length[:, None]
-    )
-    batch, seq_len, vocab_out_size = logits.size()
-
-    loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
-    
-    global_loss = bmtrain.sum_loss(loss).item() # 聚合所有rank上的损失
+唯一需要说明的是 `optim_manager`。在使用 BMTrain 后，优化器的部分相关操作需要有一些细节上的调整。我们在 `optim_manager` 帮你实现了这些细节, 你只需要通过 `add_optimizer` 将优化器和学习率调整策略收入 `optim_manager` 管理，并由 `optim_manger` 代为执行 `zero_grad()`, `backward()`, `clip_grad_norm()` 和 `step()` 等操作。
 
-    # 反向传播
-    loss.backward()
-
-    # 更新参数
-    bmtrain.optim_step(optimizer, lr_scheduler)
-
-    # ... 保存checkpoint、打印日志 ...
-```
+如果你没有使用混合精度训练，你可以不用损失缩放，只需要将 `OptimManger(loss_scale=None)` 构造函数中 `loss_scale` 置为 None 即可, 这也是 `OptimManager` 的默认构造参数。
 
-需要说明的是最后更新参数时，需要使用 `bmtrain.optim_step`, 而不能直接使用 `optimizer.step()` 和 `lr_scheduler.step()`。
+如果你使用了混合精度训练，**损失缩放**是混合精度训练中的一项常用技术，我们在 `optim_manager.backward(loss)` 帮你对 `loss` 进行了放缩，用于避免梯度下溢。只需要将 `OptimManger` 构造函数中 `loss_scale` 置为一个浮点数即可。 `loss_scale` 会在训练过程中根据梯度进行自适应的调整。
 
 <div id="性能"></div>
 
diff --git a/README.md b/README.md
index 16e50037..45ed273c 100644
--- a/README.md
+++ b/README.md
@@ -276,19 +276,16 @@ In addition, BMTrain also provides the common LRScheduler in the `bmtrain.lr_sch
 ### Part 4: Training Loop
 
 ```python
-# create a new instance of loss scaler
-loss_scaler = bmtrain.optim.LossScaler()
-# let loss_scaler handle all the optimizer and (optional) their corresponding lr_scheduler
-loss_scaler.add_optimizer(optimizer, lr_scheduler)
+# create a new instance of optimizer manager
+optim_manager = bmtrain.optim.OptimManager()
+# let optim_manager handle all the optimizer and (optional) their corresponding lr_scheduler
+optim_manager.add_optimizer(optimizer, lr_scheduler)
 # add_optimizer can be called multiple times to add other optimizers.
 
 for iteration in range(1000):
     # ... load data for each rank ...
 
-    # zero grad
-    loss_scaler.zero_grad() # calling zero_grad for each optimizer
-
-    # forward
+    # forward pass and calculate loss
     pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
     logits = model(
         enc_input,
@@ -299,14 +296,19 @@ for iteration in range(1000):
 
     loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
     
-    global_loss = bmtrain.sum_loss(loss).item() # sum the loss across all ranks
+    global_loss = bmtrain.sum_loss(loss).item() # sum the loss across all ranks. This is only used for the training log
+
+    # zero grad
+    optim_manager.zero_grad() # calling zero_grad for each optimizer
+
+    # clip grad norm
+    grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=1.0)
 
     # loss scale and backward
-    loss = loss_scaler(loss)
-    loss.backward()
+    optim_manager.backward()
 
     # optimizer step
-    loss_scaler.step()
+    optim_manager.step()
 
     # ... save checkpoint or print logs ...
 ```
@@ -315,9 +317,12 @@ The training loop part will be slightly longer, but just like a normal training
 
 You can follow the comments in the code to get an idea of what each section of code is doing.
 
-The only additional note is `loss_scaler`, *loss scale* is the technique widely used in mixed precision training to prevent gradient underflow by adding `loss_scaler(loss)` to scale the `loss` before backward. When using loss scaling, some details in optimizers should be adjusted. We have implemented all those details needed in `loss_scaler`. What you need is just letting `loss_scaler` to handle all the optimizers by `add_optimizer`, and letting `loss_scaler` do `zero_grad()` and `step()` instead.
+The only additional note is `optimizer`. After using BMTrain, some details in optimizers should be adjusted. We have implemented all those details needed in `optim_manager`. What you need is just letting `optim_manager` to handle all the optimizers by `add_optimizer`, and letting `optim_manager` do `zero_grad()`, `backward()`, `clip_grad_norm()` and `step()` instead.
+
+If you are not using the mixed-precision training, you can train without `loss_scale`. Just set `loss_scale` to None in the `__init__` function of `OptimManager(loss_scale=None)`, which is also the default.
+
+If you are using mixed-precision training, *loss scale* is the technique widely used in mixed precision training to prevent gradient underflow. By using `optim_manager.backward(loss)` to scale the `loss` before backward and set `loss_scale` to some floating number in the `__init__` function of `OptimManager`。The `loss_scale` would be adjusted adaptively based on the gradient during training.
 
-If you are not using the fused optimizer in BMTrain, you can train without `loss_scaler`. The code will be changed as below:
 
 ```python
 for iteration in range(1000):
diff --git a/bmtrain/__init__.py b/bmtrain/__init__.py
index 25c72052..14ff50de 100644
--- a/bmtrain/__init__.py
+++ b/bmtrain/__init__.py
@@ -8,7 +8,6 @@
 from .synchronize import synchronize, sum_loss, wait_loader, gather_result
 from .checkpointing import checkpoint
 from .block_layer import CheckpointBlock, TransformerBlockList
-from .backward import optim_step
 from .wrapper import BMTrainModelWrapper
 from .pipe_layer import PipelineTransformerBlockList
 from . import debug
diff --git a/bmtrain/backward.py b/bmtrain/backward.py
deleted file mode 100644
index c503c972..00000000
--- a/bmtrain/backward.py
+++ /dev/null
@@ -1,28 +0,0 @@
-from typing import Optional
-import torch
-from .global_var import config
-from .utils import print_rank
-from .lr_scheduler.warmup import WarmupLRScheduler
-
-
-def optim_step(optim : torch.optim.Optimizer, lr_scheduler : Optional[WarmupLRScheduler] = None):
-    """
-    This is a helper function to call optimizer.step() and lr_scheduler.step() and synchronize streams.
-
-    Args:
-        optim (torch.optim.Optimizer): A pytorch optimizer, e.g. torch.optim.Adam, torch.optim.SGD or bmtrain.optim.AdamOffloadOptimizer
-        lr_scheduler (Optional[WarmupLRScheduler]): A warmup lr scheduler, e.g. bmt.lr_scheduler.Noam
-    
-    This function can also handle gradient overflow by reducing the loss scale when it occurs.
-
-    """
-    
-    current_stream =  torch.cuda.current_stream()
-    # some reduce ops of distributed parameter were launched on load stream
-    current_stream.wait_stream(config['load_stream'])
-
-    optim.step()
-    if lr_scheduler is not None:
-        lr_scheduler.step()
-
-    config['load_stream'].wait_stream(current_stream)
\ No newline at end of file
diff --git a/bmtrain/optim/__init__.py b/bmtrain/optim/__init__.py
index 327d6676..15206328 100644
--- a/bmtrain/optim/__init__.py
+++ b/bmtrain/optim/__init__.py
@@ -1,3 +1,3 @@
 from .adam import AdamOptimizer
 from .adam_offload import AdamOffloadOptimizer
-from .loss_scaler import LossScaler
\ No newline at end of file
+from .optim_manager import OptimManager
\ No newline at end of file
diff --git a/bmtrain/optim/loss_scaler.py b/bmtrain/optim/optim_manager.py
similarity index 76%
rename from bmtrain/optim/loss_scaler.py
rename to bmtrain/optim/optim_manager.py
index 9fbaed37..e72fcddb 100644
--- a/bmtrain/optim/loss_scaler.py
+++ b/bmtrain/optim/optim_manager.py
@@ -27,20 +27,37 @@ def grad_rescale(param_groups, scale):
             if p.grad is not None and p.requires_grad:
                 p.grad /= scale
 
-class LossScaler:
-    """Loss scaler for mix-precision training
+class OptimManager:
+    """wait cuda stream. Optional: add loss scaler for mix-precision training
 
     Args:
-        loss_scale (float): The initial loss scale.
+        loss_scale (float): The initial loss scale. Default to None for not using loss scaling.
         loss_scale_factor (float): The loss scale factor.
         loss_scale_steps (int): The loss scale steps.
+
+    Examples:
+        >>> optim_manager = bmt.optim.OptimManager(loss_scale=1024)
+        >>> optim_manager.add_optimizer(optimizer1)
+        >>> optim_manager.add_optimizer(optimizer2, lr_scheduler2)
+        >>> for data in dataset:
+        >>>     # forward pass and calculate loss
+        >>>     optim_manager.zero_grad()
+        >>>     optim_manager.backward(loss)
+        >>>     optim_manager.clip_grad_norm(optimizer1.param_groups, max_norm=1.0, norm_type=2)
+        >>>     optim_manager.clip_grad_norm(optimizer2.param_groups, max_norm=2.0, norm_type=2)
+        >>>     optim_manager.step()
     """
     def __init__(self,
-        loss_scale : float = 1,
+        loss_scale : Optional[float] = None,
         loss_scale_factor : float = 2,
         loss_scale_steps : int = 1024,
     ):
-        self.loss_scale = loss_scale
+        if loss_scale is not None:
+            self.loss_scale = loss_scale
+            self.loss_scale_enabled = True
+        else:
+            self.loss_scale = 1
+            self.loss_scale_enabled = False
         self.steps_since_last_scale = 0
         self.loss_scale_factor = loss_scale_factor if loss_scale_factor > 1 else 1 / loss_scale_factor
         self.loss_scale_steps = loss_scale_steps
@@ -53,8 +70,8 @@ def add_optimizer(
         optimizer: torch.optim.Optimizer,
         lr_scheduler: Optional[WarmupLRScheduler] = None,
     ):
-        """Add optimizer and (optional) its corresponding lr_scheduler into loss_scaler.
-        All optimizers in the same loss_scaler share the same loss scale.
+        """Add optimizer and (optional) its corresponding lr_scheduler into optim_manager.
+        All optimizers in the same optim_manager share the same loss scale.
 
         Args:
             optim (torch.optim.Optimizer): A pytorch optimizer, e.g. torch.optim.Adam, torch.optim.SGD or bmtrain.optim.AdamOffloadOptimizer
@@ -63,14 +80,21 @@ def add_optimizer(
         self.optimizers.append(optimizer)
         self.lr_schedulers.append(lr_scheduler)
 
-    def __call__(self, loss : torch.Tensor) -> torch.Tensor:
+    def loss_scale(self, loss : torch.Tensor) -> torch.Tensor:
+        return loss * (self.loss_scale / config['world_size']) # loss scale
+
+    def backward(self, loss : torch.Tensor):
         """
         Backward with loss scale.
 
         Args:
             loss (torch.Tensor): loss
         """
-        return loss * (self.loss_scale / config['world_size'])
+        loss = self.loss_scale(loss)
+        loss.backward()
+        # some reduce ops of distributed parameter were launched on load stream
+        current_stream = torch.cuda.current_stream()
+        current_stream.wait_stream(config['load_stream'])
 
     def zero_grad(self):
         """
@@ -88,11 +112,7 @@ def step(self):
 
         This function can also handle gradient overflow by reducing the loss scale when it occurs.
         """
-        current_stream =  torch.cuda.current_stream()
-        # some reduce ops of distributed parameter were launched on load stream
-        current_stream.wait_stream(config['load_stream'])
-
-        if self.loss_scale > 1:
+        if self.loss_scale_enabled and self.loss_scale > 1:
             has_overflow = False
             for optimizer in self.optimizers:
                 try:
@@ -110,17 +130,20 @@ def step(self):
             if hasattr(optimizer, "_bmtrain_optimizer") and optimizer._bmtrain_optimizer:
                 optimizer.step(scale=self.loss_scale)
             else:
-                grad_rescale(optimizer.param_groups, self.loss_scale)
+                if self.loss_scale_enabled:
+                    grad_rescale(optimizer.param_groups, self.loss_scale)
                 optimizer.step()
 
             if lr_scheduler is not None:
                 lr_scheduler.step()
 
-        self.steps_since_last_scale += 1
+        if self.loss_scale_enabled:
+            self.steps_since_last_scale += 1
 
-        if self.steps_since_last_scale >= self.loss_scale_steps:
-            self._justify_scale(self.loss_scale * self.loss_scale_factor)
+            if self.steps_since_last_scale >= self.loss_scale_steps:
+                self._justify_scale(self.loss_scale * self.loss_scale_factor)
 
+        current_stream = torch.cuda.current_stream()
         config['load_stream'].wait_stream(current_stream)
 
     def clip_grad_norm(self, param_groups, max_norm, norm_type=2, eps=1e-6):
@@ -136,17 +159,8 @@ def clip_grad_norm(self, param_groups, max_norm, norm_type=2, eps=1e-6):
 
         Returns:
             Total norm of the parameters (viewed as a single vector).
-
-        Examples:
-            >>> optimizer = bmt.optim.AdamOffloadOptimizer(model.parameters())
-            >>> loss_scaler = bmt.optim.LossScaler()
-            >>> loss_scaler.add_optimizer(optimizer)
-            >>> # ...
-            >>> # backward_step()
-            >>> loss_scaler.clip_grad_norm(optimizer.param_groups, max_norm=1.0, norm_type=2)
-
         """
-        scale = self.loss_scale / config['world_size']
+        scale = self.loss_scale
         grads = []
         parameters = [p for group in param_groups for p in group['params']]
         for p in parameters:
diff --git a/example/train.py b/example/train.py
index 0f0454bb..7bc92400 100644
--- a/example/train.py
+++ b/example/train.py
@@ -54,8 +54,8 @@ def main():
     optimizer = bmt.optim.AdamOffloadOptimizer(model.parameters(), weight_decay=1e-2)
     lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0)
 
-    loss_scaler = bmt.optim.LossScaler(loss_scale=2**20)
-    loss_scaler.add_optimizer(optimizer, lr_scheduler)
+    optim_manager = bmt.optim.OptimManager(loss_scale=2**20)
+    optim_manager.add_optimizer(optimizer, lr_scheduler)
 
     bmt.synchronize()
     
@@ -65,7 +65,6 @@ def main():
     for iteration in range(1000):
         # load data
         st = time.time()
-        loss_scaler.zero_grad()
 
         with bmt.inspect.inspect_tensor() as inspector:
             pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
@@ -80,8 +79,9 @@ def main():
         
             global_loss = bmt.sum_loss(loss).item()
 
-            loss = loss_scaler(loss)
-            loss.backward()
+            optim_manager.zero_grad()
+
+            optim_manager.backward(loss)
         
         # print inspected tensors in the forward & backward pass
         # print parameters of the model
@@ -96,9 +96,8 @@ def main():
                     bmt.inspect.inspect_model(model, "*")
                 )
             )
-        
 
-        loss_scaler.step()
+        optim_manager.step()
 
         # record time and loss
         iteration_time = time.time() - st
@@ -113,7 +112,7 @@ def main():
                 global_loss,
                 avg_loss_recorder.value,
                 lr_scheduler.current_lr,
-                loss_scaler.loss_scale,
+                optim_manager.loss_scale,
                 avg_time_recorder.value
             )
         )

From 1cdafd939a854652a2cf22e08cf233b1164d43e6 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 2 Sep 2022 12:16:13 +0000
Subject: [PATCH 09/10] add auto training test

---
 tests/test_all.py      |   2 +
 tests/test_training.py | 439 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 441 insertions(+)
 create mode 100644 tests/test_training.py

diff --git a/tests/test_all.py b/tests/test_all.py
index af5a0253..7dd2c0b0 100644
--- a/tests/test_all.py
+++ b/tests/test_all.py
@@ -16,6 +16,8 @@
     ("model_wrapper", 4),
     ("send_recv", 4),
     ("nccl_backward", 4),
+
+    ("trainging", 4),
 ])
 
 for t, num_gpu in tq:
diff --git a/tests/test_training.py b/tests/test_training.py
new file mode 100644
index 00000000..c1aa8baf
--- /dev/null
+++ b/tests/test_training.py
@@ -0,0 +1,439 @@
+from bmtrain.optim import optim_manager
+from utils import *
+
+from typing import Optional
+import torch
+import math
+import torch.nn.functional as F
+import bmtrain as bmt
+from bmtrain.global_var import config
+import os
+
+class Attention(torch.nn.Module):
+    def __init__(self, 
+            dim_model : int, dim_head : int,
+            num_heads : int, bias : bool = True,
+            dtype = None
+        ) -> None:
+        super().__init__()
+
+        self.project_q = torch.nn.Linear(dim_model, dim_head * num_heads, bias=bias, dtype=dtype)
+        self.project_k = torch.nn.Linear(dim_model, dim_head * num_heads, bias=bias, dtype=dtype)
+        self.project_v = torch.nn.Linear(dim_model, dim_head * num_heads, bias=bias, dtype=dtype)
+
+        self.project_out = torch.nn.Linear(dim_head * num_heads, dim_model, bias=bias, dtype=dtype)
+
+        self.softmax = torch.nn.Softmax(dim=-1)
+        self.num_heads = num_heads
+        self.dim_head = dim_head
+        self.dim_model = dim_model
+    
+    def forward(self, 
+            hidden_q : torch.Tensor,        # (batch_size, seq_q, dim_model)
+            hidden_kv : torch.Tensor,       # (batch_size, seq_kv, dim_model)
+            mask : torch.BoolTensor,        # (batch_size, seq_q, seq_kv)
+            position_bias : Optional[torch.Tensor] = None,   # (batch, num_heads, seq_q, seq_kv)
+        ) -> torch.Tensor:
+        batch_size, seq_q, dim_model = hidden_q.size()
+        seq_kv = hidden_kv.size(1)
+
+        h_q : torch.Tensor = self.project_q(hidden_q)
+        h_k : torch.Tensor = self.project_k(hidden_kv)
+        h_v : torch.Tensor = self.project_v(hidden_kv)
+
+        h_q = h_q.view(batch_size, seq_q, self.num_heads, self.dim_head)
+        h_k = h_k.view(batch_size, seq_kv, self.num_heads, self.dim_head)
+        h_v = h_v.view(batch_size, seq_kv, self.num_heads, self.dim_head)
+
+        h_q = h_q.permute(0, 2, 1, 3).contiguous()
+        h_k = h_k.permute(0, 2, 1, 3).contiguous()
+        h_v = h_v.permute(0, 2, 1, 3).contiguous()
+
+        h_q = h_q.view(batch_size * self.num_heads, seq_q, self.dim_head)
+        h_k = h_k.view(batch_size * self.num_heads, seq_kv, self.dim_head)
+        h_v = h_v.view(batch_size * self.num_heads, seq_kv, self.dim_head)
+
+        score = torch.bmm(
+            h_q, h_k.transpose(1, 2)
+        )
+        score = score / math.sqrt(self.dim_head)
+
+        score = score.view(batch_size, self.num_heads, seq_q, seq_kv)
+
+        if position_bias is not None:
+            score = score + position_bias.view(batch_size, self.num_heads, seq_q, seq_kv)
+        
+        score = torch.where(
+            mask.view(batch_size, 1, seq_q, seq_kv),
+            score,
+            torch.scalar_tensor(float('-inf'), device=score.device, dtype=score.dtype)
+        )
+
+        score = torch.where(
+            mask.view(batch_size, 1, seq_q, seq_kv),
+            self.softmax(score),
+            torch.scalar_tensor(0, device=score.device, dtype=score.dtype)
+        )
+
+        score = score.view(batch_size * self.num_heads, seq_q, seq_kv)
+
+        h_out = torch.bmm(
+            score, h_v
+        )
+        h_out = h_out.view(batch_size, self.num_heads, seq_q, self.dim_head)
+        h_out = h_out.permute(0, 2, 1, 3).contiguous()
+        h_out = h_out.view(batch_size, seq_q, self.num_heads * self.dim_head)
+
+        attn_out = self.project_out(h_out)
+        return attn_out
+        
+class Feedforward(torch.nn.Module):
+    def __init__(self, dim_model : int, dim_ff : int, bias : bool = True, dtype = None) -> None:
+        super().__init__()
+
+        self.w_in = torch.nn.Linear(dim_model, dim_ff, bias = bias, dtype=dtype)
+        self.w_out = torch.nn.Linear(dim_ff, dim_model, bias = bias, dtype=dtype)
+
+        self.relu = torch.nn.ReLU()
+    
+    def forward(self, input : torch.Tensor) -> torch.Tensor:
+        return self.w_out(self.relu(self.w_in(input)))
+
+
+class TransformerEncoder(torch.nn.Module):
+    def __init__(self,
+            dim_model : int, dim_head : int, num_heads : int, dim_ff : int,
+            bias : bool = True, dtype = None
+        ) -> None:
+        super().__init__()
+
+        self.ln_attn = torch.nn.LayerNorm(dim_model, dtype=dtype)
+        self.attn = Attention(dim_model, dim_head, num_heads, bias=bias, dtype=dtype)
+
+        self.ln_ff = torch.nn.LayerNorm(dim_model, dtype=dtype)
+        self.ff = Feedforward(dim_model, dim_ff, bias=bias, dtype=dtype)
+    
+    def forward(self,
+            hidden : torch.Tensor,      # (batch, seq_len, dim_model)
+            mask : torch.BoolTensor,    # (batch, seq_len, dim_model)
+            position_bias : Optional[torch.Tensor] = None,   # (batch, num_head, seq_len, seq_len)
+        ):
+        x = self.ln_attn(hidden)
+        x = self.attn(x, x, mask, position_bias)
+        hidden = hidden + x
+
+        x = self.ln_ff(hidden)
+        x = self.ff(x)
+        hidden = hidden + x
+
+        return hidden
+    
+
+class GPT(torch.nn.Module):
+    def __init__(self,
+            num_layers : int, vocab_size : int,
+            dim_model : int, dim_head : int, num_heads : int, dim_ff : int,
+            max_distance : int,
+            bias : bool = True, dtype = None
+        ) -> None:
+        super().__init__()
+
+        self.dtype = dtype
+        self.max_distance = max_distance
+
+        self.word_emb = torch.nn.Embedding(vocab_size, dim_model, dtype=dtype)
+        self.pos_emb = torch.nn.Embedding(max_distance, dim_model, dtype=dtype)
+        self.dim_model = dim_model
+        
+        self.transformers = torch.nn.ModuleList([
+            TransformerEncoder(
+                dim_model, dim_head, num_heads, dim_ff, bias, dtype
+            )
+            for _ in range(num_layers)
+        ])
+
+        self.layernorm = torch.nn.LayerNorm(dim_model, dtype=dtype)
+
+    def forward(self,
+            input : torch.LongTensor,   # (batch, seq_len)
+            pos : torch.LongTensor,     # (batch, seq_len)
+            mask : torch.BoolTensor,    # (batch, seq_len)
+        ) -> torch.Tensor:
+
+        mask_2d = mask[:, None, :] & mask[:, :, None]   # (batch, seq_len, seq_len)
+        mask_2d = mask_2d & (pos[:, None, :] >= pos[:, :, None])
+
+        input_emb = self.pos_emb(pos) + self.word_emb(input)
+
+        out = input_emb
+        if isinstance(self.transformers, torch.nn.ModuleList):
+            for layer in self.transformers:
+                out = layer(out, mask_2d, None)
+        else:
+            out = self.transformers(out, mask_2d, None)
+        out = self.layernorm(out)
+
+        logits = F.linear(out, self.word_emb.weight) / math.sqrt(self.dim_model)
+
+        return logits
+
+def sub_train_torch(model, loss_func_cls, optimizer_cls):
+    loss_func = loss_func_cls(ignore_index=-100)
+    optimizer = optimizer_cls(model.parameters(), weight_decay=1e-2)
+    lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0)
+
+    optim_manager = bmt.optim.OptimManager(loss_scale=2**20 if model.dtype == torch.half else None)
+    optim_manager.add_optimizer(optimizer, lr_scheduler)
+
+    # use the break if i == bmt.rank() to generate different data on different rank
+    torch.manual_seed(2333)
+    batch_size = 2
+    seq_len = 512
+
+    sents = []
+    enc_lengths = []
+    enc_inputs = []
+    targetss = []
+    masks = []
+    for i in range(bmt.world_size()):
+        sent = torch.randint(0, 10240, (batch_size, seq_len + 1))
+        enc_length = torch.randint(128, seq_len, (batch_size,)).long().cuda()
+        enc_input = sent[:, :-1].long().cuda()
+        targets = sent[:, 1:].long().cuda()
+        mask = torch.arange(seq_len).long().cuda()[None, :] < enc_length[:, None]
+        targets = torch.where(
+            mask,
+            targets,
+            torch.full_like(targets, -100, dtype=torch.long)
+        )
+
+        sents.append(sent)
+        enc_lengths.append(enc_length)
+        enc_inputs.append(enc_input)
+        targetss.append(targets)
+        masks.append(mask)
+
+    sent = torch.cat(sents, dim=0)
+    enc_length = torch.cat(enc_lengths, dim=0)
+    enc_input = torch.cat(enc_inputs, dim=0)
+    targets = torch.cat(targetss, dim=0)
+    mask = torch.cat(masks, dim=0)
+
+    pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
+
+    logs = []
+    for iter in range(100):
+        optim_manager.zero_grad()
+
+        logits = model(enc_input, pos, pos < enc_length[:, None])
+
+        batch, seq_len, vocab_out_size = logits.size()
+
+        loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
+    
+        global_loss = loss.item()
+
+        loss = optim_manager.loss_scale * loss
+        loss.backward()
+
+        optim_manager.step()
+
+        bmt.print_rank("| Iter: {:6d} | loss: {:.4f} {:.4f} | lr: {:.4e} scale: {:10.4f}".format(
+            iter,
+            global_loss,
+            loss,
+            lr_scheduler.current_lr,
+            optim_manager.loss_scale,
+        ))
+        logs.append(global_loss)
+
+    summary = bmt.inspect.inspect_model(model, "*")
+    return logs, summary
+
+def sub_train(model, loss_func_cls, optimizer_cls):
+    loss_func = loss_func_cls(ignore_index=-100)
+    optimizer = optimizer_cls(model.parameters(), weight_decay=1e-2)
+    lr_scheduler = bmt.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0)
+
+    optim_manager = bmt.optim.OptimManager(loss_scale=2**20 if model.dtype == torch.half else None)
+    optim_manager.add_optimizer(optimizer, lr_scheduler)
+
+    # use the break if i == bmt.rank() to generate different data on different rank
+    torch.manual_seed(2333)
+    batch_size = 2
+    seq_len = 512
+
+    for i in range(bmt.world_size()):
+        sent = torch.randint(0, 10240, (batch_size, seq_len + 1))
+        enc_length = torch.randint(128, seq_len, (batch_size,)).long().cuda()
+        enc_input = sent[:, :-1].long().cuda()
+        targets = sent[:, 1:].long().cuda()
+        mask = torch.arange(seq_len).long().cuda()[None, :] < enc_length[:, None]
+        targets = torch.where(
+            mask,
+            targets,
+            torch.full_like(targets, -100, dtype=torch.long)
+        )
+
+        if i == bmt.rank():
+            break
+
+    pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)
+
+    logs = []
+    for iter in range(100):
+        optim_manager.zero_grad()
+
+        logits = model(enc_input, pos, pos < enc_length[:, None])
+
+        batch, seq_len, vocab_out_size = logits.size()
+
+        loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))
+    
+        global_loss = bmt.sum_loss(loss).item()
+
+        optim_manager.backward(loss)
+
+        optim_manager.step()
+
+        bmt.print_rank("| Iter: {:6d} | loss: {:.4f} {:.4f} | lr: {:.4e} scale: {:10.4f}".format(
+            iter,
+            global_loss,
+            loss,
+            lr_scheduler.current_lr,
+            optim_manager.loss_scale,
+        ))
+        logs.append(global_loss)
+
+    summary = bmt.inspect.inspect_model(model, "*")
+    return logs, summary
+    
+def train(model, loss_func, optimizer):
+    key = f"{model[0]}*{loss_func[0]}*{optimizer[0]}"
+    model = model[1]()
+    if key.startswith("torch"):
+        ret = sub_train_torch(model, loss_func[1], optimizer[1])
+    else:
+        ret = sub_train(model, loss_func[1], optimizer[1])
+    del model
+    bmt.print_rank(f"finished {key}")
+    return key, ret
+
+def test_main(test_fp16=True, test_fp32=True):
+    ckpt_path = "test_ckpt.pt"
+
+    kwargs = {
+        "num_layers": 8,
+        "vocab_size": 10240, 
+        "dim_model": 2560,
+        "dim_head": 80,
+        "num_heads": 32,
+        "dim_ff": 8192,
+        "max_distance": 1024,
+        "bias": True,
+        "dtype": None,
+    }
+
+    def make_ref_ckpt():
+        model = GPT(**kwargs)
+        if bmt.rank() == 0:
+            torch.save(model.state_dict(), ckpt_path)
+        bmt.synchronize()
+        del model
+
+    ret = {}
+    def torch_model():
+        model = GPT(**kwargs)
+        model.load_state_dict(torch.load(ckpt_path))
+        model = model.cuda()
+        return model
+
+    def wrap_model():
+        model = GPT(**kwargs)
+        wrap_model = bmt.BMTrainModelWrapper(model)
+        bmt.load(wrap_model, ckpt_path)
+        return model
+
+    def list_model():
+        model = GPT(**kwargs)
+        list_model = bmt.BMTrainModelWrapper(model)
+        list_model.transformers = bmt.TransformerBlockList([m for m in list_model.transformers])
+        bmt.load(list_model, ckpt_path)
+        return model
+
+    def pipe_model():
+        model = GPT(**kwargs)
+        pipe_model = bmt.BMTrainModelWrapper(model)
+        for m in pipe_model.transformers:
+            assert isinstance(m, bmt.CheckpointBlock)
+        pipe_model.transformers = bmt.PipelineTransformerBlockList([m for m in pipe_model.transformers])
+        bmt.load(pipe_model, ckpt_path)
+        return model
+
+    models = {
+        "torch": torch_model,
+        "wrapper": wrap_model,
+        "blocklist": list_model,
+        "pipelist": pipe_model,
+    }
+    loss_funcs = {
+        "bmt_entropy": bmt.loss.FusedCrossEntropy,
+        "torch_entropy": torch.nn.CrossEntropyLoss,
+    }
+    optimizers = {
+        "bmt_adam": bmt.optim.AdamOptimizer,
+        "bmt_adam_offload": bmt.optim.AdamOffloadOptimizer,
+        "torch_adam": torch.optim.Adam,
+    }
+
+    ret = {}
+    def add_to_check_list(m, l, o):
+        key, value = train((m, models[m]), (l, loss_funcs[l]), (o, optimizers[o]))
+        ret[key] = value
+
+    if test_fp16:
+        kwargs["dtype"] = torch.half
+        make_ref_ckpt()
+        add_to_check_list("torch", "bmt_entropy", "bmt_adam")
+        add_to_check_list("wrapper", "bmt_entropy", "bmt_adam")
+        add_to_check_list("blocklist", "bmt_entropy", "bmt_adam")
+        add_to_check_list("pipelist", "bmt_entropy", "bmt_adam")
+        add_to_check_list("blocklist", "torch_entropy", "bmt_adam")
+        add_to_check_list("blocklist", "bmt_entropy", "bmt_adam_offload")
+        if bmt.rank() == 0:
+            os.remove(ckpt_path)
+        check(ret)
+
+    if test_fp32:
+        kwargs["dtype"] = torch.float
+        make_ref_ckpt()
+        add_to_check_list("torch", "torch_entropy", "bmt_adam")
+        add_to_check_list("wrapper", "torch_entropy", "bmt_adam")
+        add_to_check_list("blocklist", "torch_entropy", "bmt_adam")
+        add_to_check_list("pipelist", "torch_entropy", "bmt_adam")
+        add_to_check_list("blocklist", "torch_entropy", "bmt_adam_offload")
+        add_to_check_list("blocklist", "torch_entropy", "torch_adam")
+        if bmt.rank() == 0:
+            os.remove(ckpt_path)
+        check(ret)
+
+def check(ret):
+    if bmt.rank() == 0:
+        for k1, v1 in ret.items():
+            for k2, v2 in ret.items():
+                print(f"checking {k1} vs. {k2}")
+                check_param(v1[1], v2[1])
+    bmt.synchronize()
+    ret.clear()
+
+def check_param(info1, info2):
+    for l1, l2 in zip(info1, info2):
+        for key in ['std', 'mean', 'max', 'min']:
+            v1 = l1[key]
+            v2 = l2[key]
+            assert_lt(abs(v1-v2), 1e-2)
+
+if __name__ == '__main__':
+    bmt.init_distributed(pipe_size=2)
+
+    test_main(test_fp16=True, test_fp32=True)
\ No newline at end of file

From 86bb536912d3c25e921b5af3dbc93249525724ac Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 2 Sep 2022 13:41:23 +0000
Subject: [PATCH 10/10] add auto clip_grad_norm test

---
 tests/test_training.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/test_training.py b/tests/test_training.py
index c1aa8baf..7342fe6c 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -236,14 +236,17 @@ def sub_train_torch(model, loss_func_cls, optimizer_cls):
         loss = optim_manager.loss_scale * loss
         loss.backward()
 
+        grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=10.0)
+
         optim_manager.step()
 
-        bmt.print_rank("| Iter: {:6d} | loss: {:.4f} {:.4f} | lr: {:.4e} scale: {:10.4f}".format(
+        bmt.print_rank("| Iter: {:6d} | loss: {:.4f} {:.4f} | lr: {:.4e} scale: {:10.4f} | grad_norm: {:.4f} |".format(
             iter,
             global_loss,
             loss,
             lr_scheduler.current_lr,
             optim_manager.loss_scale,
+            grad_norm,
         ))
         logs.append(global_loss)
 
@@ -294,14 +297,17 @@ def sub_train(model, loss_func_cls, optimizer_cls):
 
         optim_manager.backward(loss)
 
+        grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=10.0)
+
         optim_manager.step()
 
-        bmt.print_rank("| Iter: {:6d} | loss: {:.4f} {:.4f} | lr: {:.4e} scale: {:10.4f}".format(
+        bmt.print_rank("| Iter: {:6d} | loss: {:.4f} {:.4f} | lr: {:.4e} scale: {:10.4f} | grad_norm: {:.4f} |".format(
             iter,
             global_loss,
             loss,
             lr_scheduler.current_lr,
             optim_manager.loss_scale,
+            grad_norm,
         ))
         logs.append(global_loss)
 
