Add files using upload-large-folder tool

4ee152a verified over 1 year ago

5.13 kB

# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.

"""Megatron initialization."""

import torch
from datetime import timedelta

from megatron import get_args
from megatron.core import mpu, tensor_parallel
from megatron.arguments import parse_args, validate_args
from megatron.checkpointing import load_args_from_checkpoint
from megatron.global_vars import set_global_variables
from megatron.initialize import _set_random_seed, _init_autoresume
from megatron.initialize import _compile_dependencies

from .arguments import validate_moe_args



def initialize_megatron(
    extra_args_provider=None,
    args_defaults={},
    ignore_unknown_args=False,
    allow_no_cuda=False,
):
    """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 distributed env initialization
    (optionally, only when args.lazy_mpu_init == True)
    """
    if not allow_no_cuda:
        # Make sure cuda is available.
        assert torch.cuda.is_available(), "Megatron requires CUDA."

    # Parse arguments
    args = parse_args(extra_args_provider, ignore_unknown_args)

    if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False):
        assert args.load is not None, "--use-checkpoints-args requires --load argument"
        load_args_from_checkpoint(args)

    validate_args(args, args_defaults)
    validate_moe_args(args, args_defaults)

    # set global args, build tokenizer, and set adlr-autoresume,
    # tensorboard-writer, and timers.
    set_global_variables(args)

    # torch.distributed initialization
    def finish_mpu_init():
        args = get_args()
        # Pytorch distributed.
        _initialize_distributed()

        # Random seeds for reproducibility.
        if args.rank == 0:
            print("> setting random seeds to {} ...".format(args.seed))
        _set_random_seed(args.seed, args.data_parallel_random_init)

    args = get_args()
    if args.lazy_mpu_init:
        # TODO is this still a necessary option?
        args.use_cpu_initialization = True
        # delayed initialization of DDP-related stuff
        # We only set basic DDP globals
        mpu.set_tensor_model_parallel_world_size(args.tensor_model_parallel_size)
        # and return function for external DDP manager
        # to call when it has DDP initialized
        mpu.set_tensor_model_parallel_rank(args.rank)
        return finish_mpu_init
    else:
        # Megatron's MPU is the master. Complete initialization right away.
        finish_mpu_init()

        # Autoresume.
        _init_autoresume()

        # Compile dependencies.
        _compile_dependencies()

        # No continuation function
        return None

