diff --git a/bmtrain/init.py b/bmtrain/init.py index b0632d5c..43b6ca90 100644 --- a/bmtrain/init.py +++ b/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 diff --git a/tests/test.sh b/tests/test.sh new file mode 100644 index 00000000..c20dcc89 --- /dev/null +++ b/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 diff --git a/tests/test_grad_accu.py b/tests/test_grad_accu.py new file mode 100644 index 00000000..48dbe8ac --- /dev/null +++ b/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() \ No newline at end of file diff --git a/tests/test_load_ckpt.py b/tests/test_load_ckpt.py index 56298940..632f0e8e 100644 --- a/tests/test_load_ckpt.py +++ b/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() diff --git a/tests/test_requires_grad.py b/tests/test_requires_grad.py index 707bd793..abcc998c 100644 --- a/tests/test_requires_grad.py +++ b/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, '*') )