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,
)
| 42 | |
| 43 | |
| 44 | def initialize_megatron( |
| 45 | extra_args_provider=None, |
| 46 | args_defaults={}, |
| 47 | ignore_unknown_args=False, |
| 48 | allow_no_cuda=False, |
| 49 | ): |
| 50 | """Set global variables, initialize distributed, and |
| 51 | set autoresume and random seeds. |
| 52 | `allow_no_cuda` should not be set unless using megatron for cpu only |
| 53 | data processing. In general this arg should not be set unless you know |
| 54 | what you are doing. |
| 55 | Returns a function to finalize distributed env initialization |
| 56 | (optionally, only when args.lazy_mpu_init == True) |
| 57 | """ |
| 58 | if not allow_no_cuda: |
| 59 | # Make sure cuda is available. |
| 60 | assert torch.cuda.is_available(), "Megatron requires CUDA." |
| 61 | |
| 62 | # Parse args, build tokenizer, and set adlr-autoresume, |
| 63 | # tensorboard-writer, and timers. |
| 64 | set_global_variables( |
| 65 | extra_args_provider=extra_args_provider, |
| 66 | args_defaults=args_defaults, |
| 67 | ignore_unknown_args=ignore_unknown_args, |
| 68 | ) |
| 69 | |
| 70 | # torch.distributed initialization |
| 71 | def finish_mpu_init(): |
| 72 | args = get_args() |
| 73 | # Pytorch distributed. |
| 74 | _initialize_distributed() |
| 75 | |
| 76 | # Random seeds for reproducibility. |
| 77 | if args.rank == 0: |
| 78 | print("> setting random seeds to {} ...".format(args.seed)) |
| 79 | _set_random_seed(args.seed) |
| 80 | |
| 81 | args = get_args() |
| 82 | if args.lazy_mpu_init: |
| 83 | args.use_cpu_initialization = True |
| 84 | # delayed initialization of DDP-related stuff |
| 85 | # We only set basic DDP globals |
| 86 | set_tensor_model_parallel_world_size(args.tensor_model_parallel_size) |
| 87 | # and return function for external DDP manager |
| 88 | # to call when it has DDP initialized |
| 89 | set_tensor_model_parallel_rank(args.rank) |
| 90 | return finish_mpu_init |
| 91 | else: |
| 92 | # Megatron's MPU is the master. Complete initialization right away. |
| 93 | finish_mpu_init() |
| 94 | |
| 95 | # Initialize memory buffers. |
| 96 | _initialize_mem_buffs() |
| 97 | |
| 98 | # Autoresume. |
| 99 | _init_autoresume() |
| 100 | |
| 101 | # No continuation function |