Megatron-DeepSpeed

50820352 - Extract args.load prior to load_ds_checkpoint_and_setup_megatron