def checkArgs(self, args):
super().checkArgs(args)
if args.collective not in supportedCollectives:
print(
"\t ERROR: Specified collective: %s is not one of the supported collectives: %s. Make sure the input is using the correct case."
% (args.collective, str(supportedCollectives))
)
comms_utils.gracefulExit()
if args.z not in supportedCommStyle:
print(
"\t ERROR: Specified blocking: %d is not one of the supported commstyle: %s"
% (args.z, str(supportedCommStyle))
)
comms_utils.gracefulExit()
args.b = comms_utils.parsesize(args.b)
args.e = comms_utils.parsesize(args.e)
args.dtype = self.dtypeMap[args.data_type]
if args.b < 1:
print("\t Starting size: %d should atleast be 1! " % (args.b))
args.b = 1
if args.e < args.b:
print(
"\t ERROR: In COMMS-mode, the begin-size: %d is larger than the end-size: %d "
% (args.b, args.e)
)
if args.device == "cpu" and args.backend == "nccl":
raise ValueError("NCCL is not supported for device type CPU")
def checkArgs(self, args):
super().checkArgs(args)
if args.pt2pt is not None:
args.collective = "pt2pt"
if args.pt2pt not in pt2ptPatterns:
print(
"\t ERROR: Specified pt2pt pattern: %d is not one of the supported pt2pt patterns: %s"
% (args.pt2pt, str(pt2ptPatterns)))
comms_utils.gracefulExit()
args.b = comms_utils.parsesize(args.b)
args.e = comms_utils.parsesize(args.e)
args.dtype = self.dtypeMap[args.data_type]
if args.b < 1:
print("\t Starting size: %d should atleast be 1! " % (args.b))
args.b = 1
if args.e < args.b:
print(
"\t ERROR: In COMMS-mode, the begin-size: %d is larger than the end-size: %d "
% (args.b, args.e))
if args.device == "cpu" and args.backend == "nccl":
raise ValueError("NCCL is not supported for device type CPU")
if args.c == 1 and args.z == 0:
logging.warning(
"Data validation is not supported for non-blocking mode...disable validation check and proceed..."
)
args.c = 0
def checkArgs(self, args):
super().checkArgs(args)
if args.pt2pt is not None:
args.collective = "pt2pt"
if args.pt2pt not in pt2ptPatterns:
logger.error(
f"Specified pt2pt pattern: {args.pt2pt} is not one of the supported pt2pt patterns: {str(pt2ptPatterns)}"
)
comms_utils.gracefulExit()
args.b = comms_utils.parsesize(args.b)
args.e = comms_utils.parsesize(args.e)
args.dtype = self.dtypeMap[args.data_type]
if args.b < 1:
logger.warning(
f"Starting size (--b {args.b}) should be greater than 1 byte...fix and continue"
)
args.b = 1
if args.e < args.b:
logger.warning(
f"the begin-size (--b {args.b}) is larger than the end-size (--e {args.e})"
)
if args.device == "cpu" and args.backend == "nccl":
raise ValueError(
f"NCCL is not supported for device type {args.device}")
if args.c == 1 and args.z == 0 and args.collective in (
"all_reduce", "reduce", "reduce_scatter"):
logger.warning(
f"Data validation is not supported for {args.collective} in non-blocking mode, disabled and continue"
)
args.c = 0
# run a few sanity checks
if args.bitwidth < 32:
if args.device != "cuda":
logger.error(
f"collective quantization may not be fully supported for {args.device}"
)
comms_utils.checkQuantArgs(
args.collective,
args.dtype,
args.b,
args.quant_a2a_embedding_dim,
args.z,
)
def main():
### import packages ###
import sys
import argparse
supportedCommStyle = [0, 1] # 0 : non-blocking, 1 : blocking.
supportedCollectives = [
"reduce",
"all_reduce",
"all_to_all",
"all_to_allv",
] # , "scatter", "gather", "all_gather", "broadcast", "all_to_all"]
supportedNwstacks = ["pytorch-nccl", "pytorch-xla-tpu"]
supported_tpu_core_valuses = [1, 8]
dtypeMap = {
"float32": torch.float32,
"int32": torch.int32,
"float16": torch.half,
"float64": torch.double,
}
supportedDtype = list(dtypeMap.keys())
### parse arguments ###
parser = argparse.ArgumentParser(
description="PARAM-Comm Benchmark",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
# experiment related parameters
parser.add_argument(
"--backend",
type=str,
default="nccl",
help="The backend to be used in PyTorch distributed process group"
) # alternative is DLRM mode.
parser.add_argument(
"--mode", type=str, default="comms",
help="benchmark mode") # alternative is DLRM mode or comm-compute mode
parser.add_argument("--b",
type=str,
default="8",
help="minimum size, in bytes, to start with"
) # COMMS mode, begin the sweep at.
parser.add_argument("--e",
type=str,
default="64",
help="maximum size, in bytes, to end at"
) # COMMS mode, end the sweep at.
parser.add_argument("--f",
type=int,
default=2,
help="multiplication factor between sizes"
) # COMMS mode, multiplication factor.
parser.add_argument("--z",
type=int,
default=1,
help="use blocking mode for collectives"
) # 'sync/blocking' : 1 , 'async/non-blocking' : 0
parser.add_argument(
"--w", type=int, default=5,
help="number of warmup iterations") # number of warmup-iterations
parser.add_argument("--n",
type=int,
default=5,
help="number of iterations") # number of iterations
parser.add_argument(
"--collective",
type=str,
default="all_reduce",
help='Collective to benchmark, supports ' +
str(supportedCollectives)) # collective op to benchmark
parser.add_argument(
"--master-ip",
type=str,
default="127.0.0.1",
help="The master-IP to coordinate") # The master-IP to coordinate.
parser.add_argument(
"--master-port",
type=str,
default="29500",
help="The master-port to coordinate") # The master-port to coordinate.
parser.add_argument(
"--nw-stack",
type=str,
default="pytorch-nccl",
help="network stack to be used, supports " +
str(supportedNwstacks)) # The network stack to profile.
parser.add_argument(
"--dtype", type=torch.dtype, default=torch.float32
) # will be overwritten based on args.data_type and dtypeMap.
parser.add_argument("--data-type",
type=str,
default="float32",
help="the base data type, supports " +
str(supportedDtype)) # The data type
parser.add_argument(
"--num-tpu-cores",
type=int,
default=1,
help="number of TPU cores to be used") # number of TPU cores
# For comm-compute or compute mode
parser.add_argument("--kernel",
type=str,
default="gemm",
help="compute kernel") # Compute kernel: "gemm"
parser.add_argument(
"--num-compute",
type=int,
default=100,
help="one collective for every NUM_COMPUTE compute kernels"
) # Launch one coll for every n compute kernels
# For GEMM
parser.add_argument("--mm-dim",
type=int,
default=100,
help="dimension size for GEMM compute kernel"
) # Matrix multiplication dim n, A[n,n] * B [n,n]
# For emb lookup
parser.add_argument(
"--emb-dim",
type=int,
default=128,
help="dimension size for Embedding table compute kernel"
) # Embedding table dimension
parser.add_argument(
"--num-embs",
type=int,
default=100000,
help="Embedding table hash size for Embedding table compute kernel"
) # Embedding table hash size
parser.add_argument("--avg-len",
type=int,
default=28,
help="Average lookup operations per sample"
) # Average #lookup per sample
parser.add_argument("--batch-size",
type=int,
default=512,
help="number of samples reading the table concurrently"
) # #Samples reading the table concurrently
parser.add_argument(
"--root",
type=int,
default=0,
help="root process for reduce benchmark"
) # root process for reduce (and gather, scatter, bcast, etc., if support in the future)
# TODO: check the correctness of root, should be between 0 to [world_size -1]
args, leftovers = parser.parse_known_args()
args.b = comms_utils.parsesize(args.b)
args.e = comms_utils.parsesize(args.e)
if args.nw_stack not in supportedNwstacks:
print(
"\t ERROR: Specified backend: %s is not one of the supported backends: %s. Make sure the input is using the correct case."
% (args.nw_stack, str(supportedNwstacks)))
sys.exit(
) # WARNING: Assuming sys is always used, should find a platform-independent way to gracefully exit.
if args.collective not in supportedCollectives:
print(
"\t ERROR: Specified collective: %s is not one of the supported collectives: %s. Make sure the input is using the correct case."
% (args.collective, str(supportedCollectives)))
sys.exit(
) # WARNING: Assuming sys is always used, should find a platform-independent way to gracefully exit.
if args.z not in supportedCommStyle:
print(
"\t ERROR: Specified blocking: %d is not one of the supported commstyle: %s"
% (args.z, str(supportedCommStyle)))
comms_utils.gracefulExit()
if args.data_type not in supportedDtype:
print(
"\t ERROR: Specified dtype: %d is not one of the supported commstyle: %s"
% (args.data_type, str(supportedDtype)))
comms_utils.gracefulExit()
args.dtype = dtypeMap[args.data_type]
mpi_env_params = comms_utils.read_mpi_env_vars()
if mpi_env_params["global_rank"] == 0:
print("\t MPI environment: %s " % (str(mpi_env_params)))
print(
"\t backend: %s nw-stack: %s mode: %s args.b: %d args.e: %d args.f: %d args.z: %s args.master_ip: %s "
% (
args.backend,
args.nw_stack,
args.mode,
args.b,
args.e,
args.f,
args.z,
args.master_ip,
))
if args.num_tpu_cores not in supported_tpu_core_valuses:
print(
"\t ERROR: TPU core value: %d is not one of the supported values: %s "
% (args.num_tpu_cores, supported_tpu_core_valuses))
comms_utils.gracefulExit()
if args.b < 1:
print("\t Starting size: %d should atleast be 1! " % (args.b))
args.b = 1
element_size = torch.ones([1], dtype=args.dtype).element_size()
comms_world_info = comms_utils.comms_world_info_holder(
args.master_ip, args.master_port, args.num_tpu_cores, mpi_env_params)
commsParams = comms_utils.commsParamsHolder(args, element_size, benchTime)
runComms(comms_world_info, commsParams)
def test_single_size(self):
sizeStr = "1024"
size = comms_utils.parsesize(sizeStr)
self.assertEqual(1024, size)
def test_kb_size(self):
sizeStr = "5KB"
size = comms_utils.parsesize(sizeStr)
self.assertEqual(5120, size)
def test_mb_size(self):
sizeStr = "3MB"
size = comms_utils.parsesize(sizeStr)
self.assertEqual(3145728, size)
def test_gb_size(self):
sizeStr = "2GB"
size = comms_utils.parsesize(sizeStr)
# size is in bytes
self.assertEqual(2147483648, size)