def _initialize_distributed():
    """Initialize torch.distributed and core model parallel."""
    args = get_args()

    device_count = torch.cuda.device_count()
    if torch.distributed.is_initialized():

        if args.rank == 0:
            print(
                "torch distributed is already initialized, "
                "skipping initialization ...",
                flush=True,
            )
        args.rank = torch.distributed.get_rank()
        args.world_size = torch.distributed.get_world_size()

    else:

        if args.rank == 0:
            print("> initializing torch distributed ...", flush=True)
        # Manually set the device ids.
        if device_count > 0:
            device = args.rank % device_count
            if args.local_rank is not None:
                assert (
                    args.local_rank == device
                ), "expected local-rank to be the same as rank % device-count."
            else:
                args.local_rank = device
            torch.cuda.set_device(device)
        # Call the init process
        torch.distributed.init_process_group(
            backend=args.distributed_backend,
            world_size=args.world_size,
            rank=args.rank,
            timeout=timedelta(minutes=args.distributed_timeout_minutes),
        )

    # Set the tensor model-parallel, pipeline model-parallel, and
    # data-parallel communicators.
    if device_count > 0:
        if mpu.model_parallel_is_initialized():
            print("model parallel is already initialized")
        else:
            mpu.initialize_model_parallel(
                args.tensor_model_parallel_size,
                args.pipeline_model_parallel_size,
                args.virtual_pipeline_model_parallel_size,
                args.pipeline_model_parallel_split_rank,
            )

            if args.rank == 0:
                print(
                    f"> initialized tensor model parallel with size "
                    f"{mpu.get_tensor_model_parallel_world_size()}"
                )
                print(
                    f"> initialized pipeline model parallel with size "
                    f"{mpu.get_pipeline_model_parallel_world_size()}"
                )

	# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.

	"""Megatron initialization."""

	import torch
	from datetime import timedelta

	from megatron import get_args
	from megatron.core import mpu, tensor_parallel
	from megatron.arguments import parse_args, validate_args
	from megatron.checkpointing import load_args_from_checkpoint
	from megatron.global_vars import set_global_variables
	from megatron.initialize import _set_random_seed, _init_autoresume
	from megatron.initialize import _compile_dependencies

	from .arguments import validate_moe_args



	def initialize_megatron(
	extra_args_provider=None,
	args_defaults={},
	ignore_unknown_args=False,
	allow_no_cuda=False,
	):
	"""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 distributed env initialization
	(optionally, only when args.lazy_mpu_init == True)
	"""
	if not allow_no_cuda:
	# Make sure cuda is available.
	assert torch.cuda.is_available(), "Megatron requires CUDA."

	# Parse arguments
	args = parse_args(extra_args_provider, ignore_unknown_args)

	if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False):
	assert args.load is not None, "--use-checkpoints-args requires --load argument"
	load_args_from_checkpoint(args)

	validate_args(args, args_defaults)
	validate_moe_args(args, args_defaults)

	# set global args, build tokenizer, and set adlr-autoresume,
	# tensorboard-writer, and timers.
	set_global_variables(args)

	# torch.distributed initialization
	def finish_mpu_init():
	args = get_args()
	# Pytorch distributed.
	_initialize_distributed()

	# Random seeds for reproducibility.
	if args.rank == 0:
	print("> setting random seeds to {} ...".format(args.seed))
	_set_random_seed(args.seed, args.data_parallel_random_init)

	args = get_args()
	if args.lazy_mpu_init:
	# TODO is this still a necessary option?
	args.use_cpu_initialization = True
	# delayed initialization of DDP-related stuff
	# We only set basic DDP globals
	mpu.set_tensor_model_parallel_world_size(args.tensor_model_parallel_size)
	# and return function for external DDP manager
	# to call when it has DDP initialized
	mpu.set_tensor_model_parallel_rank(args.rank)
	return finish_mpu_init
	else:
	# Megatron's MPU is the master. Complete initialization right away.
	finish_mpu_init()

	# Autoresume.
	_init_autoresume()

	# Compile dependencies.
	_compile_dependencies()

	# No continuation function
	return None

	def _initialize_distributed():
	"""Initialize torch.distributed and core model parallel."""
	args = get_args()

	device_count = torch.cuda.device_count()
	if torch.distributed.is_initialized():

	if args.rank == 0:
	print(
	"torch distributed is already initialized, "
	"skipping initialization ...",
	flush=True,
	)
	args.rank = torch.distributed.get_rank()
	args.world_size = torch.distributed.get_world_size()

	else:

	if args.rank == 0:
	print("> initializing torch distributed ...", flush=True)
	# Manually set the device ids.
	if device_count > 0:
	device = args.rank % device_count
	if args.local_rank is not None:
	assert (
	args.local_rank == device
	), "expected local-rank to be the same as rank % device-count."
	else:
	args.local_rank = device
	torch.cuda.set_device(device)
	# Call the init process
	torch.distributed.init_process_group(
	backend=args.distributed_backend,
	world_size=args.world_size,
	rank=args.rank,
	timeout=timedelta(minutes=args.distributed_timeout_minutes),
	)

	# Set the tensor model-parallel, pipeline model-parallel, and
	# data-parallel communicators.
	if device_count > 0:
	if mpu.model_parallel_is_initialized():
	print("model parallel is already initialized")
	else:
	mpu.initialize_model_parallel(
	args.tensor_model_parallel_size,
	args.pipeline_model_parallel_size,
	args.virtual_pipeline_model_parallel_size,
	args.pipeline_model_parallel_split_rank,
	)

	if args.rank == 0:
	print(
	f"> initialized tensor model parallel with size "
	f"{mpu.get_tensor_model_parallel_world_size()}"
	)
	print(
	f"> initialized pipeline model parallel with size "
	f"{mpu.get_pipeline_model_parallel_world_size()}"
	)