Parse all the args.
(args_list=None, parser=None)
| 85 | |
| 86 | |
| 87 | def get_args(args_list=None, parser=None): |
| 88 | """Parse all the args.""" |
| 89 | if parser is None: |
| 90 | parser = argparse.ArgumentParser(description="sat") |
| 91 | else: |
| 92 | assert isinstance(parser, argparse.ArgumentParser) |
| 93 | parser = add_model_config_args(parser) |
| 94 | parser = add_sampling_config_args(parser) |
| 95 | parser = add_training_args(parser) |
| 96 | parser = add_evaluation_args(parser) |
| 97 | parser = add_data_args(parser) |
| 98 | |
| 99 | import deepspeed |
| 100 | |
| 101 | parser = deepspeed.add_config_arguments(parser) |
| 102 | |
| 103 | args = parser.parse_args(args_list) |
| 104 | args = process_config_to_args(args) |
| 105 | |
| 106 | if not args.train_data: |
| 107 | print_rank0("No training data specified", level="WARNING") |
| 108 | |
| 109 | assert (args.train_iters is None) or (args.epochs is None), "only one of train_iters and epochs should be set." |
| 110 | if args.train_iters is None and args.epochs is None: |
| 111 | args.train_iters = 10000 # default 10k iters |
| 112 | print_rank0("No train_iters (recommended) or epochs specified, use default 10k iters.", level="WARNING") |
| 113 | |
| 114 | args.cuda = torch.cuda.is_available() |
| 115 | |
| 116 | args.rank = int(os.getenv("RANK", "0")) |
| 117 | args.world_size = int(os.getenv("WORLD_SIZE", "1")) |
| 118 | if args.local_rank is None: |
| 119 | args.local_rank = int(os.getenv("LOCAL_RANK", "0")) # torchrun |
| 120 | |
| 121 | if args.device == -1: |
| 122 | if torch.cuda.device_count() == 0: |
| 123 | args.device = "cpu" |
| 124 | elif args.local_rank is not None: |
| 125 | args.device = args.local_rank |
| 126 | else: |
| 127 | args.device = args.rank % torch.cuda.device_count() |
| 128 | |
| 129 | if args.local_rank != args.device and args.mode != "inference": |
| 130 | raise ValueError( |
| 131 | "LOCAL_RANK (default 0) and args.device inconsistent. " |
| 132 | "This can only happens in inference mode. " |
| 133 | "Please use CUDA_VISIBLE_DEVICES=x for single-GPU training. " |
| 134 | ) |
| 135 | |
| 136 | if args.rank == 0: |
| 137 | print_rank0("using world size: {}".format(args.world_size)) |
| 138 | |
| 139 | if args.train_data_weights is not None: |
| 140 | assert len(args.train_data_weights) == len(args.train_data) |
| 141 | |
| 142 | if args.mode != "inference": # training with deepspeed |
| 143 | args.deepspeed = True |
| 144 | if args.deepspeed_config is None: # not specified |
no test coverage detected