Parse all the args.
(args_list=None, parser=None)
| 337 | return False |
| 338 | |
| 339 | def get_args(args_list=None, parser=None): |
| 340 | """Parse all the args.""" |
| 341 | if parser is None: |
| 342 | parser = argparse.ArgumentParser(description='sat') |
| 343 | else: |
| 344 | assert isinstance(parser, argparse.ArgumentParser) |
| 345 | parser = add_model_config_args(parser) |
| 346 | parser = add_training_args(parser) |
| 347 | parser = add_evaluation_args(parser) |
| 348 | parser = add_data_args(parser) |
| 349 | parser = add_tokenization_args(parser) |
| 350 | parser = add_text_generate_args(parser) |
| 351 | |
| 352 | # Include DeepSpeed configuration arguments |
| 353 | import deepspeed |
| 354 | parser = deepspeed.add_config_arguments(parser) |
| 355 | |
| 356 | args = parser.parse_args(args_list) |
| 357 | if not args.iterable_dataset_eval: |
| 358 | args.iterable_dataset_eval = args.iterable_dataset |
| 359 | else: |
| 360 | args.iterable_dataset_eval = eval(args.iterable_dataset_eval) |
| 361 | |
| 362 | if not args.train_data: |
| 363 | print_rank0('No training data specified', level='WARNING') |
| 364 | |
| 365 | assert (args.train_iters is None) or (args.epochs is None), 'only one of train_iters and epochs should be set.' |
| 366 | if args.train_iters is None and args.epochs is None: |
| 367 | args.train_iters = 10000 # default 10k iters |
| 368 | print_rank0('No train_iters (recommended) or epochs specified, use default 10k iters.', level='WARNING') |
| 369 | |
| 370 | args.cuda = torch.cuda.is_available() |
| 371 | |
| 372 | args.rank = int(os.getenv('RANK', '0')) |
| 373 | args.world_size = int(os.getenv("WORLD_SIZE", '1')) |
| 374 | if args.local_rank is None: |
| 375 | args.local_rank = int(os.getenv("LOCAL_RANK", '0')) # torchrun |
| 376 | |
| 377 | if args.device == -1: # not set manually |
| 378 | if torch.cuda.device_count() == 0: |
| 379 | args.device = 'cpu' |
| 380 | elif args.local_rank is not None: |
| 381 | args.device = args.local_rank |
| 382 | else: |
| 383 | args.device = args.rank % torch.cuda.device_count() |
| 384 | |
| 385 | # local rank should be consistent with device in DeepSpeed |
| 386 | if args.local_rank != args.device and args.mode != 'inference': |
| 387 | raise ValueError( |
| 388 | 'LOCAL_RANK (default 0) and args.device inconsistent. ' |
| 389 | 'This can only happens in inference mode. ' |
| 390 | 'Please use CUDA_VISIBLE_DEVICES=x for single-GPU training. ' |
| 391 | ) |
| 392 | |
| 393 | # args.model_parallel_size = min(args.model_parallel_size, args.world_size) |
| 394 | if args.rank == 0: |
| 395 | print_rank0('using world size: {} and model-parallel size: {} '.format( |
| 396 | args.world_size, args.model_parallel_size)) |