Set global variables, initialize distributed, and set autoresume and random seeds. `allow_no_cuda` should not be set unless using megatron for cpu only data processing. In general this arg should not be set unless you know what you are doing. Returns a function to finalize distri
(
extra_args_provider=None,
args_defaults={},
ignore_unknown_args=False,
allow_no_cuda=False,
)
| 18 | |
| 19 | |
| 20 | def initialize_megatron( |
| 21 | extra_args_provider=None, |
| 22 | args_defaults={}, |
| 23 | ignore_unknown_args=False, |
| 24 | allow_no_cuda=False, |
| 25 | ): |
| 26 | """Set global variables, initialize distributed, and |
| 27 | set autoresume and random seeds. |
| 28 | `allow_no_cuda` should not be set unless using megatron for cpu only |
| 29 | data processing. In general this arg should not be set unless you know |
| 30 | what you are doing. |
| 31 | Returns a function to finalize distributed env initialization |
| 32 | (optionally, only when args.lazy_mpu_init == True) |
| 33 | """ |
| 34 | if not allow_no_cuda: |
| 35 | # Make sure cuda is available. |
| 36 | assert torch.cuda.is_available(), "Megatron requires CUDA." |
| 37 | |
| 38 | # Parse arguments |
| 39 | args = parse_args(extra_args_provider, ignore_unknown_args) |
| 40 | |
| 41 | if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False): |
| 42 | assert args.load is not None, "--use-checkpoints-args requires --load argument" |
| 43 | load_args_from_checkpoint(args) |
| 44 | |
| 45 | validate_args(args, args_defaults) |
| 46 | validate_moe_args(args, args_defaults) |
| 47 | |
| 48 | # set global args, build tokenizer, and set adlr-autoresume, |
| 49 | # tensorboard-writer, and timers. |
| 50 | set_global_variables(args) |
| 51 | |
| 52 | # torch.distributed initialization |
| 53 | def finish_mpu_init(): |
| 54 | args = get_args() |
| 55 | # Pytorch distributed. |
| 56 | _initialize_distributed() |
| 57 | |
| 58 | # Random seeds for reproducibility. |
| 59 | if args.rank == 0: |
| 60 | print("> setting random seeds to {} ...".format(args.seed)) |
| 61 | _set_random_seed(args.seed, args.data_parallel_random_init) |
| 62 | |
| 63 | args = get_args() |
| 64 | if args.lazy_mpu_init: |
| 65 | # TODO is this still a necessary option? |
| 66 | args.use_cpu_initialization = True |
| 67 | # delayed initialization of DDP-related stuff |
| 68 | # We only set basic DDP globals |
| 69 | mpu.set_tensor_model_parallel_world_size(args.tensor_model_parallel_size) |
| 70 | # and return function for external DDP manager |
| 71 | # to call when it has DDP initialized |
| 72 | mpu.set_tensor_model_parallel_rank(args.rank) |
| 73 | return finish_mpu_init |
| 74 | else: |
| 75 | # Megatron's MPU is the master. Complete initialization right away. |
| 76 | finish_mpu_init() |
| 77 |
no test coverage detected