(cls, args: argparse.Namespace)
| 1347 | @classmethod |
| 1348 | @typechecked |
| 1349 | def main_worker(cls, args: argparse.Namespace): |
| 1350 | |
| 1351 | # 0. Init distributed process |
| 1352 | distributed_option = build_dataclass(DistributedOption, args) |
| 1353 | # Setting distributed_option.dist_rank, etc. |
| 1354 | distributed_option.init_options() |
| 1355 | |
| 1356 | # NOTE(kamo): Don't use logging before invoking logging.basicConfig() |
| 1357 | if not distributed_option.distributed or distributed_option.dist_rank == 0: |
| 1358 | if not distributed_option.distributed: |
| 1359 | _rank = "" |
| 1360 | else: |
| 1361 | _rank = ( |
| 1362 | f":{distributed_option.dist_rank}/" |
| 1363 | f"{distributed_option.dist_world_size}" |
| 1364 | ) |
| 1365 | |
| 1366 | # NOTE(kamo): |
| 1367 | # logging.basicConfig() is invoked in main_worker() instead of main() |
| 1368 | # because it can be invoked only once in a process. |
| 1369 | # FIXME(kamo): Should we use logging.getLogger()? |
| 1370 | logging.basicConfig( |
| 1371 | level=args.log_level, |
| 1372 | format=f"[{os.uname()[1].split('.')[0]}{_rank}]" |
| 1373 | f" %(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s", |
| 1374 | ) |
| 1375 | else: |
| 1376 | # Suppress logging if RANK != 0 |
| 1377 | logging.basicConfig( |
| 1378 | level="ERROR", |
| 1379 | format=f"[{os.uname()[1].split('.')[0]}" |
| 1380 | f":{distributed_option.dist_rank}/{distributed_option.dist_world_size}]" |
| 1381 | f" %(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s", |
| 1382 | ) |
| 1383 | # Invoking torch.distributed.init_process_group |
| 1384 | distributed_option.init_torch_distributed() |
| 1385 | |
| 1386 | # 1. Set random-seed |
| 1387 | set_all_random_seed(args.seed) |
| 1388 | torch.backends.cudnn.enabled = args.cudnn_enabled |
| 1389 | torch.backends.cudnn.benchmark = args.cudnn_benchmark |
| 1390 | torch.backends.cudnn.deterministic = args.cudnn_deterministic |
| 1391 | if args.detect_anomaly: |
| 1392 | logging.info("Invoking torch.autograd.set_detect_anomaly(True)") |
| 1393 | torch.autograd.set_detect_anomaly(args.detect_anomaly) |
| 1394 | |
| 1395 | if args.use_tf32: |
| 1396 | # Accelerate matmul at the cost of precision. |
| 1397 | # Only effective with Ampere GPUs and above |
| 1398 | # https://pytorch.org/docs/stable/notes/cuda.html |
| 1399 | assert not args.use_amp, "amp is not compatible with tf32" |
| 1400 | torch.backends.cuda.matmul.allow_tf32 = True |
| 1401 | torch.backends.cudnn.allow_tf32 = True |
| 1402 | logging.info("Using TensorFloat32 at the cost of matmul precision") |
| 1403 | |
| 1404 | if ( |
| 1405 | args.collect_stats |
| 1406 | and getattr(args, "model_conf", None) is not None |
no test coverage detected