Parse all the args.
(args_list=None, parser=None)
| 57 | |
| 58 | |
| 59 | def get_args(args_list=None, parser=None): |
| 60 | """Parse all the args.""" |
| 61 | if parser is None: |
| 62 | parser = argparse.ArgumentParser(description="sat") |
| 63 | else: |
| 64 | assert isinstance(parser, argparse.ArgumentParser) |
| 65 | parser = add_model_config_args(parser) |
| 66 | parser = add_sampling_config_args(parser) |
| 67 | parser = add_training_args(parser) |
| 68 | parser = add_evaluation_args(parser) |
| 69 | parser = add_data_args(parser) |
| 70 | |
| 71 | import deepspeed |
| 72 | |
| 73 | parser = deepspeed.add_config_arguments(parser) |
| 74 | |
| 75 | args = parser.parse_args(args_list) |
| 76 | args = process_config_to_args(args) |
| 77 | |
| 78 | if not args.train_data: |
| 79 | print_rank0("No training data specified", level="WARNING") |
| 80 | |
| 81 | assert (args.train_iters is None) or (args.epochs is None), "only one of train_iters and epochs should be set." |
| 82 | if args.train_iters is None and args.epochs is None: |
| 83 | args.train_iters = 10000 # default 10k iters |
| 84 | print_rank0("No train_iters (recommended) or epochs specified, use default 10k iters.", level="WARNING") |
| 85 | |
| 86 | args.cuda = torch.cuda.is_available() |
| 87 | |
| 88 | args.rank = int(os.getenv("RANK", "0")) |
| 89 | args.world_size = int(os.getenv("WORLD_SIZE", "1")) |
| 90 | if args.local_rank is None: |
| 91 | args.local_rank = int(os.getenv("LOCAL_RANK", "0")) # torchrun |
| 92 | |
| 93 | if args.device == -1: |
| 94 | if torch.cuda.device_count() == 0: |
| 95 | args.device = "cpu" |
| 96 | elif args.local_rank is not None: |
| 97 | args.device = args.local_rank |
| 98 | else: |
| 99 | args.device = args.rank % torch.cuda.device_count() |
| 100 | |
| 101 | if args.local_rank != args.device and args.mode != "inference": |
| 102 | raise ValueError( |
| 103 | "LOCAL_RANK (default 0) and args.device inconsistent. " |
| 104 | "This can only happens in inference mode. " |
| 105 | "Please use CUDA_VISIBLE_DEVICES=x for single-GPU training. " |
| 106 | ) |
| 107 | |
| 108 | if args.rank == 0: |
| 109 | print_rank0("using world size: {}".format(args.world_size)) |
| 110 | |
| 111 | if args.train_data_weights is not None: |
| 112 | assert len(args.train_data_weights) == len(args.train_data) |
| 113 | |
| 114 | if args.mode != "inference": # training with deepspeed |
| 115 | args.deepspeed = True |
| 116 | if args.deepspeed_config is None: # not specified |
no test coverage detected