assignment-distributed-training/mpi-test.py at main · mlsyscourse/assignment-distributed-training · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
from mpi4py import MPI
from mpi_wrapper import Communicator
import numpy as np
import argparse

parser = argparse.ArgumentParser()
parser.add_argument(
    "--test_case",
    type=str,
    help="MPI names for different toy examples",
    default="",
    choices=["allreduce", "allgather", "reduce_scatter", "split"],
)

if __name__ == "__main__":
    comm = MPI.COMM_WORLD
    comm = Communicator(comm)

    nprocs = comm.Get_size()
    rank = comm.Get_rank()

    if nprocs != 8:
        if rank == 0:
            print(
                "This test requires exactly 8 processes. You can modify the "
                "code in this file({__file__}) to "
                "change the number of processes and the test cases."
            )
        exit(1)

    args = parser.parse_args()

    if args.test_case == "allreduce":
        """
        Allreduce example
        """

        r = np.random.randint(0, 100, 10)
        rr = np.empty(10, dtype=int)

        print("Rank " + str(rank) + ": " + str(r))

        comm.Barrier()
        comm.Allreduce(r, rr, op=MPI.SUM)

        if rank == 0:
            print("Allreduce: " + str(rr))

    elif args.test_case == "allgather":
        """
        Allgather example
        """

        r = np.random.randint(0, 100, 2)
        rr = np.empty(16, dtype=int)

        print("Rank " + str(rank) + ": " + str(r))

        comm.Barrier()
        comm.Allgather(r, rr)

        if rank == 0:
            print("Allgather: " + str(rr))

    elif args.test_case == "reduce_scatter":
        """
        Reduce_scatter example
        """

        r = np.random.randint(0, 100, 16)
        rr = np.empty(2, dtype=int)

        print("Rank " + str(rank) + ": " + str(r))

        comm.Barrier()
        comm.Reduce_scatter(r, rr, op=MPI.SUM)

        print("Rank " + str(rank) + " After Reduce_scatter: " + str(rr))

    elif args.test_case == "split":
        """
        Split example (group-wise reduce)
        split into 4 groups based on the modulo operation:
         Group 0: (0, 4)
         Group 1: (1, 5)
         Group 2: (2, 6)
         Group 3: (3, 7)
        """

        r = np.random.randint(0, 100, 10)
        rr = np.empty(10, dtype=int)

        print("Rank " + str(rank) + ": " + str(r))

        key = rank
        color = rank % 4

        group_comm = comm.Split(key=key, color=color)

        group_comm.Barrier()
        group_comm.Allreduce(r, rr, op=MPI.SUM)

        print("Rank " + str(rank) + " After split and Allreduce: " + str(rr))

    else:
        print(f"This is rank {rank}.")