File size: 4,529 Bytes
6011e08 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import logging
import random
import numpy as np
import torch
from megatron.core import mpu, tensor_parallel
from megatron.core.config import set_experimental_flag
from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator
from megatron.training.global_vars import _build_tokenizer, set_args
logger = logging.getLogger(__name__)
def _set_random_seed(
seed_: int,
data_parallel_random_init: bool = False,
te_rng_tracker: bool = False,
inference_rng_tracker: bool = False,
use_cudagraphable_rng: bool = False,
):
"""Set random seed for reproducability."""
# Ensure that different pipeline MP stages get different seeds.
seed = seed_ + (100 * mpu.get_pipeline_model_parallel_rank())
# Ensure different data parallel ranks get different seeds
if data_parallel_random_init:
seed = seed + (10 * mpu.get_data_parallel_rank(with_context_parallel=False))
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
tensor_parallel.model_parallel_cuda_manual_seed(seed, te_rng_tracker, inference_rng_tracker, use_cudagraphable_rng)
def _initialize_distributed(args, get_embedding_ranks=None, get_position_embedding_ranks=None):
"""Initialize torch.distributed and core model parallel."""
# Set the tensor model-parallel, pipeline model-parallel, and
# data-parallel communicators.
mpu.initialize_model_parallel(
args.tensor_model_parallel_size,
args.pipeline_model_parallel_size,
args.virtual_pipeline_model_parallel_size,
pipeline_model_parallel_comm_backend=args.pipeline_model_parallel_comm_backend,
context_parallel_size=args.context_parallel_size,
hierarchical_context_parallel_sizes=args.hierarchical_context_parallel_sizes,
expert_model_parallel_size=args.expert_model_parallel_size,
num_distributed_optimizer_instances=args.num_distributed_optimizer_instances,
expert_tensor_parallel_size=args.expert_tensor_parallel_size,
distributed_timeout_minutes=args.distributed_timeout_minutes,
nccl_communicator_config_path=args.nccl_communicator_config_path,
order="tp-cp-ep-dp-pp" if not args.use_tp_pp_dp_mapping else "tp-cp-ep-pp-dp",
get_embedding_ranks=get_embedding_ranks,
get_position_embedding_ranks=get_position_embedding_ranks,
create_gloo_process_groups=args.enable_gloo_process_groups,
)
def init(args):
set_args(args)
if args.enable_experimental:
logger.info("Enable megatron experimental")
set_experimental_flag(True)
# Pytorch distributed.
_initialize_distributed(args)
# https://github.com/NVIDIA/Megatron-LM/issues/1563
assert np.__version__.startswith("1."), "Megatron does not support numpy 2.x"
# Random seeds for reproducibility.
if args.rank == 0:
logger.info(f"> setting random seeds to {args.seed} ...")
_set_random_seed(
args.seed,
args.data_parallel_random_init,
args.te_rng_tracker,
args.inference_rng_tracker,
)
_build_tokenizer(args)
# We won't use this. initialize to pass some validation in megatron.
init_num_microbatches_calculator(
args.rank,
args.rampup_batch_size,
args.global_batch_size,
args.micro_batch_size,
args.data_parallel_size,
args.decrease_batch_size_if_needed,
)
if args.deterministic_mode:
if args.rank == 0:
logger.info("> running in deterministic mode")
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.use_deterministic_algorithms(True, warn_only=False)
if args.tp_comm_overlap:
from megatron.training.initialize import _initialize_tp_communicators
_initialize_tp_communicators()
if getattr(args, "custom_megatron_init_path", None):
from slime.utils.misc import load_function
custom_init = load_function(args.custom_megatron_init_path)
custom_init(args)
# TODO shall we use a simpler method to determine which rank to init wandb?
def is_megatron_main_rank():
return (
mpu.get_data_parallel_rank(with_context_parallel=True) == 0
and mpu.get_tensor_model_parallel_rank() == 0
and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1
)
|