1234567891011121314151617181920212223242526272829303132333435363738394041424344454647 |
- from benchmarks.communication.utils import *
- from benchmarks.communication.all_reduce import run_all_reduce
- from benchmarks.communication.all_gather import run_all_gather
- from benchmarks.communication.all_to_all import run_all_to_all
- from benchmarks.communication.pt2pt import run_pt2pt
- from benchmarks.communication.broadcast import run_broadcast
- from benchmarks.communication.constants import *
- # For importing
- def main(args, rank):
- init_processes(local_rank=rank, args=args)
- ops_to_run = []
- if args.all_reduce:
- ops_to_run.append('all_reduce')
- if args.all_gather:
- ops_to_run.append('all_gather')
- if args.broadcast:
- ops_to_run.append('broadcast')
- if args.pt2pt:
- ops_to_run.append('pt2pt')
- if args.all_to_all:
- ops_to_run.append('all_to_all')
- if len(ops_to_run) == 0:
- ops_to_run = ['all_reduce', 'all_gather', 'all_to_all', 'broadcast', 'pt2pt']
- for comm_op in ops_to_run:
- if comm_op == 'all_reduce':
- run_all_reduce(local_rank=rank, args=args)
- if comm_op == 'all_gather':
- run_all_gather(local_rank=rank, args=args)
- if comm_op == 'all_to_all':
- run_all_to_all(local_rank=rank, args=args)
- if comm_op == 'pt2pt':
- run_pt2pt(local_rank=rank, args=args)
- if comm_op == 'broadcast':
- run_broadcast(local_rank=rank, args=args)
- # For directly calling benchmark
- if __name__ == "__main__":
- args = benchmark_parser().parse_args()
- rank = args.local_rank
- main(args, rank)
|