add test for grad accumulation and state_dict interface.

bmtrain/init.py

@@ -39,11 +39,13 @@ def init_distributed(
     """
     torch.backends.cudnn.enabled = False
 
-    local_rank = int(os.environ["LOCAL_RANK"])
-    rank = int(os.environ["RANK"])
-    world_size = int(os.environ["WORLD_SIZE"])
-    local_size = int(os.environ["LOCAL_WORLD_SIZE"])
-    master = os.environ["MASTER_ADDR"] + ":" + os.environ["MASTER_PORT"]
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    local_size = int(os.environ.get("LOCAL_WORLD_SIZE","1"))
+    addr = os.environ.get("MASTER_ADDR", "localhost")
+    port = os.environ.get("MASTER_PORT", "10010")
+    master = addr+":"+port
     timeout = datetime.timedelta(seconds=1800)
     rendezvous_iterator = dist.rendezvous(
         init_method, rank, world_size, timeout=timeout

tests/test.sh

@@ -0,0 +1,4 @@
+export NCCL_P2P_DISABLE=1
+export CUDA_LAUNCH_BLOCKING
+torchrun --nproc 4  test_pipe.py
+# torchrun --nproc 4 test_send_recv.py

tests/test_grad_accu.py

@@ -0,0 +1,80 @@
+import bmtrain as bmt
+import torch
+from bmtrain import config
+from bmtrain.block_layer import CheckpointBlockContext,  CheckpointBlock, TransformerBlockList
+from bmtrain.pipe_layer import PipelineTransformerBlockList
+from typing import List
+import torch.nn.functional as F
+def print_rank0(s):
+    if bmt.rank() == 0:
+        print(s)
+class Linear(bmt.DistributedModule):
+    def __init__(self, in_features : int, out_features: int, init_weight = None, init_bias = None) -> None:
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.out = {}
+        if init_weight:
+            self.weight = bmt.DistributedParameter(torch.tensor(init_weight, dtype=torch.float, device="cuda").reshape(out_features, in_features))
+        else:
+            self.weight = bmt.DistributedParameter(torch.empty(out_features, in_features, dtype=torch.float, device="cuda"), init_method=torch.nn.init.xavier_normal_)
+
+        if init_bias:
+            self.bias = bmt.DistributedParameter(torch.tensor(init_bias, dtype=torch.float, device="cuda").reshape(out_features,))
+        else:
+            self.bias = bmt.DistributedParameter(torch.empty(out_features, dtype=torch.float, device="cuda"), init_method=torch.nn.init.zeros_)
+
+    def forward(self, input):
+        ret = F.linear(input, self.weight, self.bias)
+        return ret
+
+def test_grad_accu():
+    # normal distribute module
+    m = Linear(256, 256)
+    inp = torch.randn((1, 10, 256), device="cuda")
+    logits = m(inp)
+    loss = logits.sum()
+    loss.backward()
+    grad1 = m._parameters["weight"].grad.clone()
+    logits = m(inp)
+    loss = logits.sum()
+    loss.backward()
+    grad2 = m._parameters["weight"].grad
+    assert torch.allclose(grad1*2, grad2)
+    print_rank0("grad accumulation for distribute module passed")
+    # checkpoint block
+    m = CheckpointBlock(Linear(256, 256))
+    inp = torch.randn((1, 10, 256), device="cuda")
+    logits = m(inp)
+    loss = logits.sum()
+    loss.backward()
+    bmt.synchronize()
+    grad1 = m.weight.grad.clone()
+    logits = m(inp)
+    loss = logits.sum()
+    loss.backward()
+    bmt.synchronize()
+    grad2 = m.weight.grad.clone()
+    assert torch.allclose(grad1*2, grad2)
+    print_rank0("grad accumulation for checkpointblock passed")
+    # transformer block list
+    m = TransformerBlockList([CheckpointBlock(Linear(256, 256))])
+    inp = torch.randn((1, 10, 256), device="cuda")
+    logits = m(inp)
+    loss = logits.sum()
+    loss.backward()
+    bmt.synchronize()
+    grad1 = m[0].weight.grad.clone()
+    logits = m(inp)
+    loss = logits.sum()
+    loss.backward()
+    bmt.synchronize()
+    grad2 = m[0].weight.grad
+    assert torch.allclose(grad1*2, grad2)
+    print_rank0("grad accumulation for TransformerBlockList passed")
+
+
+if __name__ == "__main__":
+    bmt.init_distributed()
+    test_grad_accu()

tests/test_load_ckpt.py

@@ -10,10 +10,10 @@ def __init__(self, in_features : int, out_features: int, bias: bool = True, dtyp
 
         self.in_features = in_features
         self.out_features = out_features
-        self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype))
+        self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=dtype, device="cuda"))
         torch.nn.init.xavier_normal_(self.weight)
         if bias:
-            self.bias = torch.nn.Parameter(torch.empty(out_features, dtype=dtype))
+            self.bias = torch.nn.Parameter(torch.empty(out_features, dtype=dtype, device="cuda"))
             torch.nn.init.zeros_(self.bias)
         else:
             self.register_parameter('bias', None)
@@ -35,30 +35,42 @@ def __init__(self, in_features : int, out_features: int, bias: bool = True, dtyp
 
     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))
+def test_main():
+    ckpt_path = "test_ckpt.pt"
+    # Transformer BlockList
+    m = Linear_Normal(256, 256).cuda()
+    m2 = bmt.TransformerBlockList([bmt.CheckpointBlock(Linear_BMT(256, 256))])
+    if bmt.rank() == 0:
+        torch.save(m.state_dict(), ckpt_path)
+    dic2 = m.state_dict()
+    dic2["0.weight"] = dic2.pop("weight")
+    dic2["0.bias"] = dic2.pop("bias")
+    m2.load_state_dict(dic2)
+    for key in m.state_dict():
+        bmt_key = f"0.{key}"
+        assert bmt_key in m2.state_dict(), "wrong key in bmtrain model"
+        assert (m2.state_dict()[bmt_key].cuda() == m.state_dict()[key]).all() , "wrong param in bmtrain model"
+    if bmt.rank() == 0:
+        os.remove(ckpt_path)
+    print("Transformer Blocklist load_state_dict and state_dict test passed")
 
-    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)
+    # CheckpointBlock 
+    m3 = bmt.CheckpointBlock(Linear_BMT(256, 256))
+    m3.load_state_dict(m.state_dict())
+    for key in m.state_dict():
+        assert key in m3.state_dict(), "wrong key in bmtrain model"
+        assert (m.state_dict()[key] == m3.state_dict()[key].cuda()).all(), "wrong param in bmtrain model"
+    print("CheckpointBlock load_state_dict and state_dict test passed")
 
-    os.remove(ckpt_path)
+    # normal Distributed module
+    m4 = Linear_BMT(256, 256)
+    m4.load_state_dict(m.state_dict())
+    for key in m.state_dict():
+        assert key in m4.state_dict(), "wrong key in bmtrain model"
+        assert (m.state_dict()[key] == m4.state_dict()[key].cuda()).all(), "wrong param in bmtrain model"
+    print("bmt.distributedmodule load_state_dict and state_dict test passed")
 
 if __name__ == "__main__":
     bmt.init_distributed()

tests/test_requires_grad.py

@@ -33,7 +33,7 @@ def run(m, a, b):
     logits = m(inp)
     loss = logits.sum()
     loss.backward()
-
+    bmt.synchronize()
     sm = bmt.inspect.format_summary(
             bmt.inspect.inspect_model(m, '*')
         )