Inference code for Persimmon-8B
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch

from .global_vars import get_args
from .global_vars import get_current_global_batch_size
from .global_vars import get_num_microbatches
from .global_vars import update_num_microbatches
from .global_vars import get_tokenizer
from .global_vars import get_tensorboard_writer
from .global_vars import get_adlr_autoresume
from .global_vars import get_timers
from .global_vars import get_global_memory_buffer

from .utils import print_rank_0, is_last_rank, print_rank_last
from .initialize import initialize_megatron

megatron/__init__.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Megatron arguments."""

import argparse
import os
from typing import Optional

import torch


class StoreDictKeyPair(argparse.Action):
    """Custom action to parse list of comma separated key-value pairs into a dictionary.
    For example:
        --my_dict a,1 b,2 c,3
    Supports a split character other than the default comma.
    """

    def __init__(self, option_strings, dest, nargs=None, split_char=",", **kwargs):
        self._nargs = nargs

        self._split_char = split_char
        super(StoreDictKeyPair, self).__init__(
            option_strings, dest, nargs=nargs, **kwargs
        )

    def __call__(self, parser, namespace, values, option_string=None):
        my_dict = dict(kv_pair.split(self._split_char) for kv_pair in values)
        setattr(namespace, self.dest, my_dict)


def parse_args(extra_args_provider=None, defaults={}, ignore_unknown_args=False):
    """Parse all arguments."""
    parser = argparse.ArgumentParser(
        description="Megatron-LM Arguments", allow_abbrev=False
    )

    parser.add_argument(
        "--config",
        type=str,
        default=None,
        help="Path to the launch config.",
    )

    # Standard arguments.
    parser = _add_network_size_args(parser)
    parser = _add_regularization_args(parser)
    parser = _add_training_args(parser)
    parser = _add_initialization_args(parser)
    parser = _add_learning_rate_args(parser)
    parser = _add_checkpointing_args(parser)
    parser = _add_mixed_precision_args(parser)
    parser = _add_distributed_args(parser)
    parser = _add_validation_args(parser)
    parser = _add_data_args(parser)
    parser = _add_autoresume_args(parser)
    parser = _add_biencoder_args(parser)
    parser = _add_inference_args(parser)

    # Custom arguments.
    if extra_args_provider is not None:
        parser = extra_args_provider(parser)

    # Parse.
    if ignore_unknown_args:
        args, _ = parser.parse_known_args()
    else:
        args = parser.parse_args()

    # Args from environment
    args.rank = int(os.getenv("RANK", "0"))
    args.world_size = int(os.getenv("WORLD_SIZE", "1"))

    return args


def validate_args(args, defaults={}):
    # Tensor model parallel size.
    args.tensor_model_parallel_size = min(
        args.tensor_model_parallel_size, args.world_size
    )
    assert (
        args.world_size % args.tensor_model_parallel_size == 0
    ), "world size" " ({}) is not divisible by tensor model parallel size ({})".format(
        args.world_size, args.tensor_model_parallel_size
    )
    # Pipeline model parallel size.
    args.pipeline_model_parallel_size = min(
        args.pipeline_model_parallel_size,
        (args.world_size // args.tensor_model_parallel_size),
    )
    args.transformer_pipeline_model_parallel_size = (
        args.pipeline_model_parallel_size - 1
        if args.standalone_embedding_stage
        else args.pipeline_model_parallel_size
    )
    # Checks.
    model_parallel_size = (
        args.pipeline_model_parallel_size * args.tensor_model_parallel_size
    )
    assert args.world_size % model_parallel_size == 0, (
        "world size is not"
        " divisible by tensor parallel size ({}) times pipeline parallel "
        "size ({})".format(
            args.world_size,
            args.tensor_model_parallel_size,
            args.pipeline_model_parallel_size,
        )
    )
    args.data_parallel_size = args.world_size // model_parallel_size
    if args.rank == 0:
        print(
            "using world size: {}, data-parallel-size: {}, "
            "tensor-model-parallel size: {}, "
            "pipeline-model-parallel size: {} ".format(
                args.world_size,
                args.data_parallel_size,
                args.tensor_model_parallel_size,
                args.pipeline_model_parallel_size,
            ),
            flush=True,
        )
    if args.pipeline_model_parallel_size > 1:
        if args.pipeline_model_parallel_split_rank is not None:
            assert (
                args.pipeline_model_parallel_split_rank
                < args.pipeline_model_parallel_size
            ), (
                "split rank needs"
                " to be less than pipeline model parallel size ({})".format(
                    args.pipeline_model_parallel_size
                )
            )

    # Deprecated arguments
    assert args.batch_size is None, (
        "--batch-size argument is no longer " "valid, use --micro-batch-size instead"
    )
    del args.batch_size
    assert args.warmup is None, (
        "--warmup argument is no longer valid, use " "--lr-warmup-fraction instead"
    )
    del args.warmup
    assert args.model_parallel_size is None, (
        "--model-parallel-size is no "
        "longer valid, use --tensor-model-parallel-size instead"
    )
    del args.model_parallel_size

    if args.checkpoint_activations:
        args.recompute_granularity = "full"
        args.recompute_method = "uniform"
        if args.rank == 0:
            print(
                "--checkpoint-activations is no longer valid, "
                "use --recompute-granularity and --recompute-method  instead. "
                "Defaulting to recompute-granularity=full and recompute-method=uniform."
            )
    del args.checkpoint_activations

    if args.recompute_activations:
        args.recompute_granularity = "selective"
    del args.recompute_activations

    # Set input defaults.
    for key in defaults:
        # For default to be valid, it should not be provided in the
        # arguments that are passed to the program. We check this by
        # ensuring the arg is set to None.
        if getattr(args, key) is not None:
            if args.rank == 0:
                print(
                    "WARNING: overriding default arguments for {key}:{v} \
                       with {key}:{v2}".format(
                        key=key, v=defaults[key], v2=getattr(args, key)
                    ),
                    flush=True,
                )
        else:
            setattr(args, key, defaults[key])

    # Batch size.
    assert args.micro_batch_size is not None
    assert args.micro_batch_size > 0
    if args.global_batch_size is None:
        args.global_batch_size = args.micro_batch_size * args.data_parallel_size
        if args.rank == 0:
            print(
                "setting global batch size to {}".format(args.global_batch_size),
                flush=True,
            )
    assert args.global_batch_size > 0
    if args.num_layers_per_virtual_pipeline_stage is not None:
        assert args.pipeline_model_parallel_size > 2, (
            "pipeline-model-parallel size should be greater than 2 with "
            "interleaved schedule"
        )
        assert args.num_layers % args.num_layers_per_virtual_pipeline_stage == 0, (
            "number of layers is not divisible by number of layers per virtual "
            "pipeline stage"
        )
        args.virtual_pipeline_model_parallel_size = (
            args.num_layers // args.transformer_pipeline_model_parallel_size
        ) // args.num_layers_per_virtual_pipeline_stage
    else:
        args.virtual_pipeline_model_parallel_size = None

    # Parameters dtype.
    args.params_dtype = torch.float
    if args.fp16:
        assert not args.bf16
        args.params_dtype = torch.half
    if args.bf16:
        assert not args.fp16
        args.params_dtype = torch.bfloat16
        # bfloat16 requires gradient accumulation and all-reduce to
        # be done in fp32.
        if not args.accumulate_allreduce_grads_in_fp32:
            args.accumulate_allreduce_grads_in_fp32 = True
            if args.rank == 0:
                print(
                    "accumulate and all-reduce gradients in fp32 for "
                    "bfloat16 data type.",
                    flush=True,
                )

    if args.rank == 0:
        print("using {} for parameters ...".format(args.params_dtype), flush=True)

    # If we do accumulation and all-reduces in fp32, we need to have local DDP
    # and we should make sure use-contiguous-buffers-in-local-ddp is not off.
    if args.accumulate_allreduce_grads_in_fp32:
        assert args.DDP_impl == "local"
        assert args.use_contiguous_buffers_in_local_ddp
    else:
        if args.gradient_accumulation_fusion:
            args.gradient_accumulation_fusion = False
            if args.rank == 0:
                print(
                    "Gradient accumulation fusion to linear layer weight "
                    "gradient computation is supported only with fp32 "
                    "gradient accumulation. Setting gradient_accumulation_fusion "
                    "to False",
                    flush=True,
                )

    # If we use the distributed optimizer, we need to have local DDP
    # and we should make sure use-contiguous-buffers-in-local-ddp is on.
    if args.use_distributed_optimizer:
        assert args.DDP_impl == "local"
        assert args.use_contiguous_buffers_in_local_ddp

    # For torch DDP, we do not use contiguous buffer
    if args.DDP_impl == "torch":
        args.use_contiguous_buffers_in_local_ddp = False

    if args.dataloader_type is None:
        args.dataloader_type = "single"

    # Consumed tokens.
    args.consumed_train_samples = 0
    args.consumed_valid_samples = 0

    # Iteration-based training.
    if args.train_iters:
        # If we use iteration-based training, make sure the
        # sample-based options are off.
        assert args.train_samples is None, "expected iteration-based training"
        assert (
            args.lr_decay_samples is None
        ), "expected iteration-based learning rate decay"
        assert (
            args.lr_warmup_samples == 0
        ), "expected iteration-based learning rate warmup"
        assert (
            args.rampup_batch_size is None
        ), "expected no batch-size rampup for iteration-based training"
        if args.lr_warmup_fraction is not None:
            assert (
                args.lr_warmup_iters == 0
            ), "can only specify one of lr-warmup-fraction and lr-warmup-iters"

    # Sample-based training.
    if args.train_samples:
        # If we use sample-based training, make sure the
        # iteration-based options are off.
        assert args.train_iters is None, "expected sample-based training"
        assert args.lr_decay_iters is None, "expected sample-based learning rate decay"
        assert args.lr_warmup_iters == 0, "expected sample-based learnig rate warmup"
        if args.lr_warmup_fraction is not None:
            assert args.lr_warmup_samples == 0, (
                "can only specify one of lr-warmup-fraction " "and lr-warmup-samples"
            )

    if args.num_layers is not None:
        if args.encoder_num_layers is not None:
            assert (
                args.num_layers == args.encoder_num_layers
            ), "cannot have num-layers and encoder-num-layers be unequal."
        args.encoder_num_layers = args.num_layers
    else:
        args.num_layers = args.encoder_num_layers

    if args.decoder_num_layers is None:
        args.decoder_num_layers = args.num_layers

    # Checks.
    if args.ffn_hidden_size is None and args.hidden_size is not None:
        args.ffn_hidden_size = 4 * args.hidden_size

    if args.kv_channels is None:
        if args.hidden_size and args.num_attention_heads:
            assert args.hidden_size % args.num_attention_heads == 0
            args.kv_channels = args.hidden_size // args.num_attention_heads

    if args.seq_length is not None:
        assert args.encoder_seq_length is None
        args.encoder_seq_length = args.seq_length
    elif args.encoder_seq_length is not None:
        assert args.seq_length is None
        args.seq_length = args.encoder_seq_length

    if args.seq_length is not None:
        assert args.max_position_embeddings >= args.seq_length
    if args.decoder_seq_length is not None:
        assert args.max_position_embeddings >= args.decoder_seq_length
    if args.lr is not None:
        assert args.min_lr <= args.lr
    if args.save is not None:
        assert args.save_interval is not None
    # Mixed precision checks.
    if args.fp16_lm_cross_entropy:
        assert args.fp16, "lm cross entropy in fp16 only support in fp16 mode."
    if args.fp32_residual_connection:
        assert (
            args.fp16 or args.bf16
        ), "residual connection in fp32 only supported when using fp16 or bf16."

    if args.weight_decay_incr_style == "constant":
        assert args.start_weight_decay is None
        assert args.end_weight_decay is None
        args.start_weight_decay = args.weight_decay
        args.end_weight_decay = args.weight_decay
    else:
        assert args.start_weight_decay is not None
        assert args.end_weight_decay is not None

    TORCH_MAJOR = int(torch.__version__.split(".")[0])
    TORCH_MINOR = int(torch.__version__.split(".")[1])
    # Persistent fused layer norm.
    if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 11):
        args.no_persist_layer_norm = True
        if args.rank == 0:
            print(
                "Persistent fused layer norm kernel is supported from "
                "pytorch v1.11 (nvidia pytorch container paired with v1.11). "
                "Defaulting to no_persist_layer_norm=True"
            )

    # Activation recomputing.
    if args.distribute_saved_activations:
        assert args.tensor_model_parallel_size > 1, (
            "can distribute "
            "recomputed activations only across tensor model "
            "parallel groups"
        )
        assert args.recompute_granularity == "full", (
            "distributed recompute activations is only "
            "application to full recompute granularity"
        )
        assert args.recompute_method is not None, (
            "for distributed recompute activations to work you "
            "need to use a recompute method "
        )
        assert TORCH_MAJOR >= 1 and TORCH_MINOR >= 10, (
            "distributed recompute activations are supported for pytorch "
            "v1.10 and above (Nvidia Pytorch container >= 21.07). Current "
            "pytorch version is v%s.%s." % (TORCH_MAJOR, TORCH_MINOR)
        )

    if args.recompute_granularity == "selective":
        assert args.recompute_method is None, (
            "recompute method is not yet supported for "
            "selective recomputing granularity"
        )

    # disable sequence parallelism when tp=1
    # to avoid change in numerics when
    # sequence_parallelism is enabled.
    if args.tensor_model_parallel_size == 1:
        args.sequence_parallel = False

    # disable async_tensor_model_parallel_allreduce when
    # model parallel memory optimization is enabled
    if args.sequence_parallel:
        args.async_tensor_model_parallel_allreduce = False

    # If we're loading a checkpoint, we can skip parameter initialization.
    if args.load is not None:
        args.perform_initialization = False

    if args.deterministic:
        assert (
            not args.use_flash_attn
        ), "Flash attention is not deterministic as of 07-25-2023"

    _print_args(args)
    return args


def _print_args(args):
    """Print arguments."""
    mpi_rank = os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", "0")
    if args.rank == 0 and mpi_rank == "0":
        print(
            "------------------------ arguments ------------------------",
            flush=True,
        )
        str_list = []
        for arg in vars(args):
            dots = "." * (48 - len(arg))
            str_list.append("  {} {} {}".format(arg, dots, getattr(args, arg)))
        for arg in sorted(str_list, key=lambda x: x.lower()):
            print(arg, flush=True)
        print(
            "-------------------- end of arguments ---------------------",
            flush=True,
        )


def _check_arg_is_not_none(args, arg):
    assert getattr(args, arg) is not None, "{} argument is None".format(arg)


def _add_inference_args(parser):
    group = parser.add_argument_group(title="inference")

    group.add_argument(
        "--inference-batch-times-seqlen-threshold",
        type=int,
        default=512,
        help="During inference, if batch-size times "
        "sequence-length is smaller than this threshold "
        "then we will not use pipelining, otherwise we will.",
    )
    group.add_argument(
        "--inference-endpoint",
        type=str,
        default="",
        help="Inference endpoint (html address) to query for running an evaluation engine. We bring up an inference server at this endpoint and query it for running an evaluation engine. Format if http://<ip>:<port>/api",
    )
    group.add_argument(
        "--report-file",
        type=str,
        default="",
        help="File to write the inference report to. If not specified, report will be printed to stdout.",
    )
    group.add_argument(
        "--run-text-eval",
        action="store_true",
        help="Run text evals during checkpointing",
    )

    return parser


def _add_network_size_args(parser):
    group = parser.add_argument_group(title="network size")

    group.add_argument(
        "--num-layers", type=int, default=None, help="Number of transformer layers."
    )
    group.add_argument(
        "--encoder-num-layers",
        type=int,
        default=None,
        help="Number of encoder transformer layers.",
    )
    group.add_argument(
        "--decoder-num-layers",
        type=int,
        default=None,
        help="Number of decoder transformer layers.",
    )
    group.add_argument(
        "--hidden-size", type=int, default=None, help="Tansformer hidden size."
    )
    group.add_argument(
        "--ffn-hidden-size",
        type=int,
        default=None,
        help="Transformer Feed-Forward Network hidden size. "
        "This is set to 4*hidden-size if not provided",
    )
    group.add_argument(
        "--num-attention-heads",
        type=int,
        default=None,
        help="Number of transformer attention heads.",
    )
    group.add_argument(
        "--attn-span-type",
        default="global_span",
        choices=["global_span", "local_span"],
        help="Attention span. Default is global",
    )
    group.add_argument(
        "--attn-block-size",
        type=int,
        default=None,
        help="Block size for self-attention",
    )
    group.add_argument(
        "--attn-halo-size", type=int, default=None, help="halo size for self-attention"
    )
    group.add_argument(
        "--use-flash-attn",
        action="store_true",
        help="use FlashAttention implementation of attention. "
        "https://arxiv.org/abs/2205.14135",
    )
    group.add_argument(
        "--use-fast-cross-entropy",
        action="store_true",
        help="use fast cross-entropy loss function "
        "implementation from FlashAttention repo.",
    )
    group.add_argument(
        "--kv-channels",
        type=int,
        default=None,
        help="Projection weights dimension in multi-head "
        "attention. This is set to "
        "   args.hidden_size // args.num_attention_heads "
        "if not provided.",
    )
    group.add_argument(
        "--max-position-embeddings",
        type=int,
        default=None,
        help="Maximum number of position embeddings to use. "
        "This is the size of position embedding.",
    )
    group.add_argument(
        "--pos-emb",
        type=str,
        default="learned",
        choices=["learned", "sinusoidal", "rotary", "alibi"],
        help="Position embedding type",
    )
    group.add_argument(
        "--rotary-pct",
        type=float,
        default=1.0,
        help="Percentage of hidden dims to apply rotary to",
    )

    # A rotary base of 25k is sufficient resolution for a sequence length of 32k. If we use longer sequences, we may
    # need larger value here. See _check_rotary_base in transformer.py for details.
    group.add_argument(
        "--rotary-emb-base",
        type=int,
        default=25000,
        help="Base for rotary position embeddings",
    )
    group.add_argument(
        "--enable-rotary-base-check",
        dest="enable_rotary_base_check",
        action="store_true",
        help="Enables checking that rotary_emb_base is sufficiently high resolution for the current sequence length.",
    )
    group.add_argument(
        "--no-enable-rotary-base-check",
        dest="enable_rotary_base_check",
        action="store_false",
        help="Disables checking that rotary_emb_base is sufficiently high resolution for the current sequence length.",
    )
    group.set_defaults(enable_rotary_base_check=None)

    group.add_argument(
        "--qk-layernorm",
        type=str,
        default="none",
        choices=[
            "none",
            "weights_per_layer",
            "weights_per_tensor_parallel_rank_and_layer",
            "weights_per_attention_head_and_layer",
        ],
        help="QK LayerNorm type",
    )

    group.add_argument(
        "--untie-embeddings",
        default=None,
        action="store_true",
        help="Share input output word embeddings",
    )
    group.add_argument(
        "--make-vocab-size-divisible-by",
        type=int,
        default=128,
        help="Pad the vocab size to be divisible by this value."
        "This is added for computational efficieny reasons.",
    )
    group.add_argument(
        "--layernorm-epsilon", type=float, default=1e-5, help="Layer norm epsilon."
    )
    group.add_argument(
        "--apply-residual-connection-post-layernorm",
        action="store_true",
        help="If set, use original BERT residula connection " "ordering.",
    )
    group.add_argument(
        "--openai-gelu",
        action="store_true",
        help="Use OpenAIs GeLU implementation. This option"
        "should not be used unless for backward compatibility"
        "reasons.",
    )
    group.add_argument(
        "--sq-relu",
        default=None,
        action="store_true",
        help="use squared relu from primer https://arxiv.org/pdf/2109.08668.pdf.",
    )
    group.add_argument(
        "--onnx-safe",
        type=bool,
        required=False,
        help="Use workarounds for known problems with " "Torch ONNX exporter",
    )
    group.add_argument(
        "--bert-no-binary-head",
        action="store_false",
        help="Disable BERT binary head.",
        dest="bert_binary_head",
    )
    group.add_argument(
        "--num-experts",
        type=int,
        default=None,
        help="Number of Experts in Switch Transformer (None means no Switch)",
    )
    return parser


def _add_regularization_args(parser):
    group = parser.add_argument_group(title="regularization")

    group.add_argument(
        "--attention-dropout",
        type=float,
        default=0.0,
        help="Post attention dropout probability.",
    )
    group.add_argument(
        "--hidden-dropout",
        type=float,
        default=0.0,
        help="Dropout probability for hidden state transformer.",
    )
    group.add_argument(
        "--weight-decay",
        type=float,
        default=0.01,
        help="Weight decay coefficient for L2 regularization.",
    )
    group.add_argument(
        "--start-weight-decay",
        type=float,
        help="Initial weight decay coefficient for L2 regularization.",
    )
    group.add_argument(
        "--end-weight-decay",
        type=float,
        help="End of run weight decay coefficient for L2 regularization.",
    )
    group.add_argument(
        "--weight-decay-incr-style",
        type=str,
        default="constant",
        choices=["constant", "linear", "cosine"],
        help="Weight decay increment function.",
    )
    group.add_argument(
        "--clip-grad",
        type=float,
        default=1.0,
        help="Gradient clipping based on global L2 norm.",
    )
    group.add_argument(
        "--adam-beta1",
        type=float,
        default=0.9,
        help="First coefficient for computing running averages "
        "of gradient and its square",
    )
    group.add_argument(
        "--adam-beta2",
        type=float,
        default=0.95,
        help="Second coefficient for computing running averages "
        "of gradient and its square",
    )
    group.add_argument(
        "--adam-eps",
        type=float,
        default=1e-08,
        help="Term added to the denominator to improve" "numerical stability",
    )
    group.add_argument(
        "--sgd-momentum", type=float, default=0.9, help="Momentum factor for sgd"
    )
    group.add_argument(
        "--label-smoothing",
        type=float,
        default=0.0,
        help="Amount of label smoothing to use.",
    )

    return parser


def _add_training_args(parser):
    group = parser.add_argument_group(title="training")
    group.add_argument(
        "--deterministic",
        action="store_true",
        help="If the run should be determinstic",
    )
    group.add_argument(
        "--run_random_id",
        type=str,
        default="",
        help="Random unique identifier for this run. "
        "Appended to all scripts created for the "
        "run, like configs and slurm ones.",
    )
    group.add_argument(
        "--micro-batch-size",
        type=int,
        default=None,
        help="Batch size per model instance (local batch size). "
        "Global batch size is local batch size times data "
        "parallel size times number of micro batches.",
    )
    group.add_argument(
        "--batch-size",
        type=int,
        default=None,
        help="Old batch size parameter, do not use. " "Use --micro-batch-size instead",
    )
    group.add_argument(
        "--global-batch-size",
        type=int,
        default=None,
        help="Training batch size. If set, it should be a "
        "multiple of micro-batch-size times data-parallel-size. "
        "If this value is None, then "
        "use micro-batch-size * data-parallel-size as the "
        "global batch size. This choice will result in 1 for "
        "number of micro-batches.",
    )
    group.add_argument(
        "--rampup-batch-size",
        nargs="*",
        default=None,
        help="Batch size ramp up with the following values:"
        "  --rampup-batch-size <start batch size> "
        "                      <batch size incerement> "
        "                      <ramp-up samples> "
        "For example:"
        "   --rampup-batch-size 16 8 300000 \ "
        "   --global-batch-size 1024"
        "will start with global batch size 16 and over "
        " (1024 - 16) / 8 = 126 intervals will increase"
        "the batch size linearly to 1024. In each interval"
        "we will use approximately 300000 / 126 = 2380 samples.",
    )
    group.add_argument(
        "--recompute-activations",
        action="store_true",
        help="recompute activation to allow for training "
        "with larger models, sequences, and batch sizes.",
    )
    group.add_argument(
        "--recompute-granularity",
        type=str,
        default=None,
        choices=["full", "selective"],
        help="Checkpoint activations to allow for training "
        "with larger models, sequences, and batch sizes. "
        "It is supported at two granularities 1) full: "
        "whole transformer layer is recomputed, "
        "2) selective: core attention part of the transformer "
        "layer is recomputed.",
    )
    group.add_argument(
        "--distribute-saved-activations",
        action="store_true",
        help="If set, distribute recomputed activations "
        "across model parallel group.",
    )
    group.add_argument(
        "--recompute-method",
        type=str,
        default=None,
        choices=["uniform", "block"],
        help="1) uniform: uniformly divide the total number of "
        "Transformer layers and recompute the input activation of "
        "each divided chunk at specified granularity, "
        "2) recompute the input activations of only a set number of "
        "individual Transformer layers per pipeline stage and do the "
        "rest without any recomputing at specified granularity"
        "default) do not apply activations recompute to any layers",
    )
    group.add_argument(
        "--recompute-num-layers",
        type=int,
        default=1,
        help="1) uniform: the number of Transformer layers in each "
        "uniformly divided recompute unit, "
        "2) block: the number of individual Transformer layers "
        "to recompute within each pipeline stage.",
    )

    # deprecated
    group.add_argument(
        "--checkpoint-activations",
        action="store_true",
        help="Checkpoint activation to allow for training "
        "with larger models, sequences, and batch sizes.",
    )
    group.add_argument(
        "--train-iters",
        type=int,
        default=None,
        help="Total number of iterations to train over all "
        "training runs. Note that either train-iters or "
        "train-samples should be provided.",
    )
    group.add_argument(
        "--train-samples",
        type=int,
        default=None,
        help="Total number of samples to train over all "
        "training runs. Note that either train-iters or "
        "train-samples should be provided.",
    )
    group.add_argument(
        "--log-interval", type=int, default=100, help="Report loss and timing interval."
    )
    group.add_argument(
        "--logs-dir", type=str, default=None, help="Directory to save logs for wandb."
    )
    group.add_argument(
        "--tensorboard-dir",
        type=str,
        default=None,
        help="Write TensorBoard logs to this directory.",
    )
    group.add_argument(
        "--no-masked-softmax-fusion",
        action="store_false",
        help="Disable fusion of query_key_value scaling, " "masking, and softmax.",
        dest="masked_softmax_fusion",
    )
    group.add_argument(
        "--no-bias-gelu-fusion",
        action="store_false",
        default=None,
        help="Disable bias and gelu fusion.",
        dest="bias_gelu_fusion",
    )
    group.add_argument(
        "--no-bias-dropout-fusion",
        action="store_false",
        help="Disable bias and dropout fusion.",
        dest="bias_dropout_fusion",
    )
    group.add_argument(
        "--optimizer",
        type=str,
        default="adam",
        choices=["adam", "sgd"],
        help="Optimizer function",
    )
    group.add_argument(
        "--dataloader-type",
        type=str,
        default=None,
        choices=["single", "cyclic"],
        help="Single pass vs multiple pass data loader",
    )
    group.add_argument(
        "--locations-in-token-stream",
        action="store_true",
        default=None,
        help="Whether to put locations in the token stream.",
    )
    group.add_argument(
        "--no-async-tensor-model-parallel-allreduce",
        action="store_false",
        help="Disable asynchronous execution of "
        "tensor-model-parallel all-reduce with weight "
        "gradient compuation of a column-linear layer.",
        dest="async_tensor_model_parallel_allreduce",
    )
    group.add_argument(
        "--no-persist-layer-norm",
        action="store_true",
        help="Disable using persistent fused layer norm kernel. "
        "This kernel supports only a set of hidden sizes. Please "
        "check persist_ln_hidden_sizes if your hidden "
        "size is supported.",
    )
    group.add_argument(
        "--sequence-parallel",
        action="store_true",
        help="Enable sequence parallel optimization.",
    )
    group.add_argument(
        "--no-gradient-accumulation-fusion",
        action="store_false",
        help="Disable fusing gradient accumulation to weight "
        "gradient computation of linear layers",
        dest="gradient_accumulation_fusion",
    )
    group.add_argument(
        "--reset-train-iters-to-zero",
        action="store_true",
        help="Restarts training from iteration 0. "
        "This feature can be used when finetuning or "
        "even when doing 2-stage training (on a new dataset).",
    )
    group.add_argument(
        "--nan-grad-check",
        action="store_true",
        help="Check grads for nans and log which node if found",
    )
    group.add_argument(
        "--continue-training-on-nan",
        action="store_true",
        default=False,
        help="Abort training if loss is NaN. "
        "When this param is False, training exits the "
        " program with exit_code 3 if loss is NaN.",
    )
    group.add_argument(
        "--validate-replicated-parameters-interval",
        type=int,
        default=1,
        help="How frequently to validate replicated parameters.",
    )
    return parser


def _add_initialization_args(parser):
    group = parser.add_argument_group(title="initialization")

    group.add_argument(
        "--seed",
        type=int,
        default=1234,
        help="Random seed used for python, numpy, " "pytorch, and cuda.",
    )
    group.add_argument(
        "--data-parallel-random-init",
        action="store_true",
        help="Enable random initialization of params " "across data parallel ranks",
    )
    group.add_argument(
        "--init-method-std",
        type=float,
        default=0.02,
        help="Standard deviation of the zero mean normal "
        "distribution used for weight initialization.",
    )
    group.add_argument(
        "--force-build-fused-kernels",
        action="store_true",
        default=False,
        help="Force JIT compile all fused kernels."
        "Note: We have now changed the import paths for the"
        " kernels so building the kernels JIT would not exactly work."
    )
    return parser


def _add_learning_rate_args(parser):
    group = parser.add_argument_group(title="learning rate")

    group.add_argument(
        "--lr",
        type=float,
        default=None,
        help="Initial learning rate. Depending on decay style "
        "and initial warmup, the learing rate at each "
        "iteration would be different.",
    )
    group.add_argument(
        "--lr-decay-style",
        type=str,
        default="cosine",
        choices=["constant", "linear", "cosine", "inverse-square-root"],
        help="Learning rate decay function.",
    )
    group.add_argument(
        "--lr-decay-iters",
        type=int,
        default=None,
        help="number of iterations to decay learning rate over,"
        " If None defaults to `--train-iters`",
    )
    group.add_argument(
        "--lr-decay-samples",
        type=int,
        default=None,
        help="number of samples to decay learning rate over,"
        " If None defaults to `--train-samples`",
    )
    group.add_argument(
        "--lr-warmup-fraction",
        type=float,
        default=None,
        help="fraction of lr-warmup-(iters/samples) to use " "for warmup (as a float)",
    )
    group.add_argument(
        "--lr-warmup-iters",
        type=int,
        default=0,
        help="number of iterations to linearly warmup " "learning rate over.",
    )
    group.add_argument(
        "--lr-warmup-samples",
        type=int,
        default=0,
        help="number of samples to linearly warmup " "learning rate over.",
    )
    group.add_argument(
        "--warmup",
        type=int,
        default=None,
        help="Old lr warmup argument, do not use. Use one of the"
        "--lr-warmup-* arguments above",
    )
    group.add_argument(
        "--min-lr",
        type=float,
        default=0.0,
        help="Minumum value for learning rate. The scheduler"
        "clip values below this threshold.",
    )
    group.add_argument(
        "--override-opt_param-scheduler",
        action="store_true",
        help="Reset the values of the scheduler (learning rate,"
        "warmup iterations, minimum learning rate, maximum "
        "number of iterations, and decay style from input "
        "arguments and ignore values from checkpoints. Note"
        "that all the above values will be reset.",
    )

    return parser


def _add_checkpointing_args(parser):
    group = parser.add_argument_group(title="checkpointing")

    group.add_argument(
        "--save",
        type=str,
        default=None,
        help="Output directory to save checkpoints to.",
    )
    group.add_argument(
        "--save-interval",
        type=int,
        default=None,
        help="Number of iterations between checkpoint saves.",
    )
    group.add_argument(
        "--no-save-optim",
        action="store_true",
        default=None,
        help="Do not save current optimizer.",
    )
    group.add_argument(
        "--no-save-rng",
        action="store_true",
        default=None,
        help="Do not save current rng state.",
    )
    group.add_argument(
        "--load",
        type=str,
        default=None,
        help="Directory containing a model checkpoint.",
    )
    group.add_argument(
        "--no-load-optim",
        action="store_true",
        default=None,
        help="Do not load optimizer when loading checkpoint.",
    )
    group.add_argument(
        "--no-load-rng",
        action="store_true",
        default=None,
        help="Do not load rng state when loading checkpoint.",
    )
    group.add_argument(
        "--finetune",
        action="store_true",
        help="Load model for finetuning. Do not load optimizer "
        "or rng state from checkpoint and set iteration to 0. Use "
        "learning rate parameters from the config not the checkpoint.",
    )
    group.add_argument(
        "--no-initialization",
        action="store_false",
        help="Do not perform initialization when building model, "
        "can reduce startup time when definitely loading from a "
        "checkpoint",
        dest="perform_initialization",
    )
    group.add_argument(
        "--keep-last-n", type=int, default=3, help="Number of checkpoints to keep."
    )
    group.add_argument(
        "--is-pipeline-conversion",
        action="store_true",
        default=False,
        help="Are you running the pipeline conversion script",
    )
    group.add_argument(
        "--load-iteration",
        type=int,
        default=None,
        help="Iteration to load from checkpoint. If None, loads latest.",
    )

    return parser


def _add_mixed_precision_args(parser):
    group = parser.add_argument_group(title="mixed precision")

    group.add_argument("--fp16", action="store_true", help="Run model in fp16 mode.")
    group.add_argument(
        "--bf16", action="store_true", help="Run model in bfloat16 mode."
    )
    group.add_argument(
        "--loss-scale",
        type=float,
        default=None,
        help="Static loss scaling, positive power of 2 "
        "values can improve fp16 convergence. If None, dynamic"
        "loss scaling is used.",
    )
    group.add_argument(
        "--initial-loss-scale",
        type=float,
        default=2**32,
        help="Initial loss-scale for dynamic loss scaling.",
    )
    group.add_argument(
        "--min-loss-scale",
        type=float,
        default=1.0,
        help="Minimum loss scale for dynamic loss scale.",
    )
    group.add_argument(
        "--loss-scale-window",
        type=float,
        default=1000,
        help="Window over which to raise/lower dynamic scale.",
    )
    group.add_argument(
        "--hysteresis",
        type=int,
        default=2,
        help="hysteresis for dynamic loss scaling",
    )
    group.add_argument(
        "--fp32-residual-connection",
        action="store_true",
        help="Move residual connections to fp32.",
    )
    group.add_argument(
        "--no-query-key-layer-scaling",
        action="store_false",
        help="Do not scale Q * K^T by 1 / layer-number.",
        dest="apply_query_key_layer_scaling",
    )
    group.add_argument(
        "--attention-softmax-in-fp32",
        action="store_true",
        help="Run attention masking and softmax in fp32. "
        "This flag is ignored unless "
        "--no-query-key-layer-scaling is specified.",
    )
    group.add_argument(
        "--accumulate-allreduce-grads-in-fp32",
        action="store_true",
        help="Gradient accumulation and all-reduce in fp32.",
    )
    group.add_argument(
        "--fp16-lm-cross-entropy",
        action="store_true",
        help="Move the cross entropy unreduced loss calculation" "for lm head to fp16.",
    )

    return parser


def _add_distributed_args(parser):
    group = parser.add_argument_group(title="distributed")

    group.add_argument(
        "--tensor-model-parallel-size",
        type=int,
        default=1,
        help="Degree of tensor model parallelism.",
    )
    group.add_argument(
        "--pipeline-model-parallel-size",
        type=int,
        default=1,
        help="Degree of pipeline model parallelism.",
    )
    group.add_argument(
        "--pipeline-model-parallel-split-rank",
        type=int,
        default=None,
        help="Rank where encoder and decoder should be split.",
    )
    group.add_argument(
        "--model-parallel-size",
        type=int,
        default=None,
        help="Old model parallel argument, do not use. Use "
        "--tensor-model-parallel-size instead.",
    )
    group.add_argument(
        "--num-layers-per-virtual-pipeline-stage",
        type=int,
        default=None,
        help="Number of layers per virtual pipeline stage",
    )
    group.add_argument(
        "--distributed-backend",
        default="nccl",
        choices=["nccl", "gloo"],
        help="Which backend to use for distributed training.",
    )
    group.add_argument(
        "--default-backend",
        type=str,
        default="nccl",
        choices=["nccl", "ucc"],
        help="Which backend to use for non-pipeline comms.",
    )
    group.add_argument(
        "--async-pipeline-comm",
        action="store_true",
        default=False,
        help="Whether to use async comms for pipelining.",
    )
    group.add_argument(
        "--cross-block-networking",
        action="store_true",
        default=False,
        help="Enabling cross block networking uses Socket (TCP) for "
        "communication groups that cross block boundaries.  It "
        "also forces the pipeline comms to use ucc instead of "
        "nccl's Socket so that overlapping is possible.",
    )
    group.add_argument(
        "--force-socket-networking",
        action="store_true",
        default=False,
        help="Always use Socket and never IB",
    )
    group.add_argument(
        "--DDP-impl",
        default="local",
        choices=["local", "torch"],
        help="which DistributedDataParallel implementation " "to use.",
    )
    group.add_argument(
        "--no-contiguous-buffers-in-local-ddp",
        action="store_false",
        help="If set, dont use " "contiguous buffer in local DDP.",
        dest="use_contiguous_buffers_in_local_ddp",
    )
    group.add_argument(
        "--no-scatter-gather-tensors-in-pipeline",
        action="store_false",
        help="Use scatter/gather to optimize communication of tensors in pipeline",
        dest="scatter_gather_tensors_in_pipeline",
    )
    group.add_argument(
        "--local_rank",
        type=int,
        default=None,
        help="local rank passed from distributed launcher.",
    )
    group.add_argument(
        "--lazy-mpu-init",
        type=bool,
        required=False,
        help="If set to True, initialize_megatron() "
        "skips DDP initialization and returns function to "
        "complete it instead. This is for external DDP manager.",
    )
    group.add_argument(
        "--place-on-cpu",
        action="store_true",
        default=None,
        help="If set, place tensors in system (ie, not GPU) memory.",
    )
    group.add_argument(
        "--standalone-embedding-stage",
        action="store_true",
        default=False,
        help="If set, *input* embedding layer "
        "is placed on its own pipeline stage, without any "
        "transformer layers. (For T5, this flag currently only "
        "affects the encoder embedding.)",
    )
    group.add_argument(
        "--use-distributed-optimizer",
        action="store_true",
        help="Use distributed optimizer.",
    )
    group.add_argument(
        "--distributed-comms-timeout",
        type=int,
        default=10,
        help="Timeout for distributed communication in minutes. "
        "This needs to be set to a large number for inference server.",
    )
    group.add_argument(
        "--sync-tensor-replicated-parameters-at-startup",
        action="store_true",
        default=False,
        help="If set, will sync tensor replicated parameters at startup. "
        "Intended to be used once to correct models trained before replication was fixed.",
    )

    return parser


def _add_validation_args(parser):
    group = parser.add_argument_group(title="validation")

    group.add_argument(
        "--eval-iters",
        type=int,
        default=100,
        help="Number of iterations to run for evaluation" "validation/test for.",
    )
    group.add_argument(
        "--eval-interval",
        type=int,
        default=1000,
        help="Interval between running evaluation on " "validation set.",
    )

    return parser


def _add_data_args(parser):
    group = parser.add_argument_group(title="data and dataloader")

    group.add_argument(
        "--data-path",
        nargs="*",
        default=None,
        help="Path to the training dataset. Accepted format:"
        "1) a single data path, 2) multiple datasets in the"
        "form: dataset1-weight dataset1-path dataset2-weight "
        "dataset2-path ...",
    )
    group.add_argument(
        "--data-file",
        type=str,
        default=None,
        help="Optional, if present overrides data-path. Path to a file "
        "which contains a listing of all the datasets in "
        "the format dataset1-weight dataset1-path <newline> ...",
    )
    group.add_argument(
        "--split",
        type=str,
        default="969, 30, 1",
        help="Comma-separated list of proportions for training,"
        " validation, and test split. For example the split "
        "`90,5,5` will use 90%% of data for training, 5%% for "
        "validation and 5%% for test.",
    )
    group.add_argument(
        "--eval-tasks-data-path",
        nargs="*",
        default=None,
        help="Space separated list of tuples of eval task names "
        "and the eval data file paths. Accepted format: "
        "<eval_task_1>,<eval_task_1_data_path> "
        "<eval_task_2>,<eval_task_2_data_path> ..."
        "Note: Currently assumes only 1 tuple and doesn't support "
        "a list.",
    )
    group.add_argument(
        "--eval-harness-tasks",
        type=str,
        default=None,
        help="Comma separated string of eval harness tasks to be evaluated.",
    )
    group.add_argument(
        "--eval-time-delay",
        type=int,
        default=120,
        help="Sleep delay (in seconds) when running evaluator to wait for new checkpoint.",
    )
    group.add_argument(
        "--eval-dir-path",
        type=str,
        default=None,
        help="Path to eval directory path",
    )
    group.add_argument(
        "--eval-max-docs",
        type=int,
        default=10000000,
        help="Total number of documents to evaluate for an eval task.",
    )
    group.add_argument(
        "--eval-last-k-ckpts",
        type=int,
        default=-1,
        help="Number of last k checkpoints to evaluate for an eval task.",
    )
    group.add_argument(
        "--eval-runner-only",
        action="store_true",
        default=False,
        help="Only run eval runner on checkpoint in checkpoint dir.",
    )
    group.add_argument(
        "--vocab-file", type=str, default=None, help="Path to the vocab file."
    )
    group.add_argument(
        "--merge-file",
        type=str,
        default=None,
        help="Path to the BPE merge file.",
    )
    group.add_argument(
        "--vocab-extra-ids",
        type=int,
        default=0,
        help="Number of additional vocabulary tokens. "
        "They are used for span masking in the T5 model",
    )
    group.add_argument(
        "--seq-length",
        type=int,
        default=None,
        help="Maximum sequence length to process.",
    )
    group.add_argument(
        "--encoder-seq-length",
        type=int,
        default=None,
        help="Maximum encoder sequence length to process."
        "This should be exclusive of --seq-length",
    )
    group.add_argument(
        "--decoder-seq-length",
        type=int,
        default=None,
        help="Maximum decoder sequence length to process.",
    )
    group.add_argument(
        "--retriever-seq-length",
        type=int,
        default=256,
        help="Maximum sequence length for the biencoder model " " for retriever",
    )
    group.add_argument(
        "--sample-rate",
        type=float,
        default=1.0,
        help="sample rate for training data. Supposed to be 0 " " < sample_rate < 1",
    )
    group.add_argument(
        "--mask-prob",
        type=float,
        default=0.15,
        help="Probability of replacing a token with mask.",
    )
    group.add_argument(
        "--short-seq-prob",
        type=float,
        default=0.1,
        help="Probability of producing a short sequence.",
    )
    group.add_argument("--mmap-warmup", action="store_true", help="Warm up mmap files.")
    group.add_argument(
        "--num-workers",
        type=int,
        default=2,
        help="Dataloader number of workers.",
    )
    group.add_argument(
        "--tokenizer-type",
        type=str,
        default=None,
        choices=[
            "BertWordPieceLowerCase",
            "BertWordPieceCase",
            "GPT2BPETokenizer",
            "SentencePiece",
        ],
        help="What type of tokenizer to use.",
    )
    group.add_argument(
        "--sp-model-file",
        type=str,
        default=None,
        help="Path to the sentence piece model file",
    )
    group.add_argument(
        "--data-impl",
        type=str,
        default="infer",
        choices=["lazy", "cached", "mmap", "infer"],
        help="Implementation of indexed datasets.",
    )
    group.add_argument(
        "--data-sleep-minutes",
        type=int,
        default=0,
        help="How many minutes for TP ranks != 0 to wait for rank=0 to load data",
    )
    group.add_argument(
        "--data-prefetch-window",
        type=int,
        default=0,
        help="How long the prefetch window for the Mosaic data loader should be. 0 means no prefetching. This only works with Mosaic remote datasets.",
    )
    group.add_argument(
        "--data-all-tensor-parallel-members-load",
        action="store_true",
        help="All tensor parallel members should load the data, not just the leader.",
    )
    group.add_argument(
        "--reset-position-ids",
        action="store_true",
        help="Reset posistion ids after end-of-document token.",
    )
    group.add_argument(
        "--reset-attention-mask",
        action="store_true",
        help="Reset self attention maske after " "end-of-document token.",
    )
    group.add_argument(
        "--eod-mask-loss",
        action="store_true",
        help="Mask loss for the end of document tokens.",
    )
    group.add_argument(
        "--data-packing",
        action="store_true",
        default=True,
        dest="data_packing",
        help="Use data packing for Megatron GPT datasets."
        "This is set to True by default as Megatron uses Packed datasets."
        "If set to False, we use non packed GPT datasets while training.",
    )
    group.add_argument(
        "--no-data-packing",
        action="store_false",
        default=True,
        dest="data_packing",
        help="To not use data packing for Megatron GPT datasets."
        "If set, we use non packed GPT datasets while training.",
    )
    group.add_argument(
        "--dump-blending-results",
        default=False,
        action="store_true",
        help="Whether to store the blending results in the checkpoint dir.",
    )
    group.add_argument(
        "--dataset-subset-size",
        type=int,
        default=None,
        help="Maximum number of samples to use from the dataset. "
        "If None, use all samples.",
    )
    group.add_argument(
        "--last-k-for-valid-in-raw-paths",
        type=int,
        default=1,
        help="Number of raw paths from the end of the list of"
        " `raw_paths` to use for validation.",
    )
    return parser


def _add_autoresume_args(parser):
    group = parser.add_argument_group(title="autoresume")

    group.add_argument(
        "--adlr-autoresume",
        action="store_true",
        help="Enable autoresume on adlr cluster.",
    )
    group.add_argument(
        "--adlr-autoresume-interval",
        type=int,
        default=1000,
        help="Intervals over which check for autoresume" "termination signal",
    )

    return parser


def _add_biencoder_args(parser):
    group = parser.add_argument_group(title="biencoder")

    # network size
    group.add_argument(
        "--ict-head-size",
        type=int,
        default=None,
        help="Size of block embeddings to be used in ICT and "
        "REALM (paper default: 128)",
    )
    group.add_argument(
        "--biencoder-projection-dim",
        type=int,
        default=0,
        help="Size of projection head used in biencoder (paper" " default: 128)",
    )
    group.add_argument(
        "--biencoder-shared-query-context-model",
        action="store_true",
        help="Whether to share the parameters of the query "
        "and context models or not",
    )

    # checkpointing
    group.add_argument(
        "--ict-load",
        type=str,
        default=None,
        help="Directory containing an ICTBertModel checkpoint",
    )
    group.add_argument(
        "--bert-load",
        type=str,
        default=None,
        help="Directory containing an BertModel checkpoint "
        "(needed to start ICT and REALM)",
    )

    # data
    group.add_argument(
        "--titles-data-path",
        type=str,
        default=None,
        help="Path to titles dataset used for ICT",
    )
    group.add_argument(
        "--query-in-block-prob",
        type=float,
        default=0.1,
        help="Probability of keeping query in block for " "ICT dataset",
    )
    group.add_argument(
        "--use-one-sent-docs",
        action="store_true",
        help="Whether to use one sentence documents in ICT",
    )
    group.add_argument(
        "--evidence-data-path",
        type=str,
        default=None,
        help="Path to Wikipedia Evidence frm DPR paper",
    )

    # training
    group.add_argument(
        "--retriever-report-topk-accuracies",
        nargs="+",
        type=int,
        default=[],
        help="Which top-k accuracies to report " "(e.g. '1 5 20')",
    )
    group.add_argument(
        "--retriever-score-scaling",
        action="store_true",
        help="Whether to scale retriever scores by inverse "
        "square root of hidden size",
    )

    # faiss index
    group.add_argument(
        "--block-data-path",
        type=str,
        default=None,
        help="Where to save/load BlockData to/from",
    )
    group.add_argument(
        "--embedding-path",
        type=str,
        default=None,
        help="Where to save/load Open-Retrieval Embedding" " data to/from",
    )

    # indexer
    group.add_argument(
        "--indexer-batch-size",
        type=int,
        default=128,
        help="How large of batches to use when doing indexing " "jobs",
    )
    group.add_argument(
        "--indexer-log-interval",
        type=int,
        default=1000,
        help="After how many batches should the indexer " "report progress",
    )
    return parser

megatron/arguments.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Input/output checkpointing."""

# pylint: disable=consider-using-f-string, invalid-name

import argparse
import os
import random
import shutil
import sys
import threading
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import requests

from megatron import mpu, update_num_microbatches
from megatron.model.module import MegatronModule
from megatron.model.utils import validate_replicated_parameters

from .global_vars import get_args
from .utils import calculate_num_params_and_params_l2_norm, print_rank_0, unwrap_model


MegatronOptimizerType = Any


_CHECKPOINT_VERSION = None

TEXT_EVAL_ENDPOINT = (
    "http://experiments.research.adept.ai/run/fine_tune_and_text_eval_flow"
)


def _get_checkpoint_specific_args() -> List[str]:
    """
    Returns a list of arguments that are specific to the 'trained' checkpoint.

    These mostly are all architecture specific args which cannot be changed
    for an already-trained checkpoint.

    Returns:
        List[str]: List of argument names.
    """

    args_specific_to_checkpoint = [
        "num_layers",
        "hidden_size",
        "ffn_hidden_size",
        "seq_length",
        "num_attention_heads",
        "kv_channels",
        "max_position_embeddings",
        "pos_emb",
        "rotary_pct",
        "rotary_emb_base",
        "untie_embeddings",
        "sq_relu",
        "bf16",
        "fp16",
        "tokenizer_type",
        "padded_vocab_size",
        "vocab_file",
        "merge_file",
        "patch_dim_h",
        "patch_dim_w",
        "encoder_num_layers",
        "bias_gelu_fusion",
        "bias_dropout_fusion",
        "sample_temp_type",
        "sample_temp_max",
        "sample_temp_min",
        "use_quantizer_after",
        "qk_layernorm",
    ]
    return args_specific_to_checkpoint


def _get_sentence_piece_tokenizer_file(
    checkpoint_args: argparse.Namespace,
) -> Optional[str]:
    """Get the sentence piece tokenizer file from the checkpoint args.

    This is a special case for the 16B-2K flagship model because
    we want to allow the checkpoints trained on AWS to be loaded
    on OCI. All other checkpoints should have correct sp_model_file.
    """

    sp_model_file = getattr(checkpoint_args, "sp_model_file", None)
    fsx_prefix = "/fsx/sentence_piece_vocabularies/vocab_new/"
    if sp_model_file is not None and sp_model_file.startswith(fsx_prefix):
        sp_model_file = sp_model_file.replace(fsx_prefix, "/mnt/weka/vocabularies/")
    # For backwards compatibility
    old_container_prefix = "/vocabularies/"
    if sp_model_file is not None and sp_model_file.startswith(old_container_prefix):
        sp_model_file = sp_model_file.replace(
            old_container_prefix, "/mnt/weka/vocabularies/"
        )
    return sp_model_file


def compare_or_set_checkpoint_args(
    checkpoint_args: argparse.Namespace,
    to_compare: bool,
    to_set: bool,
    args: argparse.Namespace,
) -> None:
    """Ensure fixed arguments for a model are the same for the input
    arguments and the one retrieved from checkpoint.

    """

    def _compare_or_set(arg_name: str, old_arg_name: Optional[str] = None) -> None:
        checkpoint_value = getattr(checkpoint_args, arg_name, None)
        if checkpoint_value is None and old_arg_name is not None:
            checkpoint_value = getattr(checkpoint_args, old_arg_name, None)

        if arg_name == "sp_model_file":
            checkpoint_value = _get_sentence_piece_tokenizer_file(checkpoint_args)

        if checkpoint_value is None:
            # If the checkpoint does not have the argument, then we assume
            # it is the same as the current argument.
            # Note: We end up skipping if the checkpoint value was actually None.
            return

        if to_compare:
            args_value = getattr(args, arg_name)
            error_message = (
                "{} value from checkpoint ({}) is not equal to the "
                "input argument value ({}).".format(
                    arg_name, checkpoint_value, args_value
                )
            )
            assert checkpoint_value == args_value, error_message
        if to_set:
            print_rank_0(f"Setting {arg_name} to {checkpoint_value} from checkpoint")
            setattr(args, arg_name, checkpoint_value)

    # Check arguments that are specific to the checkpoint.
    for arg_name in _get_checkpoint_specific_args():
        _compare_or_set(arg_name)

    if args and args.data_parallel_random_init:
        _compare_or_set("data_parallel_random_init")
    _compare_or_set("tensor_model_parallel_size")
    _compare_or_set("pipeline_model_parallel_size")


def ensure_directory_exists(filename: str) -> None:
    """Build filename's path if it does not already exists."""
    dirname = os.path.dirname(filename)
    if not os.path.exists(dirname):
        os.makedirs(dirname)


def get_checkpoint_names(
    checkpoints_path: str,
    iteration: int,
    use_distributed_optimizer: bool,
    pipeline_parallel: Optional[bool] = None,
    tensor_rank: Optional[int] = None,
    pipeline_rank: Optional[int] = None,
) -> Tuple[str, str, str]:
    """Determine the directory name for this rank's checkpoint."""
    directory = "iter_{:07d}".format(iteration)

    # Use both the tensor and pipeline MP rank.
    if pipeline_parallel is None:
        pipeline_parallel = mpu.get_pipeline_model_parallel_world_size() > 1
    if tensor_rank is None:
        tensor_rank = mpu.get_tensor_model_parallel_rank()
    if pipeline_rank is None:
        pipeline_rank = mpu.get_pipeline_model_parallel_rank()

    # Use both the tensor and pipeline MP rank. If using the distributed
    # optimizer, then the optimizer's path must additionally include the
    # data parallel rank.
    base_path = os.path.join(checkpoints_path, directory)
    if not pipeline_parallel:
        common_path = os.path.join(base_path, f"mp_rank_{tensor_rank:02d}")
    else:
        common_path = os.path.join(
            base_path, f"mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}"
        )

    if use_distributed_optimizer:
        model_name = os.path.join(common_path, "model_rng.pt")
        optim_name = os.path.join(
            common_path + "_%03d" % mpu.get_data_parallel_rank(), "optim.pt"
        )
    else:
        model_name = optim_name = os.path.join(common_path, "model_optim_rng.pt")
    return model_name, optim_name, base_path


def find_checkpoint_rank_0(
    checkpoints_path: str, iteration: int, use_distributed_optimizer: bool
) -> Tuple[str, str]:
    """Finds the checkpoint for rank 0 without knowing if we are using
    pipeline parallelism or not.

    Since the checkpoint naming scheme changes if pipeline parallelism
    is present, we need to look for both naming schemes if we don't
    know if the checkpoint has pipeline parallelism.

    """

    # Look for checkpoint with no pipelining
    model_name, optim_name, _ = get_checkpoint_names(
        checkpoints_path,
        iteration,
        use_distributed_optimizer,
        pipeline_parallel=False,
        tensor_rank=0,
        pipeline_rank=0,
    )
    if os.path.isfile(model_name):
        return model_name, optim_name

    # Look for checkpoint with pipelining
    model_name, optim_name, _ = get_checkpoint_names(
        checkpoints_path,
        iteration,
        use_distributed_optimizer,
        pipeline_parallel=True,
        tensor_rank=0,
        pipeline_rank=0,
    )
    if os.path.isfile(model_name):
        return model_name, optim_name

    raise FileNotFoundError(
        f"Could not find checkpoint for rank 0 at {checkpoints_path}"
    )


def get_checkpoint_tracker_filename(checkpoints_path: str) -> str:
    """Tracker file records the latest checkpoint reached during training.
    When continuing a training job, we usually want to start from here.
    """
    return os.path.join(checkpoints_path, "latest_checkpointed_iteration.txt")


def read_metadata(tracker_filename: str) -> int:
    # Read the tracker file and set the iteration
    iteration: int = 0
    max_iter: int = -1
    with open(tracker_filename, "r") as f:
        metastring = f.read().strip()
        try:
            iteration = int(metastring)
        except ValueError:
            print_rank_0(
                "ERROR: Invalid metadata file {}. Exiting".format(tracker_filename)
            )
            sys.exit()
    # NOTE: it isn't entirely clear why a metadata file with a iteration of 0 can't exist
    # cleanup later
    assert iteration > 0, "error parsing metadata file {}".format(tracker_filename)

    # NOTE: also very unclear how the different ranks would ever load different values either!!
    # Get the max iteration retrieved across the ranks.
    if torch.distributed.is_initialized():
        iters_cuda = torch.cuda.LongTensor([iteration])
        torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX)
        max_iter = iters_cuda[0].item()

        # We should now have all the same iteration.
        # If not, print a warning and chose the maximum
        # iteration across all ranks.
        if iteration != max_iter:
            print(
                "WARNING: on rank {} found iteration {} in the "
                "metadata while max iteration across the ranks "
                "is {}, replacing it with max iteration.".format(
                    torch.distributed.get_rank(), iteration, max_iter
                ),
                flush=True,
            )
    else:
        # When loading a checkpoint outside of training (for example,
        # when editing it), we might not have torch distributed
        # initialized, in this case, just assume we have the latest
        max_iter = iteration

    assert max_iter >= 0
    return max_iter


def get_rng_state() -> List[Optional[Dict[str, Any]]]:
    """collect rng state across data parallel ranks"""
    args = get_args()
    rng_state = {
        "random_rng_state": random.getstate(),
        "np_rng_state": np.random.get_state(),
        "torch_rng_state": torch.get_rng_state(),
        "cuda_rng_state": torch.cuda.get_rng_state(),
        "rng_tracker_states": mpu.get_cuda_rng_tracker().get_states(),
    }

    if (
        torch.distributed.is_initialized()
        and mpu.get_data_parallel_world_size() > 1
        and args.data_parallel_random_init
    ):
        rng_state_list: List[Optional[Dict[str, Any]]] = [
            None for i in range(mpu.get_data_parallel_world_size())
        ]
        torch.distributed.all_gather_object(
            rng_state_list, rng_state, group=mpu.get_data_parallel_group()
        )
    else:
        rng_state_list = [rng_state]

    return rng_state_list


def load_checkpoint_from_path(
    model_checkpoint_name: str,
    optim_checkpoint_name: str,
    use_distributed_optimizer: bool,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    # Load the checkpoint.
    try:
        # checkpoint_size = Path(model_checkpoint_name).stat().st_size
        # if torch.distributed.is_initialized():
        #     checkpoint_sizes = [
        #         None,
        #     ] * mpu.get_tensor_model_parallel_world_size()
        #     torch.distributed.all_gather_object(
        #         checkpoint_sizes, checkpoint_size, group=mpu.get_tensor_model_parallel_group()
        #     )
        #     valid = all(c == checkpoint_sizes[0] for c in checkpoint_sizes[1:])
        #     if not valid and torch.distributed.get_rank() == mpu.get_tensor_model_parallel_src_rank():
        #         print(f"Not all Tensor Parallel checkpoints are the same size!")
        #     valid_tensor = torch.cuda.LongTensor([valid])
        #     torch.distributed.all_reduce(valid_tensor, op=torch.distributed.ReduceOp.MIN)
        #     all_valid = bool(valid_tensor[0].item())
        #     if not all_valid:
        #         print_rank_0("Invalid checkpoint! Aborting training.")
        #         sys.exit(1)

        # else:
        #     print_rank_0("WARNING: Distributed is not initialized, cannot verify checkpoint validity.")
        model_state_dict = torch.load(model_checkpoint_name, map_location="cpu")
        if use_distributed_optimizer:
            optim_state_dict = torch.load(optim_checkpoint_name, map_location="cpu")
        else:
            optim_state_dict = model_state_dict
    except BaseException as e:
        print_rank_0("could not load the checkpoint")
        print_rank_0(e)
        sys.exit()

    return model_state_dict, optim_state_dict


def _load_base_checkpoint(
    load_dir: str,
    use_distributed_optimizer: bool,
    rank0: bool = False,
    load_iteration: Optional[int] = None,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """Load the base state_dict from the given directory

    If rank0 is true, just loads rank 0 checkpoint, ignoring arguments.
    """

    # If iteration is specified, set that
    if load_iteration:
        iteration = load_iteration
    else:
        # Read the tracker file
        tracker_filename = get_checkpoint_tracker_filename(load_dir)

        # If no tracker file, return nothing
        if not os.path.isfile(tracker_filename):
            raise FileNotFoundError(f"Could not find metadata file: {tracker_filename}")

        # set the iteration
        iteration = read_metadata(tracker_filename)

    # Checkpoint.
    if rank0:
        checkpoint_names = find_checkpoint_rank_0(
            load_dir, iteration, use_distributed_optimizer
        )
    else:
        model_name, optim_name, _ = get_checkpoint_names(
            load_dir, iteration, use_distributed_optimizer
        )
        checkpoint_names = (model_name, optim_name)
        print_rank_0(f" loading checkpoint from {load_dir} at iteration {iteration}")

    model_checkpoint_name, optim_checkpoint_name = checkpoint_names

    model_state_dict, optim_state_dict = load_checkpoint_from_path(
        model_checkpoint_name, optim_checkpoint_name, use_distributed_optimizer
    )

    return model_state_dict, optim_state_dict


def load_args_from_checkpoint(
    args: argparse.Namespace, load_arg: str = "load"
) -> argparse.Namespace:
    """Set required arguments from the checkpoint specified in the arguments.

    Will overwrite arguments that have a non-None default value, but
    will leave any arguments that default to None as set.

    Returns the same args NameSpace with the new values added/updated.

    If no checkpoint is specified in args, or if the checkpoint is
    there but invalid, the arguments will not be modified

    """
    load_dir = getattr(args, load_arg)

    if load_dir is None:
        print_rank_0("No load directory specified, using provided arguments.")
        return args

    model_state_dict, _ = _load_base_checkpoint(
        load_dir, use_distributed_optimizer=args.use_distributed_optimizer, rank0=True
    )

    # For args we only care about model state dict
    state_dict = model_state_dict

    if not state_dict:
        print_rank_0(
            "Checkpoint not found to provide arguments, using provided arguments."
        )
        return args

    if "args" not in state_dict:
        print_rank_0(
            "Checkpoint provided does not have arguments saved, using provided arguments."
        )
        return args

    checkpoint_args = state_dict["args"]
    args.iteration = state_dict["iteration"]

    compare_or_set_checkpoint_args(
        checkpoint_args, to_compare=False, to_set=True, args=args
    )

    return args


def load_state_dicts_and_update_args(
    load_arg: str = "load",
) -> Tuple[int, Dict[str, Any], Dict[str, Any]]:
    """Loads raw state dicts from checkpoint and updates starting iteration/samples in args.

    Returns:
        - iteration: iteration index to continue training from
        - model_state_dict: dictionary with raw model weights loaded from checkpoint
        - optim_state_dict: dictionary with raw optimizer state weights loaded from checkpoint
    """
    args = get_args()
    load_dir = getattr(args, load_arg)

    model_state_dict, optim_state_dict = _load_base_checkpoint(
        load_dir,
        use_distributed_optimizer=args.use_distributed_optimizer,
        rank0=False,
        load_iteration=args.load_iteration,
    )

    # Check arguments.
    if "args" in model_state_dict:
        checkpoint_args = model_state_dict["args"]
        compare_or_set_checkpoint_args(
            checkpoint_args, to_compare=True, to_set=False, args=args
        )
        if load_arg == "load":
            args.consumed_train_samples = getattr(
                checkpoint_args, "consumed_train_samples", 0
            )
            update_num_microbatches(consumed_samples=args.consumed_train_samples)
            args.consumed_valid_samples = getattr(
                checkpoint_args, "consumed_valid_samples", 0
            )

    # Don't load the VQ-VAE's rng state.
    # NOTE: unclear if anything is actually implementing the logic of the above comment
    if load_arg == "load":
        rng_states(model_state_dict)

    # Set iteration.
    if args.finetune:
        iteration = 0
        args.consumed_train_samples = 0
        args.consumed_valid_samples = 0
    else:
        try:
            iteration = model_state_dict["iteration"]
        except KeyError:
            print_rank_0(
                "A metadata file exists but unable to load iteration from checkpoint {}, exiting".format(
                    load_dir
                )
            )
            sys.exit()
    return iteration, model_state_dict, optim_state_dict


def update_model_from_state_dict(
    model: List[MegatronModule], model_state_dict: Dict[str, Any], strict: bool = True
) -> None:
    """Takes in raw model weights loaded from the last checkpoint and loads into model."""
    print_rank_0("starting update_model_from_state_dict")
    if len(model) == 1:
        model[0].load_state_dict(model_state_dict["model"], strict=strict)
    else:
        for i in range(len(model)):
            mpu.set_virtual_pipeline_model_parallel_rank(i)
            model[i].load_state_dict(model_state_dict["model%d" % i], strict=strict)


def rng_states(model_state_dict: Dict[str, Any]) -> None:
    args = get_args()
    if not args.finetune and not args.no_load_rng:
        try:
            if "rng_state" in model_state_dict:
                # access rng_state for data parallel rank
                if args.data_parallel_random_init:
                    rng_state = model_state_dict["rng_state"][
                        mpu.get_data_parallel_rank()
                    ]
                else:
                    rng_state = model_state_dict["rng_state"][0]
                random.setstate(rng_state["random_rng_state"])
                np.random.set_state(rng_state["np_rng_state"])
                torch.set_rng_state(rng_state["torch_rng_state"])
                torch.cuda.set_rng_state(rng_state["cuda_rng_state"])
                # Check for empty states array
                if not rng_state["rng_tracker_states"]:
                    raise KeyError
                mpu.get_cuda_rng_tracker().set_states(rng_state["rng_tracker_states"])
            else:  # backward compatability
                random.setstate(model_state_dict["random_rng_state"])
                np.random.set_state(model_state_dict["np_rng_state"])
                torch.set_rng_state(model_state_dict["torch_rng_state"])
                torch.cuda.set_rng_state(model_state_dict["cuda_rng_state"])
                # Check for empty states array
                if not model_state_dict["rng_tracker_states"]:
                    raise KeyError
                mpu.get_cuda_rng_tracker().set_states(
                    model_state_dict["rng_tracker_states"]
                )
        except KeyError:
            print_rank_0(
                "Unable to load rng state from checkpoint. "
                "Specify --no-load-rng or --finetune to prevent "
                "attempting to load the rng state, "
                "exiting ..."
            )
            sys.exit()


def update_model_and_optim_from_loaded_data(
    iteration: int,
    model_state_dict: Dict[str, Any],
    optim_state_dict: Dict[str, Any],
    model: List[MegatronModule],
    optimizer: MegatronOptimizerType,
    opt_param_scheduler: None,  # OptimizerParamScheduler,
    strict: bool = True,
    log_norms: bool = True,
) -> None:
    """Take in raw state dicts for model and optimizer and load them into actual model/optimizer/scheduler.

    Args:
        strict: whether to strictly enforce that the keys in the state dict of the checkpoint match the names
            of parameters and buffers in model.
    """
    print_rank_0("starting update_model_and_optim_from_loaded_data")
    args = get_args()
    update_model_from_state_dict(model, model_state_dict, strict)
    # unused in inference mode:
    # update_optimizer_and_scheduler_from_state_dict(optimizer, optim_state_dict, opt_param_scheduler)
    if torch.distributed.is_initialized():
        torch.distributed.barrier()

    print_rank_0(
        f"  successfully loaded checkpoint from {args.load} "
        f"at iteration {iteration}"
    )


def load_checkpoint(
    model: List[MegatronModule],
    optimizer: MegatronOptimizerType,
    opt_param_scheduler: None,  # OptimizerParamScheduler,
    load_arg: str = "load",
    strict: bool = True,
    log_norms: bool = True,
) -> int:
    """
    This function loads a model checkpoint, returning the iteration number.
    Args:
    - model: model object being loaded from checkpoint
    - optimizer: optimizer object being loaded from checkpoint
    - opt_param_scheduler: scheduler being loaded from checkpoint
    - load_arg
    - strict: Whether to strictly enforce name matching between checkpoint's
    state_dict and model's parameters and buffers. Default true.

    Returns:
    Iteration index of the latest checkpoint loaded.
    """
    iteration, model_state_dict, optim_state_dict = load_state_dicts_and_update_args(
        load_arg
    )
    update_model_and_optim_from_loaded_data(
        iteration,
        model_state_dict,
        optim_state_dict,
        model,
        optimizer,
        opt_param_scheduler,
        strict,
        log_norms,
    )
    return iteration

megatron/checkpointing.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023, ADEPT AI LABS INC.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Config objects.

Provides a bridge between the flat `args` namespace and having different
configs at any layer of the model.

Config objects can be initialized manually, but to match current `args`
behavior, they can be initialized using the `from_args` method.
"""

import argparse
import dataclasses
from typing import Type, TypeVar

T = TypeVar("T", bound="Config")

DISABLE_PREFIX = lambda: dataclasses.field(metadata={"no_from_args_prefix": True})


@dataclasses.dataclass
class Config:
    @classmethod
    def from_args(cls: Type[T], args: argparse.Namespace, arg_prefix: str = None) -> T:
        """Initialize a Config object using an `args` object.

        Args:
            args: Parsed arguments from `megatron.get_args()`. The field names
              in the Config object will be populated by `args` entries of the
              same name.
            arg_prefix: If provided, will use this prefix when finding the
              field in `args`. For example, if the field name is `hidden_size`
              and arg_prefix is `blah_`, the field will be populated by the arg
              `blah_hidden_size`.
              However, if the field has the metadata field `no_from_args_prefix`
              set to True, the prefix will not be added. This is to support
              fields that are globally applicable.
        """
        field_values = {}
        for field in dataclasses.fields(cls):
            if issubclass(field.type, Config):
                field_values[field.name] = field.type.from_args(args, arg_prefix)
            else:
                if arg_prefix:
                    prefixed_field_name = arg_prefix + field.name
                    if arg_prefix and not (
                        "no_from_args_prefix" in field.metadata and field.metadata["no_from_args_prefix"]
                    ):
                        arg_name = prefixed_field_name
                    else:
                        arg_name = field.name
                        if prefixed_field_name in vars(args):
                            raise ValueError(
                                f"{field.name} has no_from_args_prefix set, but {prefixed_field_name} exists in args. This is likely a mistake."
                            )
                else:
                    arg_name = field.name

                if arg_name not in vars(args):
                    raise ValueError(
                        f"{arg_name} not found in args when attempting to construct {cls.__name__} from args."
                    )
                field_values[field.name] = vars(args)[arg_name]
        return cls(**field_values)

megatron/config.py
persimmon-ai-labs/adept-inference
import os
import pathlib
import signal
import multiprocessing
import time
import importlib
import sys
from torch.utils import cpp_extension
from megatron.fused_kernels.megatron_fused_kernels.load_kernel import (
    KernelBuildInfo,
    ALL_BUILD_KERNELS,
    JIT_BUILD_KERNELS,
)


class CompilationTimeoutError(Exception):
    pass


def _load_kernel(kernel_build_info):
    kernel_build_info.load()


def load(force_build_fused_kernels=False):
    """Load fused kernels."""
    if force_build_fused_kernels:
        for kernel_build_info in ALL_BUILD_KERNELS.values():
            _load_kernel(kernel_build_info)
    else:
        # Just comile the kernels that we need to JIT compile.
        for kernel_name in JIT_BUILD_KERNELS:
            kernel_build_info = ALL_BUILD_KERNELS[kernel_name]
            _load_kernel(kernel_build_info)

megatron/fused_kernels/__init__.py
persimmon-ai-labs/adept-inference
import pathlib
from pybind11.setup_helpers import Pybind11Extension

srcpath = pathlib.Path(__file__).parent.absolute()


def get_helpers_extension():
    """Get the helpers pybind11 extension."""
    return [
        Pybind11Extension(
            name="helpers",
            sources=[str(srcpath / "helpers.cpp")],
            language="c++",
        )
    ]

megatron/fused_kernels/megatron_fused_kernels/load_cpp_extensions.py
persimmon-ai-labs/adept-inference
import os
import pathlib
from dataclasses import dataclass
from typing import List
from torch.utils import cpp_extension
from pathlib import Path


def _create_build_dir(buildpath):
    try:
        os.mkdir(buildpath)
    except OSError:
        if not os.path.isdir(buildpath):
            print(f"Creation of the build directory {buildpath} failed")


@dataclass
class KernelBuildInfo:
    name: str
    sources: List[Path]
    build_directory: str
    extra_cflags: List[str]
    extra_cuda_cflags: List[str]
    verbose: bool

    def __init__(self, name, sources, build_directory, extra_cflags, extra_cuda_cflags, verbose):
        self.name = name
        self.extra_cflags = extra_cflags
        self.extra_cuda_cflags = extra_cuda_cflags
        self.verbose = verbose
        if not isinstance(build_directory, Path):
            build_directory = Path(build_directory)
        self.build_directory = build_directory
        for i, source in enumerate(sources):
            if not isinstance(source, Path):
                sources[i] = Path(source)
        self.sources = sources

    def __repr__(self):
        return f"KernelBuildInfo(name={self.name}, sources={self.sources}, build_directory={self.build_directory}, extra_cflags={self.extra_cflags}, extra_cuda_cflags={self.extra_cuda_cflags}, verbose={self.verbose})\n"

    def to_setup_cpp_extension(self) -> cpp_extension.CUDAExtension:
        sources = [str(source) for source in self.sources]
        return cpp_extension.CUDAExtension(
            name=self.name,
            sources=sources,
            extra_compile_args={
                "cxx": self.extra_cflags,
                "nvcc": self.extra_cuda_cflags,
            },
            is_python_module=True,
        )

    def load(self):
        os.environ["TORCH_CUDA_ARCH_LIST"] = ""
        _create_build_dir(self.build_directory)
        _ = cpp_extension.load(
            name=self.name,
            sources=self.sources,
            build_directory=Path(self.build_directory),
            extra_cflags=self.extra_cflags,
            extra_cuda_cflags=self.extra_cuda_cflags,
            verbose=self.verbose,
        )


srcpath = pathlib.Path(__file__).parent.absolute()

BUILD_DIR = srcpath / "build"

BASE_CFLAGS = [
    "-O3",
    "-llibtorch_python",
]
BASE_CUDA_CFLAGS = [
    "-O3",
    "--use_fast_math",
    "-gencode",
    "arch=compute_80,code=sm_80",
]

BASE_MASKED_SOFTMAX_FUSION_CUDA_CFLAGS = BASE_CUDA_CFLAGS + [
    "-U__CUDA_NO_HALF_OPERATORS__",
    "-U__CUDA_NO_HALF_CONVERSIONS__",
    "--expt-relaxed-constexpr",
    "--expt-extended-lambda",
]

# These are the kernels that we need to build JIT as they have
# some issues in installing with pip. Note: They are defined in the
# ALL_BUILD_KERNELS dictionary below.
JIT_BUILD_KERNELS = []

ALL_BUILD_KERNELS = {
    "scaled_upper_triang_masked_softmax": KernelBuildInfo(
        name="scaled_upper_triang_masked_softmax_cuda",
        sources=[
            srcpath / "scaled_upper_triang_masked_softmax.cpp",
            srcpath / "scaled_upper_triang_masked_softmax_cuda.cu",
        ],
        build_directory=BUILD_DIR,
        extra_cflags=BASE_CFLAGS,
        extra_cuda_cflags=BASE_MASKED_SOFTMAX_FUSION_CUDA_CFLAGS,
        verbose=True,
    ),
    "scaled_masked_softmax_cuda": KernelBuildInfo(
        name="scaled_masked_softmax_cuda",
        sources=[srcpath / "scaled_masked_softmax.cpp", srcpath / "scaled_masked_softmax_cuda.cu"],
        build_directory=BUILD_DIR,
        extra_cflags=BASE_CFLAGS,
        extra_cuda_cflags=BASE_MASKED_SOFTMAX_FUSION_CUDA_CFLAGS,
        verbose=True,
    ),
    "scaled_softmax_cuda": KernelBuildInfo(
        name="scaled_softmax_cuda",
        sources=[srcpath / "scaled_softmax.cpp", srcpath / "scaled_softmax_cuda.cu"],
        build_directory=BUILD_DIR,
        extra_cflags=BASE_CFLAGS,
        extra_cuda_cflags=BASE_MASKED_SOFTMAX_FUSION_CUDA_CFLAGS,
        verbose=True,
    ),
    "fused_mix_prec_layer_norm_cuda": KernelBuildInfo(
        name="fused_mix_prec_layer_norm_cuda",
        sources=[srcpath / "layer_norm_cuda.cpp", srcpath / "layer_norm_cuda_kernel.cu"],
        build_directory=BUILD_DIR,
        extra_cflags=BASE_CFLAGS,
        extra_cuda_cflags=BASE_CUDA_CFLAGS + ["-maxrregcount=50"],
        verbose=True,
    ),
    "fused_dense_cuda": KernelBuildInfo(
        name="fused_dense_cuda",
        sources=[srcpath / "fused_weight_gradient_dense.cpp", srcpath / "fused_weight_gradient_dense_cuda.cu"],
        build_directory=BUILD_DIR,
        extra_cflags=BASE_CFLAGS,
        extra_cuda_cflags=BASE_CUDA_CFLAGS,
        verbose=True,
    ),
}

megatron/fused_kernels/megatron_fused_kernels/load_kernel.py
persimmon-ai-labs/adept-inference
from setuptools import setup, find_packages
from torch.utils import cpp_extension
from megatron_fused_kernels.load_kernel import ALL_BUILD_KERNELS, JIT_BUILD_KERNELS
from megatron_fused_kernels.load_cpp_extensions import get_helpers_extension


def get_kernel_extensions():
    # remove 'fused_dense_cuda' from the list of kernels to build
    # as it is not used in the current version of Megatron-LM
    # and it causes compilation errors when doing ninja installs.
    # We'll need to JIT compile it instead.
    kernels_to_build = [
        build_info.to_setup_cpp_extension() for k, build_info in ALL_BUILD_KERNELS.items() if k not in JIT_BUILD_KERNELS
    ]
    return kernels_to_build


setup(
    name="megatron_fused_kernels",
    packages=find_packages(exclude=("tests",)),
    packages_dir={"megatron_fused_kernels": "megatron_fused_kernels"},
    ext_modules=get_kernel_extensions() + get_helpers_extension(),
    cmdclass={"build_ext": cpp_extension.BuildExtension},
    zip_safe=False,
    author="Adept AI",
    author_email="",
)

megatron/fused_kernels/setup.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Megatron global variables."""

import functools
import operator
import os
import sys
import time
from functools import reduce
from pathlib import Path
import requests
import torch
import yaml

from megatron.tokenizer import build_tokenizer

from .microbatches import build_num_microbatches_calculator

_GLOBAL_ARGS = None
_GLOBAL_NUM_MICROBATCHES_CALCULATOR = None
_GLOBAL_TOKENIZER = None
_GLOBAL_TENSORBOARD_WRITER = None
_GLOBAL_ADLR_AUTORESUME = None
_GLOBAL_TIMERS = None

_GLOBAL_MEMORY_BUFFER = None


def get_args():
    """Return arguments."""
    _ensure_var_is_initialized(_GLOBAL_ARGS, "args")
    return _GLOBAL_ARGS


def get_num_microbatches():
    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get()


def get_current_global_batch_size():
    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get_current_global_batch_size()


def update_num_microbatches(consumed_samples, consistency_check=True):
    _GLOBAL_NUM_MICROBATCHES_CALCULATOR.update(consumed_samples, consistency_check)


def get_tokenizer():
    """Return tokenizer."""
    _ensure_var_is_initialized(_GLOBAL_TOKENIZER, "tokenizer")
    return _GLOBAL_TOKENIZER


def get_tensorboard_writer():
    """Return tensorboard writer. It can be None so no need
    to check if it is initialized."""
    return _GLOBAL_TENSORBOARD_WRITER


def get_adlr_autoresume():
    """ADLR autoresume object. It can be None so no need
    to check if it is initialized."""
    return _GLOBAL_ADLR_AUTORESUME


def get_timers():
    """Return timers."""
    _ensure_var_is_initialized(_GLOBAL_TIMERS, "timers")
    return _GLOBAL_TIMERS


def get_global_memory_buffer():
    _ensure_var_is_initialized(_GLOBAL_MEMORY_BUFFER, "global memory buffer")
    return _GLOBAL_MEMORY_BUFFER


def set_global_variables(args):
    """Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers."""

    assert args is not None

    _ensure_var_is_not_initialized(_GLOBAL_ARGS, "args")
    set_args(args)

    _build_num_microbatches_calculator(args)
    if args.vocab_file or args.sp_model_file:
        _ = _build_tokenizer(args)
    _set_tensorboard_writer(args)
    _set_adlr_autoresume(args)
    _set_timers()
    _set_global_memory_buffer()


def set_args(args):
    global _GLOBAL_ARGS
    _GLOBAL_ARGS = args


def _build_num_microbatches_calculator(args):

    global _GLOBAL_NUM_MICROBATCHES_CALCULATOR
    _ensure_var_is_not_initialized(
        _GLOBAL_NUM_MICROBATCHES_CALCULATOR, "num microbatches calculator"
    )

    _GLOBAL_NUM_MICROBATCHES_CALCULATOR = build_num_microbatches_calculator(args)


def _build_tokenizer(args):
    """Initialize tokenizer."""
    global _GLOBAL_TOKENIZER
    _ensure_var_is_not_initialized(_GLOBAL_TOKENIZER, "tokenizer")
    _GLOBAL_TOKENIZER = build_tokenizer(args)
    return _GLOBAL_TOKENIZER


def rebuild_tokenizer(args):
    global _GLOBAL_TOKENIZER
    _GLOBAL_TOKENIZER = None
    return _build_tokenizer(args)


def _set_tensorboard_writer(args):
    """Set tensorboard writer."""
    global _GLOBAL_TENSORBOARD_WRITER
    _ensure_var_is_not_initialized(_GLOBAL_TENSORBOARD_WRITER, "tensorboard writer")

    if (
        hasattr(args, "tensorboard_dir")
        and args.tensorboard_dir
        and args.rank == (args.world_size - 1)
    ):
        try:
            from torch.utils.tensorboard import (
                SummaryWriter,
            )  # pylint: disable=import-outside-toplevel

            print("> setting tensorboard ...")
            _GLOBAL_TENSORBOARD_WRITER = SummaryWriter(
                log_dir=args.tensorboard_dir,
                max_queue=args.tensorboard_queue_size,
            )
        except ModuleNotFoundError:
            print(
                "WARNING: TensorBoard writing requested but is not "
                "available (are you using PyTorch 1.1.0 or later?), "
                "no TensorBoard logs will be written.",
                flush=True,
            )


def _set_adlr_autoresume(args):
    """Initialize ADLR autoresume."""
    global _GLOBAL_ADLR_AUTORESUME
    _ensure_var_is_not_initialized(_GLOBAL_ADLR_AUTORESUME, "adlr autoresume")

    if args.adlr_autoresume:
        if args.rank == 0:
            print("enabling autoresume ...", flush=True)
        sys.path.append(os.environ.get("SUBMIT_SCRIPTS", "."))
        try:
            from userlib.auto_resume import (
                AutoResume,
            )  # pylint: disable=import-outside-toplevel
        except BaseException:  # pylint: disable=broad-except
            print("ADLR autoresume is not available, exiting ...")
            sys.exit()

        _GLOBAL_ADLR_AUTORESUME = AutoResume


def _set_timers():
    """Initialize timers."""
    global _GLOBAL_TIMERS
    _ensure_var_is_not_initialized(_GLOBAL_TIMERS, "timers")
    _GLOBAL_TIMERS = Timers()


def _set_global_memory_buffer():
    """Initialize global buffer"""
    global _GLOBAL_MEMORY_BUFFER
    _ensure_var_is_not_initialized(_GLOBAL_MEMORY_BUFFER, "global memory buffer")
    _GLOBAL_MEMORY_BUFFER = GlobalMemoryBuffer()


def _ensure_var_is_initialized(var, name):
    """Make sure the input variable is not None."""
    assert var is not None, f"{name} is not initialized."


def _ensure_var_is_not_initialized(var, name):
    """Make sure the input variable is not None."""
    assert var is None, f"{name} is already initialized."


class _Timer:
    """Timer."""

    def __init__(self, name):
        self.name_ = name
        self.elapsed_ = 0.0
        self.started_ = False
        self.start_time = time.time()

    def start(self):
        """Start the timer."""
        # this import has to be here because of circular dependencies.
        from megatron.mpu import (
            get_data_parallel_group,
        )  # pylint: disable=import-outside-toplevel

        assert not self.started_, "timer has already been started"
        torch.distributed.barrier(get_data_parallel_group())
        torch.cuda.synchronize()
        self.start_time = time.time()
        self.started_ = True

    def stop(self):
        """Stop the timer."""
        # this import has to be here because of circular dependencies.
        from megatron.mpu import (
            get_data_parallel_group,
        )  # pylint: disable=import-outside-toplevel

        assert self.started_, "timer is not started"
        torch.distributed.barrier(get_data_parallel_group())
        torch.cuda.synchronize()
        self.elapsed_ += time.time() - self.start_time
        self.started_ = False

    def reset(self):
        """Reset timer."""
        self.elapsed_ = 0.0
        self.started_ = False

    def elapsed(self, reset=True):
        """Calculate the elapsed time."""
        started = self.started_
        # If the timing in progress, end it first.
        if self.started_:
            self.stop()
        # Get the elapsed time.
        elapsed = self.elapsed_
        # Reset the elapsed time
        if reset:
            self.reset()
        # If timing was in progress, set it back.
        if started:
            self.start()
        return elapsed


class Timers:
    """Group of timers."""

    def __init__(self):
        self.timers = {}

    def __call__(self, name):
        if name not in self.timers:
            self.timers[name] = _Timer(name)
        return self.timers[name]

    def write(self, names, iteration, normalizer=1.0, reset=False):
        """Write timers to a tensorboard writer"""
        # currently when using add_scalars,
        # torch.utils.add_scalars makes each timer its own run, which
        # polutes the runs list, so we just add each as a scalar
        assert normalizer > 0.0
        for name in names:
            value = self.timers[name].elapsed(reset=reset) / normalizer
            key = f"timers/{name}-(s)"

    def log(self, names, normalizer=1.0, reset=True):
        """Log a group of timers."""
        assert normalizer > 0.0
        string = "time (ms)"
        for name in names:
            elapsed_time = self.timers[name].elapsed(reset=reset) * 1000.0 / normalizer
            string += f" | {name}: {elapsed_time:.2f}"
        if torch.distributed.is_initialized():
            if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1):
                print(string, flush=True)
        else:
            print(string, flush=True)


class GlobalMemoryBuffer:
    """Global buffer to avoid dynamic memory allocations.
    Caller should ensure that buffers of the same name
    are not used concurrently."""

    def __init__(self):
        self.buffer = {}

    def get_tensor(self, tensor_shape, dtype, name):
        required_len = reduce(operator.mul, tensor_shape, 1)
        if (
            self.buffer.get((name, dtype), None) is None
            or self.buffer[(name, dtype)].numel() < required_len
        ):
            self.buffer[(name, dtype)] = torch.empty(
                required_len,
                dtype=dtype,
                device=torch.cuda.current_device(),
                requires_grad=False,
            )

        return self.buffer[(name, dtype)][0:required_len].view(*tensor_shape)

megatron/global_vars.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Megatron initialization."""

import argparse
import git
import random
import os
import sys
import time

import numpy as np
import torch
from datetime import timedelta

from megatron import fused_kernels
from megatron import get_adlr_autoresume
from megatron import get_args
from megatron import get_tensorboard_writer
from megatron import mpu
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.mpu import (
    set_tensor_model_parallel_rank,
    set_tensor_model_parallel_world_size,
)
from megatron.model.transformer import bias_dropout_add_fused_train
from megatron.model.fused_bias_gelu import bias_gelu
import requests
import json

ENDPOINT = "http://experiments.research.adept.ai/run/slurm_job_slack_monitor"

# This is needed here to avoid a circular dependency on the main version
def _local_print_rank_0(message):
    """If distributed is initialized, print only on rank 0."""
    if torch.distributed.is_initialized():
        if torch.distributed.get_rank() == 0:
            print(message, flush=True)
    else:
        print(message, flush=True)


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."

    args = parse_args(extra_args_provider, ignore_unknown_args=ignore_unknown_args)
    if not args.untie_embeddings:
        # Set untie-embeddings to None by default so it's not checked when a model checkpoint which doesn't use this param, is loaded.
        # Note: We might need to do this for a lot of other args as well.
        args.untie_embeddings = None

    # Initialize basic torch distributed
    _initialize_torch_distributed(args)

    if args.load:
        # if a checkpoint is specified, we always load arguments from it
        args = load_args_from_checkpoint(args)
    args = validate_args(args, args_defaults)

    if args.deterministic:
        torch.use_deterministic_algorithms(mode=True)

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

    # Initialize megatron's customizations to the distributed setup (including parallelism)
    _initialize_megatron_distributed()


def _initialize_torch_distributed(args: argparse.Namespace):
    """Initialize torch.distributed and mpu."""

    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
    if args.cross_block_networking or args.force_socket_networking:
        os.environ["NCCL_NET"] = "Socket"
    else:
        os.environ["NCCL_NET"] = "IB"
    torch.distributed.init_process_group(
        backend=args.default_backend,
        world_size=args.world_size,
        rank=args.rank,
        timeout=timedelta(minutes=args.distributed_comms_timeout),
    )

    mpu.force_communicator_creation(
        rank=torch.distributed.get_rank(),
        world_size=torch.distributed.get_world_size(),
        all_reduce=True,
        all_gather=True,
        barrier=True,
    )
    _local_print_rank_0("_initialize_torch_distributed init")


def _initialize_megatron_distributed():
    """Initialize distributed."""
    # Parse arguments

    # torch.distributed initialization
    args = get_args()

    # Set the tensor model-parallel, pipeline model-parallel, and
    # data-parallel communicators.
    if torch.cuda.device_count() > 0:
        if mpu.model_parallel_is_initialized():
            _local_print_rank_0(
                "mpu.initialize_model_parallel skipped, model parallel is already initialized"
            )
        else:
            mpu.initialize_model_parallel(
                tensor_model_parallel_size_=args.tensor_model_parallel_size,
                pipeline_model_parallel_size_=args.pipeline_model_parallel_size,
                virtual_pipeline_model_parallel_size_=args.virtual_pipeline_model_parallel_size,
                pipeline_model_parallel_split_rank_=args.pipeline_model_parallel_split_rank,
                default_backend=args.default_backend,
                p2p_backend="nccl",
                cross_block_networking=args.cross_block_networking,
                share_word_embeddings=not args.untie_embeddings,
                force_socket_networking=args.force_socket_networking,
            )

            _local_print_rank_0("mpu.initialize_model_parallel init")

    # 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)

    # Autoresume.
    _init_autoresume()

    # Compile dependencies.
    _compile_dependencies()
    _local_print_rank_0("_compile_dependencies init")


def _compile_dependencies():

    args = get_args()

    # ==================
    # Load fused kernels
    # ==================

    compilation_succeeded = torch.tensor([0]).cuda()  # Communicate success/failure
    # Always build on rank zero first.
    assert (
        args.force_build_fused_kernels == False
    ), "force_build_fused_kernels is not supported anymore."
    if torch.distributed.get_rank() == 0:
        start_time = time.time()
        try:
            fused_kernels.load(args.force_build_fused_kernels)
            compilation_succeeded[0] = 1
            torch.distributed.broadcast(compilation_succeeded, 0)
        except fused_kernels.CompilationTimeoutError as e:
            print("ERROR", e, flush=True)
            torch.distributed.broadcast(compilation_succeeded, 0)
            sys.exit(1)
    else:
        torch.distributed.broadcast(compilation_succeeded, 0)
        if compilation_succeeded[0] == 0:
            print(
                "ERROR: fused kernel compilation on rank 0 timed out, exiting...",
                flush=True,
            )
            sys.exit(1)
        fused_kernels.load(args.force_build_fused_kernels)

    # Simple barrier to make sure all ranks have passed the
    # compilation phase successfully before moving on to the
    # rest of the program. We think this might ensure that
    # the lock is released.
    torch.distributed.barrier()
    if torch.distributed.get_rank() == 0:
        print(
            ">>> done with compiling and loading fused kernels. "
            "Compilation time: {:.3f} seconds".format(time.time() - start_time),
            flush=True,
        )


def _init_autoresume():
    """Set autoresume start time."""
    autoresume = get_adlr_autoresume()
    if autoresume:
        torch.distributed.barrier()
        autoresume.init()
        torch.distributed.barrier()


def _set_random_seed(seed_, data_parallel_random_init=False):
    """Set random seed for reproducability."""
    if seed_ is not None and seed_ > 0:
        # 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())
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        if torch.cuda.device_count() > 0:
            mpu.model_parallel_cuda_manual_seed(seed)
    else:
        raise ValueError("Seed ({}) should be a positive integer.".format(seed_))


def set_jit_fusion_options():
    """Set PyTorch JIT layer fusion options."""
    # flags required to enable jit fusion kernels
    TORCH_MAJOR = int(torch.__version__.split(".")[0])
    TORCH_MINOR = int(torch.__version__.split(".")[1])
    if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10):
        # nvfuser
        torch._C._jit_set_profiling_executor(True)
        torch._C._jit_set_profiling_mode(True)
        torch._C._jit_override_can_fuse_on_cpu(False)
        torch._C._jit_override_can_fuse_on_gpu(False)
        torch._C._jit_set_texpr_fuser_enabled(False)
        torch._C._jit_set_nvfuser_enabled(True)
        torch._C._debug_set_autodiff_subgraph_inlining(False)
    else:
        # legacy pytorch fuser
        torch._C._jit_set_profiling_mode(False)
        torch._C._jit_set_profiling_executor(False)
        torch._C._jit_override_can_fuse_on_cpu(True)
        torch._C._jit_override_can_fuse_on_gpu(True)

megatron/initialize.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import torch


# A dictionary of all the memory buffers allocated.
_MEM_BUFFS = dict()


def allocate_mem_buff(name, numel, dtype, track_usage):
    """Allocate a memory buffer."""
    assert name not in _MEM_BUFFS, "memory buffer {} already allocated.".format(name)
    _MEM_BUFFS[name] = MemoryBuffer(name, numel, dtype, track_usage)
    return _MEM_BUFFS[name]


class MemoryBuffer:
    """Contiguous memory buffer.
    Allocate a contiguous memory of type `dtype` and size `numel`. It is
    used to reduce memory fragmentation.

    Usage: After the allocation, the `_start` index is set tot the first
           index of the memory. A memory chunk starting from `_start` index
           can be `allocated` for an input tensor, with the elements of the
           tensor being coppied. The buffer can be reused by resetting the
           `_start` index.

    """

    def __init__(self, name, numel, dtype, track_usage):
        if torch.distributed.get_rank() == 0:
            element_size = torch.tensor([], dtype=dtype).element_size()
            print(
                "> building the {} memory buffer with {} num elements "
                "and {} dtype ({:.1f} MB)...".format(name, numel, dtype, numel * element_size / 1024 / 1024),
                flush=True,
            )
        self.name = name
        self.numel = numel
        self.dtype = dtype
        self.data = torch.empty(self.numel, dtype=self.dtype, device=torch.cuda.current_device(), requires_grad=False)

        # Index tracking the start of the free memory.
        self._start = 0

        # Values used for tracking usage.
        self.track_usage = track_usage
        if self.track_usage:
            self.in_use_value = 0.0
            self.total_value = 0.0

    def reset(self):
        """Reset the buffer start index to the beginning of the buffer."""
        self._start = 0

    def is_in_use(self):
        """Whether the current buffer hold on to any memory."""
        return self._start > 0

    def numel_in_use(self):
        """Return number of elements in use."""
        return self._start

    def add(self, tensor):
        """Allocate a chunk of memory from the buffer to tensor and copy
        the values."""
        assert tensor.dtype == self.dtype, "Input tensor type {} different from buffer type {}".format(
            tensor.dtype, self.dtype
        )
        # Number of elements of the input tensor.
        tensor_numel = torch.numel(tensor)
        new_start = self._start + tensor_numel
        assert new_start <= self.numel, "Not enough memory left in the buffer ({} > {})".format(
            tensor_numel, self.numel - self._start
        )
        # New tensor is a view into the memory.
        new_tensor = self.data[self._start : new_start]
        self._start = new_start
        new_tensor = new_tensor.view(tensor.shape)
        new_tensor.copy_(tensor)
        # Return a pointer to the new tensor.
        return new_tensor

    def get_data(self):
        """Return the data currently in use."""
        if self.track_usage:
            self.in_use_value += float(self._start)
            self.total_value += float(self.numel)
        return self.data[: self._start]

    def print_average_usage(self):
        """Print memory usage average over time. We would like this value
        to be as high as possible."""
        assert self.track_usage, "You need to enable track usage."
        if torch.distributed.get_rank() == 0:
            print(
                " > usage of {} memory buffer: {:.2f} %".format(
                    self.name, self.in_use_value * 100.0 / self.total_value
                ),
                flush=True,
            )

megatron/memory.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Megatron number of micro-batches calculators."""

from abc import ABC
from abc import abstractmethod


def build_num_microbatches_calculator(args):

    # Constant num micro-batches.
    if args.rampup_batch_size is None:
        num_microbatches_calculator = ConstantNumMicroBatches(
            args.global_batch_size, args.micro_batch_size, args.data_parallel_size
        )
        if args.rank == 0:
            print(
                "setting number of micro-batches to constant {}".format(num_microbatches_calculator.get()), flush=True
            )

    else:
        assert len(args.rampup_batch_size) == 3, (
            "expected the following "
            "format: --rampup-batch-size <start batch size> "
            "<batch size incerement> <ramp-up samples>"
        )
        start_batch_size = int(args.rampup_batch_size[0])
        batch_size_increment = int(args.rampup_batch_size[1])
        ramup_samples = int(args.rampup_batch_size[2])
        if args.rank == 0:
            print(
                "will use batch size rampup starting from global batch "
                "size {} to global batch size {} with batch size increments "
                "{} over {} samples.".format(
                    start_batch_size, args.global_batch_size, batch_size_increment, ramup_samples
                ),
                flush=True,
            )
        num_microbatches_calculator = RampupBatchsizeNumMicroBatches(
            start_batch_size,
            batch_size_increment,
            ramup_samples,
            args.global_batch_size,
            args.micro_batch_size,
            args.data_parallel_size,
        )

    return num_microbatches_calculator


class NumMicroBatchesCalculator(ABC):
    def __init__(self):
        self.num_micro_batches = None
        self.current_global_batch_size = None

    def get(self):
        return self.num_micro_batches

    def get_current_global_batch_size(self):
        return self.current_global_batch_size

    @abstractmethod
    def update(self, consumed_samples, consistency_check):
        pass


class ConstantNumMicroBatches(NumMicroBatchesCalculator):
    def __init__(self, global_batch_size, micro_batch_size, data_parallel_size):
        micro_batch_times_data_parallel = micro_batch_size * data_parallel_size
        assert (
            global_batch_size % micro_batch_times_data_parallel == 0
        ), "global batch size ({}) is not divisible by micro batch size ({})" " times data parallel size ({})".format(
            global_batch_size, micro_batch_size, data_parallel_size
        )
        self.num_micro_batches = global_batch_size // micro_batch_times_data_parallel
        assert self.num_micro_batches >= 1
        self.current_global_batch_size = global_batch_size

    def update(self, consumed_samples, consistency_check):
        pass


class RampupBatchsizeNumMicroBatches(NumMicroBatchesCalculator):
    def __init__(
        self,
        start_batch_size,
        batch_size_increment,
        ramup_samples,
        global_batch_size,
        micro_batch_size,
        data_parallel_size,
    ):
        """Batch size ramp up.
        Over
          steps = (global-batch-size - start-batch-size) / batch_size_increment
        increment batch size from start-batch-size to global-batch-size using
          rampup-samples / steps
        samples.
        Arguments:
            start_batch_size: global batch size to start with
            batch_size_increment: global batch size increments
            ramup_samples: number of samples to use ramp up global
               batch size from `start_batch_size` to `global_batch_size`
            global_batch_size: global batch size post rampup
            micro_batch_size: micro batch size
            data_parallel_size: data parallel size.
        """

        self.micro_batch_size = micro_batch_size
        self.data_parallel_size = data_parallel_size
        self.micro_batch_times_data_parallel_size = self.micro_batch_size * self.data_parallel_size
        assert self.micro_batch_times_data_parallel_size > 0

        assert start_batch_size > 0
        self.start_batch_size = start_batch_size

        assert global_batch_size > 0
        self.global_batch_size = global_batch_size
        diff_batch_size = self.global_batch_size - self.start_batch_size
        assert diff_batch_size >= 0
        assert batch_size_increment > 0
        self.batch_size_increment = batch_size_increment
        assert diff_batch_size % batch_size_increment == 0, (
            "expected "
            "global batch size interval ({}) to be divisible by global batch "
            "size increment ({})".format(diff_batch_size, batch_size_increment)
        )

        num_increments = diff_batch_size // self.batch_size_increment
        self.ramup_samples = ramup_samples
        assert self.ramup_samples >= 0
        self.rampup_samples_per_increment = self.ramup_samples / num_increments

        # Initialize number of microbatches.
        self.update(0, False)

    def update(self, consumed_samples, consistency_check):

        if consumed_samples > self.ramup_samples:
            self.current_global_batch_size = self.global_batch_size
        else:
            steps = int(consumed_samples / self.rampup_samples_per_increment)
            self.current_global_batch_size = self.start_batch_size + steps * self.batch_size_increment
            assert self.current_global_batch_size <= self.global_batch_size

        if consistency_check:
            assert self.current_global_batch_size % self.micro_batch_times_data_parallel_size == 0, (
                "current global "
                "batch size ({}) is not divisible by micro-batch-size ({}) times"
                "data parallel size ({})".format(
                    self.current_global_batch_size, self.micro_batch_size, self.data_parallel_size
                )
            )
        self.num_micro_batches = self.current_global_batch_size // self.micro_batch_times_data_parallel_size

megatron/microbatches.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .fused_layer_norm import MixedFusedLayerNorm as LayerNorm

from .distributed import DistributedDataParallel
from .gpt_model import GPTModel
from .language_model import get_language_model
from .module import Float16Module
from .enums import ModelType

megatron/model/__init__.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from abc import ABC
from abc import abstractmethod
import math
import collections

import torch
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors

from megatron import mpu, get_args
from .module import MegatronModule


class MemoryBuffer:
    def __init__(self, numel, numel_padded, dtype):
        self.numel = numel
        self.numel_padded = numel_padded
        self.dtype = dtype
        self.data = torch.zeros(
            self.numel_padded, dtype=self.dtype, device=torch.cuda.current_device(), requires_grad=False
        )

    def zero(self):
        """Reset the buffer to zero."""
        self.data.zero_()

    def get(self, shape, start_index):
        """Return a tensor with the input `shape` as a view into the
        1-D data starting at `start_index`."""
        end_index = start_index + shape.numel()
        assert end_index <= self.numel, "requested tensor is out of the buffer range."
        buffer_tensor = self.data[start_index:end_index]
        buffer_tensor = buffer_tensor.view(shape)
        return buffer_tensor


class DistributedDataParallelBase(MegatronModule, ABC):
    """Abstract class for DDP."""

    def __init__(self, module):
        super(DistributedDataParallelBase, self).__init__()
        # Keep a pointer to the model.
        self.module = module

    @abstractmethod
    def allreduce_gradients(self):
        pass

    def forward(self, *inputs, **kwargs):
        return self.module(*inputs, **kwargs)

    def state_dict(self, destination=None, prefix="", keep_vars=False):
        return self.module.state_dict(destination, prefix, keep_vars)

    def state_dict_for_save_checkpoint(self, destination=None, prefix="", keep_vars=False):
        return self.module.state_dict_for_save_checkpoint(destination, prefix, keep_vars)

    def load_state_dict(self, state_dict, strict=True):
        self.module.load_state_dict(state_dict, strict=strict)


class DistributedDataParallel(DistributedDataParallelBase):
    """DDP with contiguous buffers options to storre and accumulate gradients.
    This class:
        - has the potential to reduce memory fragmentation.
        - provides the option to do the gradient accumulation
          in a type other than the params type (for example fp32)

    Arguments:
        module: input model.
        accumulate_allreduce_grads_in_fp32: if true do the gradient accumulation
            and the gradient all-reduce all in in float32. If this option is
            true, we require `use_contiguous_buffers` to be true too.
        use_contiguous_buffers: if true, use a contiguous buffer to store the
            gradients.
    """

    def __init__(self, module, accumulate_allreduce_grads_in_fp32, use_contiguous_buffers):

        super(DistributedDataParallel, self).__init__(module)

        self.accumulate_allreduce_grads_in_fp32 = accumulate_allreduce_grads_in_fp32
        self.use_contiguous_buffers = use_contiguous_buffers
        # If we are using fp32-accumulate-allreduce explicitly
        # this means we need main grads in a continous buffer.
        if self.accumulate_allreduce_grads_in_fp32:
            assert self.use_contiguous_buffers

        # ===================================
        # Rest of this part applies only to
        # the case we use continuous buffers.
        # ===================================
        self._grad_buffers = None
        self._grad_buffer_param_index_map = None
        if self.use_contiguous_buffers:
            self._grad_buffers = {}
            self._grad_buffer_param_index_map = {}
            data_parallel_world_size = mpu.get_data_parallel_world_size()

            # Simple function to define buffer type.
            def _get_buffer_type(param):
                return torch.float if self.accumulate_allreduce_grads_in_fp32 else param.dtype

            # First calculate total number of elements per type.
            type_num_elements = {}
            for param in self.module.parameters():
                if param.requires_grad:
                    dtype = _get_buffer_type(param)
                    type_num_elements[dtype] = type_num_elements.get(dtype, 0) + param.data.nelement()

            # Allocate the buffer.
            for dtype, num_elements in type_num_elements.items():

                # If using distributed optimizer, pad memory buffer to be
                # multiple of data_parallel_world_size. (This padding is done
                # due to a constraint with the reduce_scatter op, which requires
                # all tensors have equal size. See: optimizer.py.)
                num_elements_padded = data_parallel_world_size * int(math.ceil(num_elements / data_parallel_world_size))

                # Allocate grad buffer.
                self._grad_buffers[dtype] = MemoryBuffer(num_elements, num_elements_padded, dtype)

            # Assume the back prop order is reverse the params order,
            # store the start index for the gradients.
            for param in self.module.parameters():
                if param.requires_grad:
                    dtype = _get_buffer_type(param)
                    type_num_elements[dtype] -= param.data.nelement()
                    param.main_grad = self._grad_buffers[dtype].get(param.data.shape, type_num_elements[dtype])
                    if dtype not in self._grad_buffer_param_index_map:
                        self._grad_buffer_param_index_map[dtype] = {}
                    self._grad_buffer_param_index_map[dtype][param] = (
                        type_num_elements[dtype],
                        type_num_elements[dtype] + param.data.nelement(),
                    )

            # Backward hook.
            # Accumalation function for the gradients. We need
            # to store them so they don't go out of scope.
            self.grad_accs = []
            # Loop over all the parameters in the model.
            for param in self.module.parameters():
                if param.requires_grad:
                    # Expand so we get access to grad_fn.
                    param_tmp = param.expand_as(param)
                    # Get the gradient accumulator functtion.
                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
                    grad_acc.register_hook(self._make_param_hook(param))
                    self.grad_accs.append(grad_acc)

    def _make_param_hook(self, param):
        """Create the all-reduce hook for backprop."""
        # Hook used for back-prop.
        def param_hook(*unused):
            # Add the gradient to the buffer.
            if param.grad is not None:
                # The gradient function of linear layers is fused with GEMMs
                param.main_grad.add_(param.grad.data)
                # Now we can deallocate grad memory.
                param.grad = None

        return param_hook

    def zero_grad_buffer(self):
        """Set the grad buffer data to zero. Needs to be called at the
        begining of each iteration."""
        assert self._grad_buffers is not None, "buffers are not initialized."
        for _, buffer_ in self._grad_buffers.items():
            buffer_.zero()

    def broadcast_params(self):
        for param in self.module.parameters():
            torch.distributed.broadcast(
                param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group()
            )

    def allreduce_gradients(self):
        """Reduce gradients across data parallel ranks."""
        # If we have buffers, simply reduce the data in the buffer.
        if self._grad_buffers is not None:
            args = get_args()
            for _, buffer_ in self._grad_buffers.items():
                if args.nan_grad_check:
                    buffer_has_nan = buffer_.data.isnan().any()
                    if buffer_has_nan:
                        print(socket.gethostname(), "has a NaN!")
                torch.distributed.all_reduce(buffer_.data, group=mpu.get_data_parallel_group())
        else:
            # Otherwise, bucketize and all-reduce
            buckets = collections.defaultdict(list)
            # Pack the buckets.
            for param in self.module.parameters():
                if param.requires_grad and param.grad is not None:
                    tp = param.data.type()
                    buckets[tp].append(param)
                    param.main_grad = param.grad

            # For each bucket, all-reduce and copy all-reduced grads.
            for tp, bucket in buckets.items():
                grads = [param.grad.data for param in bucket]
                coalesced = _flatten_dense_tensors(grads)
                torch.distributed.all_reduce(coalesced, group=mpu.get_data_parallel_group())
                for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):
                    buf.copy_(synced)

        # Sync tensor parallel replicated params.
        # These parameters should be the same across tensor parallel ranks because they are initialized to be the same,
        # and the gradients going into them are the same. However, during the data parallel all-reduce, the gradients
        # for the different parameters are added together in different orders. Because floating point addition is not
        # associative, the resulting parameters can be slightly different.
        # To fix this, we just broadcast the result that ended up in the first tensor parallel rank to the other ones.
        tp_replicated_buckets = collections.defaultdict(list)
        # Pack the buckets.
        for param in self.module.parameters():
            if mpu.param_is_tensor_parallel_replicated(param):
                tp_replicated_buckets[param.main_grad.dtype].append(param.main_grad)

        # For each bucket, broadcast.
        for bucket in tp_replicated_buckets.values():
            coalesced = _flatten_dense_tensors(bucket)
            torch.distributed.broadcast(
                coalesced, src=mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group()
            )
            for buf, synced in zip(bucket, _unflatten_dense_tensors(coalesced, bucket)):
                buf.copy_(synced)

megatron/model/distributed.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import enum


class ModelType(enum.Enum):
    encoder_or_decoder = 1
    encoder_and_decoder = 2


class LayerType(enum.Enum):
    encoder = 1
    decoder = 2


class AttnType(enum.Enum):
    self_attn = 1
    cross_attn = 2


class AttnMaskType(enum.Enum):
    padding = 1
    causal = 2

megatron/model/enums.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.

import torch


###### BIAS SQRELU FUSION/ NO AUTOGRAD ################


@torch.jit.script
def bias_sqrelu(bias, y):
    x = bias + y
    relud_x = torch.relu(x)
    return relud_x * relud_x


@torch.jit.script
def bias_sqrelu_back(g, bias, y):
    x = bias + y
    return g * 2 * torch.relu(x)


class SqReLUFunction(torch.autograd.Function):
    @staticmethod
    # bias is an optional argument
    def forward(ctx, input, bias):
        ctx.save_for_backward(input, bias)
        return bias_sqrelu(bias, input)

    @staticmethod
    def backward(ctx, grad_output):
        input, bias = ctx.saved_tensors
        tmp = bias_sqrelu_back(grad_output, bias, input)
        return tmp, tmp


bias_sqrelu_impl = SqReLUFunction.apply

megatron/model/fused_bias_sqrelu.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""This code is copied fron NVIDIA apex:
      https://github.com/NVIDIA/apex
   with some changes. """

import numbers
import torch
from torch.nn.parameter import Parameter
from torch.nn import init
import importlib

from megatron.mpu import make_viewless_tensor
from megatron.mpu import set_tensor_model_parallel_attributes

try:
    from apex.contrib.layer_norm.layer_norm import FastLayerNormFN

    HAVE_PERSIST_LAYER_NORM = True
except:
    HAVE_PERSIST_LAYER_NORM = False

global fused_mix_prec_layer_norm_cuda
fused_mix_prec_layer_norm_cuda = None


class FusedLayerNormAffineFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, bias, normalized_shape, eps):

        ctx.normalized_shape = normalized_shape
        ctx.eps = eps
        input_ = input.contiguous()
        weight_ = weight.contiguous()
        bias_ = bias.contiguous()
        output, mean, invvar = fused_mix_prec_layer_norm_cuda.forward_affine(
            input_, ctx.normalized_shape, weight_, bias_, ctx.eps
        )
        ctx.save_for_backward(input_, weight_, bias_, mean, invvar)

        return output

    @staticmethod
    def backward(ctx, grad_output):

        input_, weight_, bias_, mean, invvar = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None
        grad_input, grad_weight, grad_bias = fused_mix_prec_layer_norm_cuda.backward_affine(
            grad_output.contiguous(), mean, invvar, input_, ctx.normalized_shape, weight_, bias_, ctx.eps
        )

        return grad_input, grad_weight, grad_bias, None, None


class MixedFusedLayerNorm(torch.nn.Module):
    def __init__(
        self,
        normalized_shape,
        eps=1e-5,
        no_persist_layer_norm=True,
        sequence_parallel=False,
        is_tensor_parallel_unique=False,
        reduce_tensor_parallel_grads=False,
    ):
        super(MixedFusedLayerNorm, self).__init__()

        global fused_mix_prec_layer_norm_cuda
        fused_mix_prec_layer_norm_cuda = importlib.import_module(
            "megatron_fused_kernels.fused_mix_prec_layer_norm_cuda"
        )

        # List of hiddens sizes supported in the persistent layer norm kernel
        # If the hidden size is not supported, fall back to the non-persistent
        # kernel.
        persist_ln_hidden_sizes = [
            1024,
            1536,
            2048,
            2304,
            3072,
            3840,
            4096,
            5120,
            6144,
            8192,
            10240,
            12288,
            12800,
            15360,
            16384,
            18432,
            20480,
            24576,
            25600,
            30720,
            32768,
            40960,
            49152,
            65536,
        ]
        if normalized_shape not in persist_ln_hidden_sizes or not HAVE_PERSIST_LAYER_NORM:
            no_persist_layer_norm = True

        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        self.normalized_shape = torch.Size(normalized_shape)
        self.eps = eps

        self.weight = Parameter(torch.Tensor(*normalized_shape))
        set_tensor_model_parallel_attributes(
            self.weight,
            is_tensor_parallel_unique=is_tensor_parallel_unique,
            reduce_tensor_parallel_grads=reduce_tensor_parallel_grads,
        )

        self.bias = Parameter(torch.Tensor(*normalized_shape))
        set_tensor_model_parallel_attributes(
            self.bias,
            is_tensor_parallel_unique=is_tensor_parallel_unique,
            reduce_tensor_parallel_grads=reduce_tensor_parallel_grads,
        )

        self.reset_parameters()
        self.no_persist_layer_norm = no_persist_layer_norm
        self.sequence_parallel = sequence_parallel

        # set sequence parallelism flag on weight and bias parameters
        setattr(self.weight, "sequence_parallel", self.sequence_parallel)
        setattr(self.bias, "sequence_parallel", self.sequence_parallel)

    def reset_parameters(self):

        init.ones_(self.weight)
        init.zeros_(self.bias)

    def forward(self, input):
        # Doing this shape check here is important because the two kernels below don't handle shapes the same way,
        # so we need to check that we're passing them what we expect.
        assert len(self.normalized_shape) == 1
        if input.shape[-1] != self.normalized_shape[0]:
            raise ValueError(f"Expected input shape ending in {self.normalized_shape=}, but got {input.shape=}")

        if self.no_persist_layer_norm:
            return FusedLayerNormAffineFunction.apply(input, self.weight, self.bias, self.normalized_shape, self.eps)
        else:
            output = FastLayerNormFN.apply(input, self.weight, self.bias, self.eps)

            # Apex's fast layer norm function outputs a 'view' tensor (i.e., has
            # a populated '_base' field). This will result in schedule.py's
            # deallocate_output_tensor() throwing an error, so a viewless tensor is
            # created to prevent this.
            output = make_viewless_tensor(inp=output, requires_grad=input.requires_grad, keep_graph=True)

            return output

megatron/model/fused_layer_norm.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import torch
import torch.nn as nn
from megatron.model.enums import AttnMaskType


class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
    """
    Fused operation which performs following three operations in sequence
    1. Scale the tensor.
    2. Apply upper triangular mask (typically used in gpt models).
    3. Perform softmax.
    """

    @staticmethod
    def forward(ctx, inputs, scale):
        from megatron_fused_kernels import scaled_upper_triang_masked_softmax_cuda

        scale_t = torch.tensor([scale])
        softmax_results = scaled_upper_triang_masked_softmax_cuda.forward(
            inputs, scale_t[0]
        )

        ctx.save_for_backward(softmax_results, scale_t)
        return softmax_results

    @staticmethod
    def backward(ctx, output_grads):
        from megatron_fused_kernels import scaled_upper_triang_masked_softmax_cuda

        softmax_results, scale_t = ctx.saved_tensors
        input_grads = scaled_upper_triang_masked_softmax_cuda.backward(
            output_grads, softmax_results, scale_t[0]
        )

        return input_grads, None


class ScaledSoftmax(torch.autograd.Function):
    """
    Fused operation which performs following two operations in sequence
    1. Scale the tensor.
    2. Perform softmax.
    """

    @staticmethod
    def forward(ctx, inputs, scale):
        from megatron_fused_kernels import scaled_softmax_cuda

        scale_t = torch.tensor([scale])

        softmax_results = scaled_softmax_cuda.forward(inputs, scale_t[0])
        ctx.save_for_backward(softmax_results, scale_t)
        return softmax_results

    @staticmethod
    def backward(ctx, output_grads):
        from megatron_fused_kernels import scaled_softmax_cuda

        softmax_results, scale_t = ctx.saved_tensors

        input_grads = scaled_softmax_cuda.backward(
            output_grads, softmax_results, scale_t[0]
        )
        return input_grads, None, None


class FusedScaleSoftmax(nn.Module):
    """
    fused operation: scaling + mask + softmax

    Arguments:
        input_in_fp16: flag to indicate if input in fp16 data format.
        input_in_bf16: flag to indicate if input in bf16 data format.
        attn_mask_type: attention mask type (pad or causal)
        scaled_masked_softmax_fusion: flag to indicate user want to use softmax fusion
        softmax_in_fp32: if true, softmax in performed at fp32 precision.
        scale: scaling factor used in input tensor scaling.
    """

    def __init__(
        self,
        input_in_fp16,
        input_in_bf16,
        attn_mask_type,
        scaled_masked_softmax_fusion,
        softmax_in_fp32,
        scale,
    ):
        super(FusedScaleSoftmax, self).__init__()
        self.input_in_fp16 = input_in_fp16
        self.input_in_bf16 = input_in_bf16
        assert not (
            self.input_in_fp16 and self.input_in_bf16
        ), "both fp16 and bf16 flags cannot be active at the same time."
        self.input_in_float16 = self.input_in_fp16 or self.input_in_bf16
        self.attn_mask_type = attn_mask_type
        self.scaled_masked_softmax_fusion = scaled_masked_softmax_fusion
        self.softmax_in_fp32 = softmax_in_fp32
        self.scale = scale

        assert (
            self.scale is None or softmax_in_fp32
        ), "softmax should be in fp32 when scaled"

    def forward(self, input):
        # [b, np, sq, sk]
        assert input.dim() == 4

        if self.is_kernel_available(*input.size()):
            return self.forward_fused_softmax(input)
        else:
            return self.forward_torch_softmax(input)

    def is_kernel_available(self, b, np, sq, sk):
        attn_batches = b * np
        if (
            self.scaled_masked_softmax_fusion  # user want to fuse
            and self.input_in_float16  # input must be fp16
            and 16 < sk <= 8192  # sk must be 16 ~ 8192
            and sq % 4 == 0  # sq must be divisor of 4
            and sk % 4 == 0  # sk must be divisor of 4
            and attn_batches % 4 == 0  # np * b must be divisor of 4
        ):
            if 0 <= sk <= 8192:
                batch_per_block = self.get_batch_per_block(sq, sk, b, np)

                if self.attn_mask_type == AttnMaskType.causal:
                    if attn_batches % batch_per_block == 0:
                        return True
                else:
                    if sq % batch_per_block == 0 and sk % batch_per_block == 0:
                        return True
        return False

    def forward_fused_softmax(self, input):
        b, np, sq, sk = input.size()
        scale = self.scale if self.scale is not None else 1.0

        if self.attn_mask_type == AttnMaskType.causal:
            assert sq == sk, "causal mask is only for self attention"

            # input is 3D tensor (attn_batches, sq, sk)
            input = input.view(-1, sq, sk)
            probs = ScaledUpperTriangMaskedSoftmax.apply(input, scale)
            return probs.view(b, np, sq, sk)
        else:
            # input is 4D tensor (b, np, sq, sk)
            return ScaledSoftmax.apply(input, scale)

    def forward_torch_softmax(self, input):
        b, np, sq, sk = input.size()
        if self.input_in_float16 and self.softmax_in_fp32:
            input = input.float()

        if self.scale is not None:
            input = input * self.scale

        if self.attn_mask_type == AttnMaskType.causal:
            assert sq == sk, "causal mask is only for self attention"
            attention_mask = torch.tril(
                torch.ones((1, input.shape[-1], input.shape[-1]), device=input.device)
            ).view(1, 1, input.shape[-1], input.shape[-1])
            # invert mask so that 0 -> 1 and 1 -> 0
            attention_mask = attention_mask < 0.5

            # assign very low logit to masked entries
            input.masked_fill_(attention_mask, -10000.0)

        probs = torch.nn.Softmax(dim=-1)(input)

        if self.input_in_float16 and self.softmax_in_fp32:
            if self.input_in_fp16:
                probs = probs.half()
            else:
                probs = probs.bfloat16()

        return probs

    @staticmethod
    def get_batch_per_block(sq, sk, b, np):
        from megatron_fused_kernels import scaled_masked_softmax_cuda

        return scaled_masked_softmax_cuda.get_batch_per_block(sq, sk, b, np)

megatron/model/fused_softmax.py
persimmon-ai-labs/adept-inference
# coding=utf-8
#
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""GPT-2 model."""

import torch

from megatron import get_args
from megatron import mpu
from .module import MegatronModule

from .enums import AttnMaskType
from .language_model import parallel_lm_logits
from .language_model import get_language_model
from .utils import init_method_normal
from .utils import scaled_init_method_normal

try:
    from megatron.mpu.cross_entropy_parallel import CrossEntropyLossParallel
    from einops import rearrange
except ImportError:
    CrossEntropyLossParallel = None
    rearrange = None


def post_language_model_processing(
    lm_output, labels, logit_weights, parallel_output, fp16_lm_cross_entropy, use_fast_cross_entropy
):

    # Output. Format [s b h]
    output = parallel_lm_logits(lm_output, logit_weights, parallel_output)
    if labels is None:
        # [s b h] => [b s h]
        return output.transpose(0, 1).contiguous()
    else:
        # [b s] => [s b]
        labels = labels.transpose(0, 1).contiguous()
        if use_fast_cross_entropy:
            loss_fn = CrossEntropyLossParallel(reduction="none", inplace_backward=True)
            loss = rearrange(
                loss_fn(rearrange(output, "s b ... -> (s b) ..."), rearrange(labels, "s b ... -> (s b) ...")),
                "(s b) ... -> s b ...",
                s=output.shape[0],
            )
        elif fp16_lm_cross_entropy:
            assert output.dtype == torch.half
            loss = mpu.vocab_parallel_cross_entropy(output, labels)
        else:
            loss = mpu.vocab_parallel_cross_entropy(output.float(), labels)

        # [s b] => [b, s]
        loss = loss.transpose(0, 1).contiguous()

        argmax = mpu.layers.argmax_parallel(output)
        accuracy = labels == argmax

        accuracy = accuracy.transpose(0, 1).contiguous()
        argmax = argmax.transpose(0, 1).contiguous()

        return {"argmax": argmax, "loss": loss, "accuracy": accuracy}


class GPTModel(MegatronModule):
    """GPT-2 Language model."""

    def __init__(
        self,
        num_tokentypes=0,
        parallel_output=True,
        pre_process=True,
        post_process=True,
        continuous_embed_input_size=None,
    ):
        args = get_args()
        super(GPTModel, self).__init__(share_word_embeddings=not args.untie_embeddings)
        self.parallel_output = parallel_output
        self.pre_process = pre_process
        self.post_process = post_process
        self.fp16_lm_cross_entropy = args.fp16_lm_cross_entropy
        self.share_word_embeddings = not args.untie_embeddings
        self.use_fast_cross_entropy = args.use_fast_cross_entropy
        if self.use_fast_cross_entropy:
            if CrossEntropyLossParallel is None:
                raise ImportError("xentropy CUDA extension is not installed")
            if rearrange is None:
                raise ImportError("einops is not installed")

        self.language_model, self._language_model_key = get_language_model(
            num_tokentypes=num_tokentypes,
            add_pooler=False,
            encoder_attn_mask_type=AttnMaskType.causal,
            init_method=init_method_normal(args.init_method_std),
            scaled_init_method=scaled_init_method_normal(args.init_method_std, args.num_layers),
            pre_process=self.pre_process,
            post_process=self.post_process,
            continuous_embed_input_size=continuous_embed_input_size,
        )

        self.initialize_word_embeddings(init_method_normal)

    def set_input_tensor(self, input_tensor):
        """See megatron.model.transformer.set_input_tensor()"""
        self.language_model.set_input_tensor(input_tensor)

    def forward(
        self,
        input_ids,
        position_ids,
        labels=None,
        tokentype_ids=None,
        inference_params=None,
        lm_logits_mask=None,
    ):
        lm_output = self.language_model(input_ids, position_ids, inference_params=inference_params)

        if lm_logits_mask is not None:
            # lm_logits_mask has shape [s] while lm_output has shape [s, b, h].
            lm_output = lm_output[lm_logits_mask, :, :]

        del tokentype_ids
        if self.post_process:
            return post_language_model_processing(
                lm_output,
                labels,
                self.word_embeddings_weight(),
                self.parallel_output,
                self.fp16_lm_cross_entropy,
                use_fast_cross_entropy=self.use_fast_cross_entropy,
            )
        else:
            return lm_output

    def state_dict_for_save_checkpoint(self, destination=None, prefix="", keep_vars=False):

        state_dict_ = {}
        state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint(
            destination, prefix, keep_vars
        )
        # Save word_embeddings.
        # When it is the last stage of pipelining or if we are not sharing word
        # embeddings, then we need to save the variable to the state dict.
        if self.post_process and (not self.pre_process or not self.share_word_embeddings):
            state_dict_[self._word_embeddings_for_head_key] = self.word_embeddings.state_dict(
                destination, prefix, keep_vars
            )
        return state_dict_

    def load_state_dict(self, state_dict, strict=True):
        """Customized load."""

        # Load word_embeddings.
        if (self.post_process and not self.pre_process) or not self.share_word_embeddings:
            if self._word_embeddings_for_head_key in state_dict:
                self.word_embeddings.load_state_dict(state_dict[self._word_embeddings_for_head_key], strict=strict)
        if self._language_model_key in state_dict:
            state_dict = state_dict[self._language_model_key]
        self.language_model.load_state_dict(state_dict, strict=strict)

megatron/model/gpt_model.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Transformer based language model."""

from typing import List, Optional

import torch

from megatron import get_args, mpu
from megatron.model.enums import AttnMaskType, LayerType
from megatron.model.transformer import ParallelTransformer, ParallelTransformerConfig
from megatron.model.utils import (
    check_shapes,
    get_linear_layer,
    init_method_normal,
    scaled_init_method_normal,
)
from .module import MegatronModule


def parallel_lm_logits(input_, word_embeddings_weight, parallel_output, bias=None):
    """LM logits using word embedding weights."""
    args = get_args()
    # Parallel logits.
    if args.async_tensor_model_parallel_allreduce or args.sequence_parallel:
        input_parallel = input_
        model_parallel = mpu.get_tensor_model_parallel_world_size() > 1
        async_grad_allreduce = (
            args.async_tensor_model_parallel_allreduce and model_parallel and not args.sequence_parallel
        )
    else:
        input_parallel = mpu.reduce_backward_from_tensor_model_parallel_region(input_)
        async_grad_allreduce = False

    # Matrix multiply.
    logits_parallel = mpu.LinearWithGradAccumulationAndAsyncCommunication.apply(
        input_parallel,
        word_embeddings_weight,
        bias,
        args.gradient_accumulation_fusion,
        async_grad_allreduce,
        args.sequence_parallel,
    )
    # Gather if needed.

    if parallel_output:
        return logits_parallel

    return mpu.gather_from_tensor_model_parallel_region(logits_parallel)


def get_language_model(
    num_tokentypes,
    add_pooler,
    encoder_attn_mask_type,
    init_method=None,
    scaled_init_method=None,
    add_encoder=True,
    add_decoder=False,
    decoder_attn_mask_type=AttnMaskType.causal,
    pre_process=True,
    post_process=True,
    continuous_embed_input_size=None,
):
    """Build language model and return along with the key to save."""
    args = get_args()

    if init_method is None:
        init_method = init_method_normal(args.init_method_std)

    if scaled_init_method is None:
        scaled_init_method = scaled_init_method_normal(args.init_method_std, args.num_layers)

    # Language model.
    language_model = TransformerLanguageModel(
        init_method,
        scaled_init_method,
        encoder_attn_mask_type,
        num_tokentypes=num_tokentypes,
        add_encoder=add_encoder,
        add_decoder=add_decoder,
        decoder_attn_mask_type=decoder_attn_mask_type,
        add_pooler=add_pooler,
        pre_process=pre_process,
        post_process=post_process,
        continuous_embed_input_size=continuous_embed_input_size,
    )
    # key used for checkpoints.
    language_model_key = "language_model"

    return language_model, language_model_key


class Embedding(MegatronModule):
    """Language model embeddings.

    Arguments:
        hidden_size: hidden size
        vocab_size: vocabulary size
        max_sequence_length: maximum size of sequence. This
                             is used for positional embedding
        embedding_dropout_prob: dropout probability for embeddings
        init_method: weight initialization method
        num_tokentypes: size of the token-type embeddings. 0 value
                        will ignore this embedding
    """

    def __init__(
        self,
        hidden_size,
        vocab_size,
        max_sequence_length,
        embedding_dropout_prob,
        init_method,
        num_tokentypes=0,
        pos_emb=None,
        continuous_embed_input_size=None,
    ):
        super(Embedding, self).__init__()

        self.hidden_size = hidden_size
        self.init_method = init_method
        self.num_tokentypes = num_tokentypes
        self.pos_emb = pos_emb
        self.continuous_embed_input_size = continuous_embed_input_size

        args = get_args()

        from megatron import print_rank_0  # pylint: disable=import-outside-toplevel

        # Word embeddings (parallel).
        self.word_embeddings = mpu.VocabParallelEmbedding(
            vocab_size, self.hidden_size, init_method=self.init_method, init_rng_key="input_word_embeddings"
        )
        self._word_embeddings_key = "word_embeddings"

        # Position embedding (serial).
        if pos_emb == "sinusoidal" or pos_emb == "learned":
            if pos_emb == "sinusoidal":
                print_rank_0(
                    "WARNING: sinusoidal positional embedding argument is deprecated. It never "
                    "meant sinusoidal and always meant learned. Please use 'learned' in the future."
                )
            self.position_embeddings = mpu.VocabParallelEmbedding(
                max_sequence_length,
                self.hidden_size,
                init_method=self.init_method,
                init_rng_key="input_position_embeddings",
            )
            self._position_embeddings_key = "position_embeddings"
        else:
            self.position_embeddings = None

        self.continuous_embed_linear = None
        if continuous_embed_input_size:
            if continuous_embed_input_size != hidden_size:
                assert False, (
                    f"All our current configs have continuous_embed_input_size == hidden_size, but got "
                    f"{continuous_embed_input_size=} and {hidden_size=}. This is probably a mistake. "
                    f"(Feel free to delete this assert if we're doing this on purpose.)"
                )
                print_rank_0(
                    f"Using linear layer to project from continuous embedding size of {continuous_embed_input_size} to hidden size of {hidden_size}."
                )
                self.continuous_embed_linear = mpu.ParallelLinear(
                    self.continuous_embed_input_size,
                    self.hidden_size,
                    bias=False,
                    init_method=self.init_method,
                    init_rng_key="input_continuous_embed_linear",
                    sequence_parallel=args.sequence_parallel,
                )
                self._continuous_embed_linear_key = "continuous_embed_linear"
            else:
                print_rank_0(
                    f"Skipping linear layer to project from continuous embedding size of {continuous_embed_input_size} to hidden size of {hidden_size}."
                )

        # Token type embedding.
        # Add this as an optional field that can be added through
        # method call so we can load a pretrain model without
        # token types and add them as needed.
        self._tokentype_embeddings_key = "tokentype_embeddings"
        if self.num_tokentypes > 0:
            self.tokentype_embeddings = torch.nn.Embedding(self.num_tokentypes, self.hidden_size)
            # Initialize the token-type embeddings.
            if args.perform_initialization:
                self.init_method(self.tokentype_embeddings.weight)
        else:
            self.tokentype_embeddings = None

        self.fp32_residual_connection = args.fp32_residual_connection
        self.sequence_parallel = args.sequence_parallel
        # Embeddings dropout
        self.embedding_dropout = torch.nn.Dropout(embedding_dropout_prob)

    def zero_parameters(self):
        """Zero out all parameters in embedding."""
        self.word_embeddings.weight.data.fill_(0)
        self.word_embeddings.weight.shared = True
        if self.pos_emb == "sinusoidal" or self.pos_emb == "learned":
            self.position_embeddings.weight.data.fill_(0)
            self.position_embeddings.weight.shared = True
        if self.num_tokentypes > 0:
            self.tokentype_embeddings.weight.data.fill_(0)
            self.tokentype_embeddings.weight.shared = True
        if self.continuous_embed_linear:
            raise ValueError("Zeroing out parameters not supported when continuous_embed_linear is used.")

    def add_tokentype_embeddings(self, num_tokentypes):
        """Add token-type embedding. This function is provided so we can add
        token-type embeddings in case the pretrained model does not have it.
        This allows us to load the model normally and then add this embedding.
        """
        if self.tokentype_embeddings is not None:
            raise Exception("tokentype embeddings is already initialized")
        from megatron import print_rank_0  # pylint: disable=import-outside-toplevel

        print_rank_0("adding embedding for {} tokentypes".format(num_tokentypes))
        self.num_tokentypes = num_tokentypes
        self.tokentype_embeddings = torch.nn.Embedding(num_tokentypes, self.hidden_size)
        # Initialize the token-type embeddings.
        args = get_args()
        self.init_method(self.tokentype_embeddings.weight)

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        tokentype_ids: Optional[torch.Tensor] = None,
        continuous_embeddings: Optional[List[torch.Tensor]] = None,
    ):
        """Embeddings layer forward pass.

        Args:
            input_ids: input token ids. Shape: [batch_size, sequence_length].
            position_ids: Input position ids. Shape: [batch_size, sequence_length].
            tokentype_ids: token type ids. Shape: [batch_size, sequence_length].
            continuous_embeddings: continuous embeddings.
        """
        # Embeddings.
        word_embeddings: torch.Tensor = self.word_embeddings(input_ids)

        assert continuous_embeddings is None

        if self.pos_emb == "sinusoidal" or self.pos_emb == "learned":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings = word_embeddings + position_embeddings
        else:
            embeddings = word_embeddings

        if tokentype_ids is not None:
            assert self.tokentype_embeddings is not None
            embeddings = embeddings + self.tokentype_embeddings(tokentype_ids)
        else:
            assert self.tokentype_embeddings is None

        # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
        embeddings = embeddings.transpose(0, 1).contiguous()

        # If the input flag for fp32 residual connection is set, convert for float.
        if self.fp32_residual_connection:
            embeddings = embeddings.float()

        # Dropout.
        if self.sequence_parallel:
            embeddings = mpu.scatter_to_sequence_parallel_region(embeddings)
            with mpu.get_cuda_rng_tracker().fork():
                embeddings = self.embedding_dropout(embeddings)
        else:
            embeddings = self.embedding_dropout(embeddings)

        return embeddings

    def state_dict_for_save_checkpoint(self, destination=None, prefix="", keep_vars=False):
        """For easy load."""

        state_dict_ = {}
        state_dict_[self._word_embeddings_key] = self.word_embeddings.state_dict(destination, prefix, keep_vars)
        if self.pos_emb == "sinusoidal" or self.pos_emb == "learned":
            state_dict_[self._position_embeddings_key] = self.position_embeddings.state_dict(
                destination, prefix, keep_vars
            )
        if self.num_tokentypes > 0:
            state_dict_[self._tokentype_embeddings_key] = self.tokentype_embeddings.state_dict(
                destination, prefix, keep_vars
            )
        if self.continuous_embed_linear is not None:
            state_dict_[self._continuous_embed_linear_key] = self.continuous_embed_linear.state_dict(
                destination, prefix, keep_vars
            )

        return state_dict_

    def load_state_dict(self, state_dict, strict=True):
        """Customized load."""
        from megatron import print_rank_0  # pylint: disable=import-outside-toplevel

        # Word embedding.
        if self._word_embeddings_key in state_dict:
            state_dict_ = state_dict[self._word_embeddings_key]
            self.word_embeddings.load_state_dict(state_dict_, strict=strict)
        else:
            # for backward compatibility.
            state_dict_ = {}
            for key in state_dict.keys():
                if "word_embeddings" in key:
                    state_dict_[key.split("word_embeddings.")[1]] = state_dict[key]
                    print_rank_0("Load word embeddings in state dict by splitting original key")
            if len(state_dict_.keys()) > 1:
                self.word_embeddings.load_state_dict(state_dict_, strict=strict)

        # Position embedding.
        if hasattr(self, "_position_embeddings_key"):
            if self._position_embeddings_key in state_dict:
                state_dict_ = state_dict[self._position_embeddings_key]
            else:
                # for backward compatibility.
                state_dict_ = {}
                for key in state_dict.keys():
                    if "position_embeddings" in key:
                        state_dict_[key.split("position_embeddings.")[1]] = state_dict[key]
            self.position_embeddings.load_state_dict(state_dict_, strict=strict)

        # Tokentype embedding.
        if self.num_tokentypes > 0:
            state_dict_ = {}
            if self._tokentype_embeddings_key in state_dict:
                state_dict_ = state_dict[self._tokentype_embeddings_key]
            else:
                # for backward compatibility.
                for key in state_dict.keys():
                    if "tokentype_embeddings" in key:
                        state_dict_[key.split("tokentype_embeddings.")[1]] = state_dict[key]
            if len(state_dict_.keys()) > 0:
                self.tokentype_embeddings.load_state_dict(state_dict_, strict=strict)
            else:
                print_rank_0("***WARNING*** expected tokentype embeddings in the " "checkpoint but could not find it")

        # Continuous embedding linear layer.
        if self.continuous_embed_linear:
            self.continuous_embed_linear.load_state_dict(state_dict[self._continuous_embed_linear_key], strict=strict)


class TransformerLanguageModel(MegatronModule):
    """Transformer language model.

    Arguments:
        transformer_hparams: transformer hyperparameters
        vocab_size: vocabulary size
        max_sequence_length: maximum size of sequence. This
                             is used for positional embedding
        embedding_dropout_prob: dropout probability for embeddings
        num_tokentypes: size of the token-type embeddings. 0 value
                        will ignore this embedding
    """

    def __init__(
        self,
        init_method,
        output_layer_init_method,
        encoder_attn_mask_type,
        num_tokentypes=0,
        add_encoder=True,
        add_decoder=False,
        decoder_attn_mask_type=AttnMaskType.causal,
        add_pooler=False,
        pre_process=True,
        post_process=True,
        continuous_embed_input_size=None,
    ):
        super(TransformerLanguageModel, self).__init__()
        args = get_args()

        self.pre_process = pre_process
        self.post_process = post_process
        self.hidden_size = args.hidden_size
        self.num_tokentypes = num_tokentypes
        self.init_method = init_method
        self.add_encoder = add_encoder
        self.encoder_attn_mask_type = encoder_attn_mask_type
        self.add_decoder = add_decoder
        self.decoder_attn_mask_type = decoder_attn_mask_type
        self.add_pooler = add_pooler
        self.encoder_hidden_state = None
        self.pos_emb = args.pos_emb

        # Embeddings.
        if self.pre_process:
            self.embedding = Embedding(
                self.hidden_size,
                args.padded_vocab_size,
                args.max_position_embeddings,
                args.hidden_dropout,
                self.init_method,
                self.num_tokentypes,
                pos_emb=args.pos_emb,
                continuous_embed_input_size=continuous_embed_input_size,
            )
            self._embedding_key = "embedding"

        # Transformer.
        # Encoder (usually set to True, False if part of an encoder-decoder
        # architecture and in encoder-only stage).
        if self.add_encoder:
            self._encoder_key = "encoder"
            self.encoder = ParallelTransformer(
                ParallelTransformerConfig.from_args(args),
                self.init_method,
                output_layer_init_method,
                self_attn_mask_type=self.encoder_attn_mask_type,
                pre_process=self.pre_process,
                post_process=self.post_process,
                module_name=self._encoder_key,
            )
        else:
            self.encoder = None

        # Decoder (usually set to False, True if part of an encoder-decoder
        # architecture and in decoder-only stage).
        if self.add_decoder:
            self._decoder_key = "decoder"
            self.decoder = ParallelTransformer(
                ParallelTransformerConfig.from_args(args),
                self.init_method,
                output_layer_init_method,
                layer_type=LayerType.decoder,
                self_attn_mask_type=self.decoder_attn_mask_type,
                pre_process=self.pre_process,
                post_process=self.post_process,
                module_name=self._decoder_key,
            )
        else:
            self.decoder = None


    def set_input_tensor(self, input_tensor):
        """See megatron.model.transformer.set_input_tensor()"""

        # This is usually handled in schedules.py but some inference code still
        # gives us non-lists or None
        if not isinstance(input_tensor, list):
            input_tensor = [input_tensor]

        if self.add_encoder and self.add_decoder:
            assert (
                len(input_tensor) == 1
            ), "input_tensor should only be length 1 for stage with both encoder and decoder"
            self.encoder.set_input_tensor(input_tensor[0])
        elif self.add_encoder:
            assert len(input_tensor) == 1, "input_tensor should only be length 1 for stage with only encoder"
            self.encoder.set_input_tensor(input_tensor[0])
        elif self.add_decoder:
            if len(input_tensor) == 2:
                self.decoder.set_input_tensor(input_tensor[0])
                self.encoder_hidden_state = input_tensor[1]
            elif len(input_tensor) == 1:
                self.decoder.set_input_tensor(None)
                self.encoder_hidden_state = input_tensor[0]
            else:
                raise Exception("input_tensor must have either length 1 or 2")
        else:
            raise Exception("Stage must have at least either encoder or decoder")

    def forward(
        self,
        enc_input_ids,
        enc_position_ids,
        dec_input_ids=None,
        dec_position_ids=None,
        tokentype_ids=None,
        inference_params=None,
        pooling_sequence_index=0,
        enc_hidden_states=None,
        output_enc_hidden=False,
        continuous_embeddings: Optional[List[torch.Tensor]] = None,
    ):
        args = get_args()

        # Encoder embedding.
        if self.pre_process and self.add_encoder:
            encoder_input = self.embedding(
                enc_input_ids,
                enc_position_ids,
                tokentype_ids=tokentype_ids,
                continuous_embeddings=continuous_embeddings,
            )
        else:
            encoder_input = None

        # Run encoder.
        if enc_hidden_states is None:
            if self.encoder is not None:
                encoder_output = self.encoder(encoder_input, inference_params=inference_params)
            else:
                encoder_output = self.encoder_hidden_state
        else:
            assert False, "T5 models are no longer supported"

        # output_enc_hidden refers to when we just need the encoder's
        # output. For example, it is helpful to compute
        # similarity between two sequences by average pooling
        if not self.add_decoder or output_enc_hidden:
            return encoder_output

        # Decoder embedding.
        if self.pre_process:
            decoder_input = self.embedding(dec_input_ids, dec_position_ids)
        else:
            decoder_input = None

        # Run decoder.
        decoder_output = self.decoder(
            decoder_input,
            encoder_output=encoder_output,
            inference_params=inference_params,
        )

        return decoder_output, encoder_output

    def state_dict_for_save_checkpoint(self, destination=None, prefix="", keep_vars=False):
        """For easy load."""

        state_dict_ = {}
        if self.pre_process:
            state_dict_[self._embedding_key] = self.embedding.state_dict_for_save_checkpoint(
                destination, prefix, keep_vars
            )
        if self.add_encoder:
            state_dict_[self._encoder_key] = self.encoder.state_dict_for_save_checkpoint(destination, prefix, keep_vars)
        if self.post_process:
            if self.add_pooler:
                state_dict_[self._pooler_key] = self.pooler.state_dict_for_save_checkpoint(
                    destination, prefix, keep_vars
                )
        if self.add_decoder:
            state_dict_[self._decoder_key] = self.decoder.state_dict_for_save_checkpoint(destination, prefix, keep_vars)

        return state_dict_

    def load_state_dict(self, state_dict, strict=True):
        """Customized load."""

        args = get_args()
        # Embedding.
        if self.pre_process:
            if self._embedding_key in state_dict:
                state_dict_ = state_dict[self._embedding_key]
            else:
                # for backward compatibility.
                state_dict_ = {}
                for key in state_dict.keys():
                    if "_embeddings" in key:
                        state_dict_[key] = state_dict[key]
            self.embedding.load_state_dict(state_dict_, strict=strict)

        # Encoder.
        if self.add_encoder:
            if self._encoder_key in state_dict:
                state_dict_ = state_dict[self._encoder_key]
            # For backward compatibility.
            elif "transformer" in state_dict:
                state_dict_ = state_dict["transformer"]
            else:
                # For backward compatibility.
                state_dict_ = {}
                for key in state_dict.keys():
                    if "transformer." in key:
                        state_dict_[key.split("transformer.")[1]] = state_dict[key]

            # For backward compatibility.
            state_dict_self_attention = {}
            for key in state_dict_.keys():
                if ".attention." in key:
                    state_dict_self_attention[key.replace(".attention.", ".self_attention.")] = state_dict_[key]
                else:
                    state_dict_self_attention[key] = state_dict_[key]
            state_dict_ = state_dict_self_attention

            # Compatibility for pipeline conversion. Add random layernorm weight and bias to the model
            if args.is_pipeline_conversion:
                if "final_layernorm.weight" not in state_dict_:
                    from megatron import print_rank_0  # pylint: disable=import-outside-toplevel

                    print_rank_0(
                        "Create a random final_layernorm weight and bias tensor when not present in the checkpoint"
                    )
                    state_dict_["final_layernorm.weight"] = torch.rand((self.hidden_size))
                    state_dict_["final_layernorm.bias"] = torch.rand((self.hidden_size))

            self.encoder.load_state_dict(state_dict_, strict=strict)

        # Decoder.
        if self.add_decoder:
            assert "decoder" in state_dict, "could not find data for pooler in the checkpoint"
            self.decoder.load_state_dict(state_dict[self._decoder_key], strict=strict)

megatron/model/language_model.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Megatron Module"""

import argparse
from typing import Any, Callable, Dict, Optional, TypeVar
import torch
from torch.autograd import Variable
from torch.nn.parameter import Parameter

from megatron import get_args
from megatron import mpu


_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor)


def param_is_not_shared(param: Parameter) -> bool:
    return not hasattr(param, "shared") or not param.shared


class MegatronModule(torch.nn.Module):  # type: ignore
    """Megatron specific extensions of torch Module with support
    for pipelining."""

    def __init__(self, share_word_embeddings: bool = True) -> None:
        super().__init__()
        self.share_word_embeddings = share_word_embeddings
        # Initialize the key name for the output word embeddings.
        # This variable is referred at multiple places inside the code,
        # and based on untie_embeddings, it may be present or not.
        # When loading we need to check if they key is present, so this
        # serves as a central place to refer to the variable.
        self._word_embeddings_for_head_key = "word_embeddings_for_head"

    def state_dict_for_save_checkpoint(
        self, destination: Optional[Dict[str, Any]] = None, prefix: str = "", keep_vars: bool = False
    ) -> Dict[str, Any]:
        """Use this function to override the state dict for
        saving checkpoints."""
        state_dict: Dict[str, Any] = self.state_dict(destination, prefix, keep_vars)
        return state_dict

    def word_embeddings_weight(self) -> Parameter:
        """
        This function is responsible for returning the right word_embeddings var
        based on which pipeline stage this is called from and also if we are
        sharing word_embeddings between the input and the output.
        Megatron uses two vars called post_process and pre_process.

        When there is pipelining,
        - post_process is True for last pipeline stage
        - pre_process is True for all stages except the last

        When there is no pipelining,
        - post_process and pre_process is True

        We create a separate var called self.word_embeddings when we are at the
        last pipeline stage or we have share_embeddings False (untie_embeddings is True.)
        """
        args = get_args()

        # When not using pipelining, return the input embedding var if we
        # have share_word_embeddings to be True, and return the new embedding var,
        # word_embeddings, if False.
        if args.pipeline_model_parallel_size == 1:
            if self.share_word_embeddings:
                weight: Parameter = self.language_model.embedding.word_embeddings.weight
                return weight
            else:
                weight = self.word_embeddings.weight
                return weight

        # First condition triggered for last pipeline stage
        if self.post_process and not self.pre_process:
            weight = self.word_embeddings.weight
            return weight
        # Triggered when we are not in the last pipeline stage or when we are
        # sharing word embeddings, which will matter when there is no pipelining.
        elif self.pre_process or self.share_word_embeddings:
            weight = self.language_model.embedding.word_embeddings.weight
            return weight
        # This will trigger for cases when there is no pipelining and
        # no sharing of word_embedding.
        else:
            weight = self.word_embeddings.weight
            return weight

    def initialize_word_embeddings(self, init_method_normal: Callable[..., Any]) -> None:
        args = get_args()

        # This function just initializes the word embeddings in the final stage
        # when we are using pipeline parallelism. Nothing to do if we aren't
        # using pipeline parallelism.

        if args.pipeline_model_parallel_size == 1:
            if not self.share_word_embeddings:
                self.word_embeddings = mpu.VocabParallelEmbedding(
                    args.padded_vocab_size,
                    args.hidden_size,
                    init_method=init_method_normal(args.init_method_std),
                    init_rng_key="output_word_embeddings",
                )
            return

        # Parameters are shared between the word embeddings layers, and the
        # heads at the end of the model. In a pipelined setup with more than
        # one stage, the initial embedding layer and the head are on different
        # workers, so we do the following:
        # 1. Create a second copy of word_embeddings on the last stage, with
        #    initial parameters of 0.0.
        # 2. Do an all-reduce between the first and last stage to ensure that
        #    the two copies of word_embeddings start off with the same
        #    parameter values.
        # 3. In the training loop, before an all-reduce between the grads of
        #    the two word_embeddings layers to ensure that every applied weight
        #    update is the same on both stages.
        if self.post_process and (not self.pre_process or mpu.is_pipeline_last_stage()):
            assert not mpu.is_pipeline_first_stage()
            # set word_embeddings weights to 0 here, then copy first
            # stage's weights using all_reduce below.
            self.word_embeddings = mpu.VocabParallelEmbedding(
                args.padded_vocab_size,
                args.hidden_size,
                init_method=init_method_normal(args.init_method_std),
                init_rng_key="output_word_embeddings",
            )
            if self.share_word_embeddings:
                self.word_embeddings.weight.data.fill_(0)
                self.word_embeddings.weight.shared = True

        # Zero out initial weights for decoder embedding.
        # NOTE: We don't currently support T5 with the interleaved schedule.
        # NOTE: Can we remove the pre_process flag check here?
        if not mpu.is_pipeline_first_stage(ignore_virtual=True) and self.pre_process:
            self.language_model.embedding.zero_parameters()

        if not torch.distributed.is_initialized():
            if not getattr(MegatronModule, "embedding_warning_printed", False):
                print(
                    "WARNING! Distributed processes aren't initialized, so "
                    "word embeddings in the last layer are not initialized. "
                    "If you are just manipulating a model this is fine, but "
                    "this needs to be handled manually. If you are training "
                    "something is definitely wrong."
                )
                MegatronModule.embedding_warning_printed = True
            return

        # Ensure that first and last stages have the same initial parameter
        # values.
        if not args.is_pipeline_conversion:
            if self.share_word_embeddings and mpu.is_rank_in_embedding_group():
                torch.distributed.all_reduce(self.word_embeddings_weight().data, group=mpu.get_embedding_group())

            # Ensure that encoder(first stage) and decoder(split stage) position
            # embeddings have the same initial parameter values
            # NOTE: We don't currently support T5 with the interleaved schedule.
            if mpu.is_rank_in_position_embedding_group() and args.pipeline_model_parallel_split_rank is not None:
                self.language_model.embedding.cuda()
                position_embeddings = self.language_model.embedding.position_embeddings
                torch.distributed.all_reduce(position_embeddings.weight.data, group=mpu.get_position_embedding_group())


T1 = TypeVar("T1")
T2 = TypeVar("T2")


def conversion_helper(val: Any, conversion: Callable[[Any], Any]) -> Any:
    """Apply conversion to val. Recursively apply conversion if `val`
    #is a nested tuple/list structure."""
    if not isinstance(val, (tuple, list)):
        return conversion(val)
    rtn = [conversion_helper(v, conversion) for v in val]
    if isinstance(val, tuple):
        return tuple(rtn)
    return rtn


def fp32_to_float16(val: Any, float16_convertor: Callable[[Any], Any]) -> Any:
    """Convert fp32 `val` to fp16/bf16"""

    def half_conversion(val: Any) -> Any:
        val_typecheck = val
        if isinstance(val_typecheck, (Parameter, Variable)):
            val_typecheck = val.data
        if isinstance(val_typecheck, _FLOAT_TYPES):
            val = float16_convertor(val)
        return val

    return conversion_helper(val, half_conversion)


def float16_to_fp32(val: Any) -> Any:
    """Convert fp16/bf16 `val` to fp32"""

    def float_conversion(val: Any) -> Any:
        val_typecheck = val
        if isinstance(val_typecheck, (Parameter, Variable)):
            val_typecheck = val.data
        if isinstance(val_typecheck, (_BF16_TYPES, _HALF_TYPES)):
            val = val.float()
        return val

    return conversion_helper(val, float_conversion)


class Float16Module(MegatronModule):
    def __init__(self, module: torch.nn.Module, args: argparse.Namespace) -> None:
        super().__init__()

        if args.fp16:
            self.add_module("module", module.half())

            def float16_convertor(val: Any) -> Any:
                return val.half()

        elif args.bf16:
            self.add_module("module", module.bfloat16())

            def float16_convertor(val: Any) -> Any:
                return val.bfloat16()

        else:
            raise Exception("should not be here")

        self.float16_convertor = float16_convertor

    def set_input_tensor(self, input_tensor: torch.Tensor) -> Any:
        return self.module.set_input_tensor(input_tensor)

    def forward(self, *inputs: Any, **kwargs: Any) -> Any:
        if mpu.is_pipeline_first_stage():
            inputs = fp32_to_float16(inputs, self.float16_convertor)
        outputs = self.module(*inputs, **kwargs)
        if mpu.is_pipeline_last_stage():
            outputs = float16_to_fp32(outputs)
        return outputs

    def state_dict(
        self, destination: Optional[Dict[str, Any]] = None, prefix: str = "", keep_vars: bool = False
    ) -> Dict[str, Any]:
        return self.module.state_dict(destination, prefix, keep_vars)  # type: ignore

    def state_dict_for_save_checkpoint(
        self, destination: Optional[Dict[str, Any]] = None, prefix: str = "", keep_vars: bool = False
    ) -> Dict[str, Any]:
        return self.module.state_dict_for_save_checkpoint(destination, prefix, keep_vars)  # type: ignore

    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True) -> None:
        self.module.load_state_dict(state_dict, strict=strict)

megatron/model/module.py
persimmon-ai-labs/adept-inference
# Copyright (c) 2023, ADEPT AI LABS INC.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Positional embeddings modules."""
import torch
import math


class RotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, base=10000, precision=torch.bfloat16):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.seq_len_cached = None
        self.cos_cached = None
        self.sin_cached = None
        self.precision = precision

    def forward(self, x, seq_dim=1, seq_len=None):
        if seq_len is None:
            seq_len = x.shape[seq_dim]
        if seq_len != self.seq_len_cached:
            self.seq_len_cached = seq_len
        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
        if self.precision == torch.bfloat16:
            emb = emb.float()
        self.cos_cached = emb.cos()[:, None, None, :]
        self.sin_cached = emb.sin()[:, None, None, :]
        if self.precision == torch.bfloat16:
            self.cos_cached = self.cos_cached.bfloat16()
            self.sin_cached = self.sin_cached.bfloat16()
        return self.cos_cached, self.sin_cached


def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=x1.ndim - 1)  # dim=-1 triggers a bug in earlier torch versions


@torch.jit.script
def apply_rotary_pos_emb(q, k, cos, sin, offset: int = 0):
    cos, sin = (
        cos[offset : q.shape[0] + offset, ...],
        sin[offset : q.shape[0] + offset, ...],
    )
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)


def apply_rotary_pos_emb_torch(q, k, cos, sin, offset: int = 0):  # jitting fails with bf16
    cos, sin = (
        cos[offset : q.shape[0] + offset, ...],
        sin[offset : q.shape[0] + offset, ...],
    )
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)

megatron/model/positional_embeddings.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Transformer."""
from typing import Tuple

import math
from contextlib import nullcontext
import dataclasses
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from megatron import get_args, get_global_memory_buffer
from megatron.config import Config, DISABLE_PREFIX
from megatron import mpu
from .module import MegatronModule
from megatron.model.enums import AttnMaskType, ModelType, LayerType, AttnType
from megatron.model import LayerNorm
from megatron.model.fused_softmax import FusedScaleSoftmax
from megatron.model.fused_bias_gelu import bias_gelu_impl
from megatron.model.fused_bias_sqrelu import bias_sqrelu_impl
from megatron.model.positional_embeddings import (
    RotaryEmbedding,
    apply_rotary_pos_emb,
    apply_rotary_pos_emb_torch,
)
from megatron.model.utils import openai_gelu, erf_gelu

try:
    from einops import rearrange, repeat
except ImportError:
    rearrange, repeat = None, None

try:
    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func
    from flash_attn.flash_attn_interface import flash_attn_varlen_kvpacked_func
    from flash_attn.layers.rotary import RotaryEmbedding as RotaryEmbeddingFlash
except ImportError:
    flash_attn_varlen_qkvpacked_func = None
    flash_attn_varlen_kvpacked_func = None
    RotaryEmbeddingFlash = None

# FT's single-query attention kernel from flash-attention repo
try:
    import ft_attention
except ImportError:
    ft_attention = None


""" We use the following notation throughout this file:
     h: hidden size
     n: number of attention heads
     p: number of model parallel partitions
     np: n/p
     hp: h/p
     hn: h/n
     b: batch size
     s: sequence length
     l: number of layers
    Transformer takes input of size [s, b, h] and returns a
    tensor of the same size. We use the following arguments:
        hyperparameters: transformer hyperparameters
"""


class DropPath(MegatronModule):
    """Drop paths (Stochastic Depth) per sample
    (when applied in main path of residual blocks).
    """

    def __init__(self, drop_prob=0.0):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, hidden_state):
        if self.drop_prob == 0.0 or not self.training:
            return hidden_state
        keep_prob = 1 - self.drop_prob
        # work with diff dim tensors, not just 2D ConvNets
        shape = (hidden_state.shape[0],) + (1,) * (hidden_state.ndim - 1)
        random_tensor = keep_prob + torch.rand(shape, dtype=hidden_state.dtype, device=hidden_state.device)
        random_tensor.floor_()  # binarize
        output = hidden_state.div(keep_prob) * random_tensor
        return output


@dataclasses.dataclass
class ParallelMLPConfig(Config):
    bias_gelu_fusion: bool
    ffn_hidden_size: int
    hidden_size: int
    onnx_safe: bool = DISABLE_PREFIX()
    openai_gelu: bool = DISABLE_PREFIX()
    sequence_parallel: bool = DISABLE_PREFIX()
    sq_relu: bool


class ParallelMLP(MegatronModule):
    """MLP.

    MLP will take the input with h hidden state, project it to 4*h
    hidden dimension, perform nonlinear transformation, and project the
    state back into h hidden dimension.
    """

    def __init__(
        self,
        parallel_mlp_config: ParallelMLPConfig,
        init_method,
        output_layer_init_method,
        init_rng_key: str,
    ):
        super(ParallelMLP, self).__init__()

        self.sequence_parallel = parallel_mlp_config.sequence_parallel

        # Project to 4h.
        self.dense_h_to_4h = mpu.ColumnParallelLinear(
            parallel_mlp_config.hidden_size,
            parallel_mlp_config.ffn_hidden_size,
            gather_output=False,
            init_method=init_method,
            skip_bias_add=True,
            init_rng_key=f"{init_rng_key}_dense_h_to_4h",
            sequence_parallel=self.sequence_parallel,
        )

        self.bias_gelu_fusion = parallel_mlp_config.bias_gelu_fusion
        self.bias_sqrelu_fusion = parallel_mlp_config.sq_relu
        self.activation_func = F.gelu
        if parallel_mlp_config.openai_gelu:
            self.activation_func = openai_gelu
        elif parallel_mlp_config.onnx_safe:
            self.activation_func = erf_gelu

        # Project back to h.
        self.dense_4h_to_h = mpu.RowParallelLinear(
            parallel_mlp_config.ffn_hidden_size,
            parallel_mlp_config.hidden_size,
            input_is_parallel=True,
            init_method=output_layer_init_method,
            skip_bias_add=True,
            init_rng_key=f"{init_rng_key}_dense_4h_to_h",
            sequence_parallel=self.sequence_parallel,
        )

    def forward(self, hidden_states):

        # [s, b, 4hp]
        intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states)

        if self.bias_gelu_fusion:
            intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)
        elif self.bias_sqrelu_fusion:
            intermediate_parallel = bias_sqrelu_impl(intermediate_parallel, bias_parallel)
        else:
            intermediate_parallel = self.activation_func(intermediate_parallel + bias_parallel)

        # [s, b, h]
        output, output_bias = self.dense_4h_to_h(intermediate_parallel)
        return output, output_bias


@dataclasses.dataclass
class CoreAttentionConfig(Config):
    apply_query_key_layer_scaling: bool = DISABLE_PREFIX()
    attention_dropout: float
    attention_softmax_in_fp32: bool = DISABLE_PREFIX()
    bf16: bool = DISABLE_PREFIX()
    fp16: bool = DISABLE_PREFIX()
    kv_channels: int
    masked_softmax_fusion: bool = DISABLE_PREFIX()
    num_attention_heads: int
    pos_emb: str
    sequence_parallel: bool = DISABLE_PREFIX()
    tensor_model_parallel_size: int = DISABLE_PREFIX()


class CoreAttention(MegatronModule):
    def __init__(
        self,
        core_attention_config: CoreAttentionConfig,
        layer_number,
        attn_mask_type=AttnMaskType.padding,
    ):
        super(CoreAttention, self).__init__()
        self.fp16 = core_attention_config.fp16
        self.bf16 = core_attention_config.bf16

        self.apply_query_key_layer_scaling = core_attention_config.apply_query_key_layer_scaling
        if self.apply_query_key_layer_scaling:
            self.attention_softmax_in_fp32 = True
        self.layer_number = max(1, layer_number)
        self.attn_mask_type = attn_mask_type
        self.sequence_parallel = core_attention_config.sequence_parallel

        projection_size = core_attention_config.kv_channels * core_attention_config.num_attention_heads

        # Per attention head and per partition values.
        world_size = mpu.get_tensor_model_parallel_world_size()
        self.hidden_size_per_partition = mpu.divide(projection_size, world_size)
        self.hidden_size_per_attention_head = mpu.divide(projection_size, core_attention_config.num_attention_heads)
        self.num_attention_heads_per_partition = mpu.divide(core_attention_config.num_attention_heads, world_size)

        coeff = None
        self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
        if self.apply_query_key_layer_scaling:
            coeff = self.layer_number
            self.norm_factor *= coeff

        self.scale_mask_softmax = FusedScaleSoftmax(
            self.fp16,
            self.bf16,
            self.attn_mask_type,
            core_attention_config.masked_softmax_fusion,
            self.attention_softmax_in_fp32,
            coeff,
        )

        # Dropout. Note that for a single iteration, this layer will generate
        # different outputs on different number of parallel partitions but
        # on average it should not be partition dependent.
        self.attention_dropout = torch.nn.Dropout(core_attention_config.attention_dropout)

    def compute_attention_scores(self, query_layer, key_layer, matmul_result, output_size, transpose_sb=False):
        """Compute attention probabilities.

        Args:
            query_layer (torch.Tensor): query tensor
            key_layer (torch.Tensor): key tensor
            matmul_result (torch.Tensor): matmul result
            output_size (tuple): output size
            transpose_sb (bool): transpose the first two dimensions of q and k

        Returns:
            torch.Tensor: attention probabilities
        """
        # Raw attention scores. [b * np, sq, sk]
        matmul_result = torch.baddbmm(
            matmul_result,
            query_layer.transpose(0, 1) if transpose_sb else query_layer,
            key_layer.transpose(0, 1).transpose(1, 2) if transpose_sb else key_layer.transpose(1, 2),
            beta=0.0,
            alpha=(1.0 / self.norm_factor),
        )

        # change view to [b, np, sq, sk]
        attention_scores = matmul_result.view(*output_size)

        # ===========================
        # Attention probs and dropout
        # ===========================

        # attention scores [b, np, sq, sk]
        attention_probs = self.scale_mask_softmax(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.

        if not self.sequence_parallel:
            with mpu.get_cuda_rng_tracker().fork():
                attention_probs = self.attention_dropout(attention_probs)
        else:
            attention_probs = self.attention_dropout(attention_probs)
        return attention_probs

    def forward(self, query_layer, key_layer, value_layer):

        # ===================================
        # Raw attention scores. [b, np, s, s]
        # ===================================
        # [b, np, sq, sk]
        output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))

        # [sq, b, np, hn] -> [sq, b * np, hn]
        query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
        # [sk, b, np, hn] -> [sk, b * np, hn]
        key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)

        # preallocting input tensor: [b * np, sq, sk]
        matmul_input_buffer = get_global_memory_buffer().get_tensor(
            (output_size[0] * output_size[1], output_size[2], output_size[3]), query_layer.dtype, "mpu"
        )

        attention_probs = self.compute_attention_scores(
            query_layer, key_layer, matmul_input_buffer, output_size, transpose_sb=True
        )

        # =========================
        # Context layer. [sq, b, hp]
        # =========================

        # value_layer -> context layer.
        # [sk, b, np, hn] --> [b, np, sq, hn]

        # context layer shape: [b, np, sq, hn]
        output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))

        # change view [sk, b * np, hn]
        value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)

        # change view [b * np, sq, sk]
        attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)

        # matmul: [b * np, sq, hn]
        context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))

        # change view [b, np, sq, hn]
        context_layer = context_layer.view(*output_size)

        # [b, np, sq, hn] --> [sq, b, np, hn]
        context_layer = context_layer.permute(2, 0, 1, 3).contiguous()

        # [sq, b, np, hn] --> [sq, b, hp]
        new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
        context_layer = context_layer.view(*new_context_layer_shape)

        return context_layer


class FlashSelfAttention(nn.Module):
    """Implement the scaled dot product attention with softmax.
    Arguments
    ---------
        softmax_scale: The temperature to use for the softmax attention.
                      (default: 1/sqrt(d_keys) where d_keys is computed at
                      runtime)
        attention_dropout: The dropout rate to apply to the attention
                           (default: 0.0)
    """

    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):
        super().__init__()
        assert flash_attn_varlen_qkvpacked_func is not None, "FlashAttention is not installed"
        assert flash_attn_varlen_kvpacked_func is not None, "FlashAttention is not installed"
        self.causal = causal
        self.softmax_scale = softmax_scale
        self.dropout_p = attention_dropout

    def forward(self, qkv=None, kv=None, q=None):
        """Implements the multihead softmax attention.
        In the case of self-attention, the qkv matrix must be provided.
        In the case of cross-attention, however, the kv and q matrices must
            be provided separately.
        Arguments
        ---------
            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D)
            kv: The tensor containing the key and value. (B, S, 2, H, D)
            q: The tensor containing the query. (B, S, H, D)
        """
        if qkv is not None:
            assert kv is None and q is None
            return self._self_attn(qkv)
        else:
            assert kv is not None and q is not None
            return self._cross_attn(kv, q)

    def _cross_attn(self, kv, q):
        assert kv.dtype in [torch.float16, torch.bfloat16] and q.dtype in [torch.float16, torch.bfloat16]
        assert kv.is_cuda and q.is_cuda

        batch_size, seqlen_kv = kv.shape[0], kv.shape[1]
        kv = rearrange(kv, "b s ... -> (b s) ...")

        seqlen_q = q.shape[1]
        q = rearrange(q, "b s ... -> (b s) ...")

        cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q.device)
        cu_seqlens_kv = torch.arange(
            0, (batch_size + 1) * seqlen_kv, step=seqlen_kv, dtype=torch.int32, device=kv.device
        )
        output = flash_attn_varlen_kvpacked_func(
            q,
            kv,
            cu_seqlens_q,
            cu_seqlens_kv,
            seqlen_q,
            seqlen_kv,
            self.dropout_p if self.training else 0.0,
            softmax_scale=self.softmax_scale,
            causal=self.causal,
        )
        output = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
        return output

    def _self_attn(self, qkv):
        assert qkv.dtype in [torch.float16, torch.bfloat16]
        assert qkv.is_cuda
        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
        qkv = rearrange(qkv, "b s ... -> (b s) ...")
        max_s = seqlen
        cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32, device=qkv.device)
        output = flash_attn_varlen_qkvpacked_func(
            qkv,
            cu_seqlens,
            max_s,
            self.dropout_p if self.training else 0.0,
            softmax_scale=self.softmax_scale,
            causal=self.causal,
        )
        output = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
        return output


@dataclasses.dataclass
class ParallelAttentionConfig(Config):
    attention_dropout: float
    bf16: bool = DISABLE_PREFIX()
    core_attention_config: CoreAttentionConfig
    fp16: bool = DISABLE_PREFIX()
    hidden_size: int
    kv_channels: int
    num_attention_heads: int
    params_dtype: torch.dtype = DISABLE_PREFIX()
    pos_emb: str
    recompute_granularity: str = DISABLE_PREFIX()
    rotary_emb_base: int
    enable_rotary_base_check: bool = DISABLE_PREFIX()
    rotary_pct: float
    use_flash_attn: bool = DISABLE_PREFIX()
    qk_layernorm: str = DISABLE_PREFIX()
    layernorm_epsilon: float = DISABLE_PREFIX()
    no_persist_layer_norm: bool = DISABLE_PREFIX()
    sequence_parallel: bool = DISABLE_PREFIX()


class LayerNormPerAttentionHead(nn.Module):
    def __init__(self, layernorm_builder, num_attention_heads_per_partition):
        super().__init__()
        self.num_attention_heads_per_partition = num_attention_heads_per_partition
        self.layernorms = torch.nn.ModuleList(
            [layernorm_builder() for _ in range(self.num_attention_heads_per_partition)]
        )

    def forward(self, layer_input):
        layer_heads = torch.split(layer_input, 1, dim=2)
        assert (
            len(layer_heads) == len(self.layernorms) == self.num_attention_heads_per_partition
        ), f"{len(layer_heads)}, {len(self.layernorms)}, {self.num_attention_heads_per_partition}"
        normed_heads = [ln(ln_input) for ln, ln_input in zip(self.layernorms, layer_heads)]
        return torch.cat(normed_heads, dim=2)


class ParallelAttention(MegatronModule):
    """Parallel self-attention layer abstract class.

    Self-attention layer takes input with size [s, b, h]
    and returns output of the same size.
    """

    def __init__(
        self,
        parallel_attention_config: ParallelAttentionConfig,
        init_method,
        output_layer_init_method,
        layer_number,
        module_name: str,
        attention_type=AttnType.self_attn,
        attn_mask_type=AttnMaskType.padding,
    ):
        super(ParallelAttention, self).__init__()
        self.module_name = module_name
        self.layer_number = max(1, layer_number)
        self.attention_type = attention_type
        self.attn_mask_type = attn_mask_type
        self.params_dtype = parallel_attention_config.params_dtype
        # Set positional embeddings
        self.pos_emb = parallel_attention_config.pos_emb
        self.rotary_emb_base = parallel_attention_config.rotary_emb_base
        self.enable_rotary_base_check = parallel_attention_config.enable_rotary_base_check
        self.qk_layernorm = parallel_attention_config.qk_layernorm
        self.fp16 = parallel_attention_config.fp16
        self.bf16 = parallel_attention_config.bf16
        self.use_flash_attn = parallel_attention_config.use_flash_attn
        if self.use_flash_attn:
            assert flash_attn_varlen_qkvpacked_func is not None, "FlashAttention is not installed"
            assert flash_attn_varlen_kvpacked_func is not None, "FlashAttention is not installed"
            if rearrange is None:
                raise ImportError("einops is not installed")

        if parallel_attention_config.pos_emb == "rotary" and attention_type == AttnType.cross_attn:
            raise AssertionError("Can't use rotary embeddings and cross attention.")

        projection_size = parallel_attention_config.kv_channels * parallel_attention_config.num_attention_heads

        # Per attention head and per partition values.
        world_size = mpu.get_tensor_model_parallel_world_size()
        self.hidden_size_per_attention_head = mpu.divide(projection_size, parallel_attention_config.num_attention_heads)
        self.num_attention_heads_per_partition = mpu.divide(parallel_attention_config.num_attention_heads, world_size)

        self.sequence_parallel = parallel_attention_config.sequence_parallel

        # Strided linear layer.
        if attention_type == AttnType.self_attn:
            self.query_key_value = mpu.ColumnParallelLinear(
                parallel_attention_config.hidden_size,
                3 * projection_size,
                gather_output=False,
                init_method=init_method,
                init_rng_key=f"{self.module_name}_layer_{layer_number}_query_key_value",
                sequence_parallel=self.sequence_parallel,
            )
        else:
            assert attention_type == AttnType.cross_attn
            self.query = mpu.ColumnParallelLinear(
                parallel_attention_config.hidden_size,
                projection_size,
                gather_output=False,
                init_method=init_method,
                init_rng_key=f"{self.module_name}_layer_{layer_number}_query",
                sequence_parallel=self.sequence_parallel,
            )

            self.key_value = mpu.ColumnParallelLinear(
                parallel_attention_config.hidden_size,
                2 * projection_size,
                gather_output=False,
                init_method=init_method,
                init_rng_key=f"{self.module_name}_layer_{layer_number}_key_value",
                sequence_parallel=self.sequence_parallel,
            )

        # See "Scaling Vision Transformers to 22 Billion Parameters" https://arxiv.org/abs/2302.05442
        if self.qk_layernorm != "none":

            def qk_layernorm_builder():
                # Modes available are:
                # weights_per_layer: QK LayerNorm weights are the same for every attention head on a given layer.
                # weights_per_tensor_parallel_rank_and_layer: QK LayerNorm weights are the same for every attention
                #   head on a given tensor parallel rank. This version was implemented by accident and used in the 4B
                #   flagship model. It is preserved for backwards compatibility, but should probably not be used going
                #   forward.
                # weights_per_attention_head_and_layer: QK LayerNorm weights are unique per attention head.

                # Heads are sharded by tensor parallel ranks. If we're using unique weights per head or tensor
                # parallel rank, then the weights for these LayerNorms will be unique per tensor parallel rank and
                # are marked is_tensor_parallel_unique.
                # Otherwise, they will be shared (via reduce_tensor_parallel_grads) and should not be counted
                # multiple times when calculating params norm for logging or gradient clipping.
                if self.qk_layernorm == "weights_per_layer":
                    is_tensor_parallel_unique = False
                    reduce_tensor_parallel_grads = True
                elif self.qk_layernorm == "weights_per_tensor_parallel_rank_and_layer":
                    is_tensor_parallel_unique = True
                    reduce_tensor_parallel_grads = False
                elif self.qk_layernorm == "weights_per_attention_head_and_layer":
                    is_tensor_parallel_unique = True
                    reduce_tensor_parallel_grads = False
                else:
                    raise ValueError(f"Unknown qk_layernorm mode: {self.qk_layernorm}")

                return LayerNorm(
                    self.hidden_size_per_attention_head,
                    eps=parallel_attention_config.layernorm_epsilon,
                    no_persist_layer_norm=parallel_attention_config.no_persist_layer_norm,
                    sequence_parallel=False,
                    is_tensor_parallel_unique=is_tensor_parallel_unique,
                    reduce_tensor_parallel_grads=reduce_tensor_parallel_grads,
                )

            if self.qk_layernorm == "weights_per_attention_head_and_layer":
                self.q_layernorm = LayerNormPerAttentionHead(
                    layernorm_builder=qk_layernorm_builder,
                    num_attention_heads_per_partition=self.num_attention_heads_per_partition,
                )
                self.k_layernorm = LayerNormPerAttentionHead(
                    layernorm_builder=qk_layernorm_builder,
                    num_attention_heads_per_partition=self.num_attention_heads_per_partition,
                )
            elif self.qk_layernorm in ("weights_per_layer", "weights_per_tensor_parallel_rank_and_layer"):
                self.q_layernorm = qk_layernorm_builder()
                self.k_layernorm = qk_layernorm_builder()
            else:
                raise ValueError(f"Unknown qk_layernorm mode: {self.qk_layernorm}")
        else:
            self.q_layernorm = None
            self.k_layernorm = None

        self.core_attention = CoreAttention(
            parallel_attention_config.core_attention_config,
            self.layer_number,
            self.attn_mask_type,
        )
        self.checkpoint_core_attention = parallel_attention_config.recompute_granularity == "selective"

        if self.use_flash_attn:
            is_causal = self.attn_mask_type == AttnMaskType.causal
            self.core_attention_flash = FlashSelfAttention(
                causal=is_causal, attention_dropout=parallel_attention_config.attention_dropout
            )
            if parallel_attention_config.pos_emb == "rotary":
                self.rotary_ndims = int(parallel_attention_config.rotary_pct * self.hidden_size_per_attention_head)
                self.rotary_emb = RotaryEmbeddingFlash(
                    self.rotary_ndims,
                    base=parallel_attention_config.rotary_emb_base,
                    pos_idx_in_fp32=False,  # For backward compatibility with models already trained
                )

        # Output.
        self.dense = mpu.RowParallelLinear(
            projection_size,
            parallel_attention_config.hidden_size,
            input_is_parallel=True,
            init_method=output_layer_init_method,
            skip_bias_add=True,
            init_rng_key=f"{self.module_name}_layer_{layer_number}_dense",
            sequence_parallel=self.sequence_parallel,
        )

        if parallel_attention_config.pos_emb == "rotary" and not self.use_flash_attn:
            if parallel_attention_config.rotary_pct == 1:
                self.rotary_ndims = None
            else:
                assert parallel_attention_config.rotary_pct < 1
                self.rotary_ndims = int(self.hidden_size_per_attention_head * parallel_attention_config.rotary_pct)
            rotary_dim = self.rotary_ndims if self.rotary_ndims is not None else self.hidden_size_per_attention_head
            self.rotary_emb = RotaryEmbedding(
                rotary_dim,
                base=parallel_attention_config.rotary_emb_base,
                precision=parallel_attention_config.params_dtype,
            )

    def _checkpointed_attention_forward(self, query_layer, key_layer, value_layer):
        """Forward method with activation checkpointing."""

        def custom_forward(*inputs):
            query_layer = inputs[0]
            key_layer = inputs[1]
            value_layer = inputs[2]
            output_ = self.core_attention(query_layer, key_layer, value_layer)
            return output_

        hidden_states = mpu.checkpoint(custom_forward, False, query_layer, key_layer, value_layer)

        return hidden_states

    def _allocate_memory(self, inference_max_sequence_len, batch_size):
        return torch.empty(
            inference_max_sequence_len,
            batch_size,
            self.num_attention_heads_per_partition,
            self.hidden_size_per_attention_head,
            dtype=self.params_dtype,
            device=torch.cuda.current_device(),
        )

    def _check_rotary_base(self, sequence_len: int) -> None:
        if not self.enable_rotary_base_check:
            return
        min_rotary_emb_base = math.ceil((2 * sequence_len) / (math.pi))
        if min_rotary_emb_base > self.rotary_emb_base:
            raise ValueError(
                f"To support sequence length {sequence_len}, a minimum rotary base of "
                f"{min_rotary_emb_base} is required, but current base is {self.rotary_emb_base}"
            )

    def _apply_rotary_embedding(self, query_layer, key_layer, offset=0):
        """
        query_layer, key_layer: (seqlen, batch, nheads, headdim)
        """
        if self.rotary_ndims is not None:
            # partial rotary
            query_rot, query_pass = (
                query_layer[..., : self.rotary_ndims],
                query_layer[..., self.rotary_ndims :],
            )
            key_rot, key_pass = (
                key_layer[..., : self.rotary_ndims],
                key_layer[..., self.rotary_ndims :],
            )
        else:
            # full rotary
            query_rot, key_rot = query_layer, key_layer
        apply_rotary_fn = apply_rotary_pos_emb_torch if self.bf16 else apply_rotary_pos_emb

        # For inference, add the offset value to sequence length
        sequence_len = key_layer.shape[0] + offset
        # Apply the rotary embeddings
        cos, sin = self.rotary_emb(key_layer, seq_len=sequence_len)
        query_layer, key_layer = apply_rotary_fn(query_rot, key_rot, cos, sin, offset=offset)

        # query [sq, b, np, hn], key [sk, b, np, hn]
        if self.rotary_ndims is not None:
            query_layer = torch.cat((query_layer, query_pass), dim=-1)
            key_layer = torch.cat((key_layer, key_pass), dim=-1)
        return query_layer, key_layer

    def _update_kv_cache(self, inference_params, key_layer: Tensor, value_layer: Tensor) -> Tuple[Tensor]:
        """k, v: (seqlen, batch_size, nheads, head_dim)"""
        # =================================================
        # Pre-allocate memory for key-values for inference.
        # =================================================
        if self.layer_number not in inference_params.key_value_memory_dict:
            inf_max_seq_len = inference_params.max_sequence_len
            inf_max_batch_size = inference_params.max_batch_size
            inference_key_memory = self._allocate_memory(inf_max_seq_len, inf_max_batch_size)
            inference_value_memory = self._allocate_memory(inf_max_seq_len, inf_max_batch_size)
            inference_params.key_value_memory_dict[self.layer_number] = (
                inference_key_memory,
                inference_value_memory,
            )
            preallocated = False
        else:
            # It's possible that we allocate kv_cache larger than necessary, we'll slice to
            # take a subset of the right batch size and seqlen layer
            # If the memory is pre-allocated, and we're using FT's fused kernel,
            # then k_cache has shape (b h d/packsize seqlen packsize) where packsize = 4 or 8 (fp32 vs fp16,bf16),
            # and v_cache has shape (b h s d).
            inference_key_memory, inference_value_memory = inference_params.key_value_memory_dict[self.layer_number]
            preallocated = True
        # ==================================
        # Adjust key and value for inference
        # ==================================
        seqlen_dim_k = 0 if not (inference_params.fused_ft_kernel and preallocated) else 3
        batch_dim = 1 if not (inference_params.fused_ft_kernel and preallocated) else 0
        batch_start = inference_params.batch_size_offset
        batch_end = batch_start + key_layer.size(1)
        sequence_start = inference_params.sequence_len_offset
        sequence_end = sequence_start + key_layer.size(0)
        assert batch_end <= inference_key_memory.size(batch_dim)
        assert sequence_end <= inference_key_memory.size(seqlen_dim_k)
        # Copy key and values.
        if not inference_params.fused_ft_kernel:
            inference_key_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = key_layer
            inference_value_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = value_layer
            key_layer = inference_key_memory[:sequence_end, batch_start:batch_end, ...]
            value_layer = inference_value_memory[:sequence_end, batch_start:batch_end, ...]
            return key_layer, value_layer
        else:
            assert ft_attention is not None
            assert rearrange is not None
            assert inference_params.sequence_len_offset == 0
            # FT kernel requires different layouts for the k_cache and v_cache.
            assert inference_key_memory.dtype in [torch.float16, torch.bfloat16, torch.float32]
            assert inference_value_memory.dtype == inference_key_memory.dtype
            packsize = 4 if inference_key_memory.dtype == torch.float32 else 8
            if not preallocated:
                inference_key_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = key_layer
                inference_value_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = value_layer
                inference_key_memory = rearrange(
                    inference_key_memory, "s b h (d packsize) -> b h d s packsize", packsize=packsize
                ).contiguous()
                inference_value_memory = rearrange(inference_value_memory, "s b h d -> b h s d").contiguous()
                inference_params.key_value_memory_dict[self.layer_number] = (
                    inference_key_memory,
                    inference_value_memory,
                )
            else:
                inference_key_memory[batch_start:batch_end, :, :, sequence_start:sequence_end, :] = rearrange(
                    key_layer, "s b h (d packsize) -> b h d s packsize", packsize=packsize
                )
                inference_value_memory[batch_start:batch_end, :, sequence_start:sequence_end, :] = rearrange(
                    value_layer, "s b h d -> b h s d"
                )
            return key_layer, value_layer

    def forward(self, hidden_states, encoder_output=None, inference_params=None):
        """
        Wrt to inference_params, there are 4 code paths:
        1. inference_params=None: training, we don't touch inference_params at all.
        2. inference_params is not None, inference_params.sequence_len_offset == 0: prompt processing.
            We do the same ops as during training, except we store the key_layer and value_layer to the kv cache
            in inference_params.
        3. inference_params is not None, inference_params.sequence_len_offset > 0: iterative decoding.
            3a. inference_params.fused_ft_kernel == False: we do QKV projection, rotary embedding,
                add k and v to the kv_cache, then load the kv_cache from inference_params, do single-query attention,
                then do output projection.
            3b. inference_params.fused_ft_kernel == True: we do QKV projection (but no rotary embedding),
                load the kv_cache from inference_params, then call ft_attention (which fuses rotary embedding,
                adding k and v to the kv_cache, and single-query attention), then do output projection.
        """
        # hidden_states: [sq, b, h]

        # =====================
        # Query, Key, and Value
        # =====================

        if self.attention_type == AttnType.self_attn:
            # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
            mixed_x_layer, _ = self.query_key_value(hidden_states)

            # [sq, b, (np * 3 * hn)] --> [sq, b, np, 3 * hn]
            new_tensor_shape = mixed_x_layer.size()[:-1] + (
                self.num_attention_heads_per_partition,
                3 * self.hidden_size_per_attention_head,
            )
            mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)

            # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
            (query_layer, key_layer, value_layer) = mpu.split_tensor_along_last_dim(mixed_x_layer, 3)

            if self.q_layernorm:
                assert self.k_layernorm
                query_layer = self.q_layernorm(query_layer)
                key_layer = self.k_layernorm(key_layer)

                if self.use_flash_attn:
                    mixed_x_layer = torch.cat([query_layer, key_layer, value_layer], dim=-1)
                    assert mixed_x_layer.shape == new_tensor_shape

            if self.use_flash_attn:  # FlashAttention needs to reshape things
                qkv = rearrange(mixed_x_layer, "s b h (three d) -> b s three h d", three=3).contiguous()
        else:
            # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)]
            mixed_kv_layer, _ = self.key_value(encoder_output)

            # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn]
            new_tensor_shape = mixed_kv_layer.size()[:-1] + (
                self.num_attention_heads_per_partition,
                2 * self.hidden_size_per_attention_head,
            )
            mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape)

            # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn]
            (key_layer, value_layer) = mpu.split_tensor_along_last_dim(mixed_kv_layer, 2)

            assert not self.q_layernorm and not self.k_layernorm, "qk_layernorm not yet implemented for cross-attention"

            # Attention head [sq, b, h] --> [sq, b, hp]
            query_layer, _ = self.query(hidden_states)
            # [sq, b, hp] --> [sq, b, np, hn]
            new_tensor_shape = query_layer.size()[:-1] + (
                self.num_attention_heads_per_partition,
                self.hidden_size_per_attention_head,
            )
            query_layer = query_layer.view(*new_tensor_shape)
            if self.use_flash_attn:  # FlashAttention needs to reshape things
                kv = rearrange(mixed_kv_layer, "s b h (two d) -> b s two h d", two=2).contiguous()
                q = rearrange(query_layer, "s b h d -> b s h d").contiguous()

        if self.pos_emb == "rotary":
            if (
                inference_params is None
                or inference_params.sequence_len_offset == 0
                or not inference_params.fused_ft_kernel
            ):  # Code path 1, 2, & 3a above
                offset = 0 if inference_params is None else inference_params.sequence_len_offset
                sequence_len = key_layer.shape[0] + offset
                self._check_rotary_base(sequence_len)
                if not self.use_flash_attn:
                    query_layer, key_layer = self._apply_rotary_embedding(
                        query_layer=query_layer, key_layer=key_layer, offset=offset
                    )
                else:
                    assert self.attention_type == AttnType.self_attn
                    qkv = self.rotary_emb(qkv, seqlen_offset=offset)
                    # For inference_params
                    # TD [2023-03-28]: We need to assign to query_layer here, otherwise later on when
                    # we use query_layer, we would be using the "unrotated" query_layer.
                    query_layer = rearrange(qkv[:, :, 0], "b s h d -> s b h d")
                    key_layer = rearrange(qkv[:, :, 1], "b s h d -> s b h d")
                    value_layer = rearrange(qkv[:, :, 2], "b s h d -> s b h d")

        # Code path 2 & 3a
        if inference_params is not None and (
            inference_params.sequence_len_offset == 0 or not inference_params.fused_ft_kernel
        ):
            key_layer, value_layer = self._update_kv_cache(inference_params, key_layer, value_layer)

        # ==================================
        # core attention computation
        # ==================================

        if (
            inference_params is None
            or inference_params.sequence_len_offset == 0
            or not inference_params.fused_ft_kernel
        ):  # Code path 1, 2, & 3a above
            # For iterative decoding (code path 3a), we don't use FlashAttention (since it has causal=True and
            # will mask out things)
            if not self.use_flash_attn or (inference_params is not None and inference_params.sequence_len_offset > 0):
                if self.checkpoint_core_attention:
                    context_layer = self._checkpointed_attention_forward(query_layer, key_layer, value_layer)
                else:
                    context_layer = self.core_attention(query_layer, key_layer, value_layer)
            else:
                if self.attention_type == AttnType.self_attn:
                    context_layer = rearrange(self.core_attention_flash(qkv=qkv), "b s h d -> s b (h d)").contiguous()
                else:
                    context_layer = rearrange(
                        self.core_attention_flash(kv=kv, q=q), "b s h d -> s b (h d)"
                    ).contiguous()
        else:
            # Code path 3b
            assert ft_attention is not None
            assert query_layer.shape[0] == key_layer.shape[0] == value_layer.shape[0] == 1
            q, k, v = [rearrange(x, "1 b h d -> b h d").contiguous() for x in [query_layer, key_layer, value_layer]]
            k_cache, v_cache = inference_params.key_value_memory_dict[self.layer_number]
            batch_start = inference_params.batch_size_offset
            batch_end = batch_start + q.size(0)
            assert batch_end <= k_cache.size(0)
            k_cache = k_cache[batch_start:batch_end]
            v_cache = v_cache[batch_start:batch_end]
            lengths_per_sample = (
                inference_params.lengths_per_sample[batch_start:batch_end]
                if inference_params.lengths_per_sample is not None
                else None
            )
            # For now, we're emulating the behavior where rotary_emb calculation is done in bf16
            # This is because we have a lot of models already trained with bf16 rotary, and the
            # output of bf16 rotary is quite different from fp32 rotary.
            # If rotary_bf16, compute rotary_cos / sin and pass to ft_attention.
            # If not rotary_bf16, then just set rotary_cos / sin to None.
            if self.pos_emb == "rotary":
                dtype = self.rotary_emb.inv_freq.dtype
                if lengths_per_sample is not None:
                    # angle should have shape (batch, rotary_dim // 2) since that's what ft_attention expects
                    angle = self.rotary_emb.inv_freq * rearrange(lengths_per_sample.to(dtype), "b -> b 1")
                else:
                    angle = repeat(
                        self.rotary_emb.inv_freq * inference_params.sequence_len_offset, "d -> b d", b=q.size(0)
                    ).contiguous()
                rotary_cos, rotary_sin = angle.cos(), angle.sin()
            else:
                rotary_cos, rotary_sin = None, None
            rotary_emb_dim = 0 if self.pos_emb != "rotary" else self.rotary_ndims
            rotary_emb_base = float(self.rotary_emb_base) if self.pos_emb == "rotary" else 0.0
            context_layer = ft_attention.single_query_attention(
                q,
                k,
                v,
                k_cache,
                v_cache,
                lengths_per_sample,
                rotary_cos,
                rotary_sin,
                None,  # nnz_head_idx, unrelated option that we don't ever need
                inference_params.sequence_len_offset,
                rotary_emb_dim,
                rotary_emb_base,
            )
            context_layer = rearrange(context_layer, "b h d -> 1 b (h d)")

        # =================
        # Output. [sq, b, h]
        # =================

        output, bias = self.dense(context_layer)

        return output, bias


def bias_dropout_add(x, bias, residual, prob, training):
    # type: (Tensor, Tensor, Tensor, float, bool) -> Tensor
    out = torch.nn.functional.dropout(x + bias, p=prob, training=training)
    out = residual + out
    return out


def get_bias_dropout_add(training):
    def _bias_dropout_add(x, bias, residual, prob):
        return bias_dropout_add(x, bias, residual, prob, training)

    return _bias_dropout_add


@torch.jit.script
def bias_dropout_add_fused_train(
    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
) -> torch.Tensor:
    return bias_dropout_add(x, bias, residual, prob, True)


@torch.jit.script
def bias_dropout_add_fused_inference(
    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
) -> torch.Tensor:
    return bias_dropout_add(x, bias, residual, prob, False)


@dataclasses.dataclass
class ParallelTransformerLayerConfig(Config):
    bf16: bool = DISABLE_PREFIX()
    apply_residual_connection_post_layernorm: bool = DISABLE_PREFIX()
    fp32_residual_connection: bool = DISABLE_PREFIX()
    hidden_size: int
    layernorm_epsilon: float = DISABLE_PREFIX()
    no_persist_layer_norm: bool = DISABLE_PREFIX()
    sequence_parallel: bool = DISABLE_PREFIX()
    hidden_dropout: float
    bias_dropout_fusion: bool = DISABLE_PREFIX()
    num_experts: int = DISABLE_PREFIX()
    parallel_mlp_config: ParallelMLPConfig
    parallel_attention_config: ParallelAttentionConfig


class ParallelTransformerLayer(MegatronModule):
    """A single transformer layer.

    Transformer layer takes input with size [s, b, h] and returns an
    output of the same size.
    """

    def __init__(
        self,
        parallel_transformer_layer_config: ParallelTransformerLayerConfig,
        init_method,
        output_layer_init_method,
        layer_number,
        module_name: str,
        layer_type=LayerType.encoder,
        self_attn_mask_type=AttnMaskType.padding,
        drop_path_rate=0.0,
    ):
        super(ParallelTransformerLayer, self).__init__()
        self.module_name = module_name
        self.layer_number = layer_number
        self.layer_type = layer_type
        self.sequence_parallel = parallel_transformer_layer_config.sequence_parallel

        self.apply_residual_connection_post_layernorm = (
            parallel_transformer_layer_config.apply_residual_connection_post_layernorm
        )

        self.bf16 = parallel_transformer_layer_config.bf16
        self.fp32_residual_connection = parallel_transformer_layer_config.fp32_residual_connection

        # Layernorm on the input data.
        self.input_layernorm = LayerNorm(
            parallel_transformer_layer_config.hidden_size,
            eps=parallel_transformer_layer_config.layernorm_epsilon,
            no_persist_layer_norm=parallel_transformer_layer_config.no_persist_layer_norm,
            sequence_parallel=self.sequence_parallel,
        )

        # Self attention.
        self.self_attention = ParallelAttention(
            parallel_transformer_layer_config.parallel_attention_config,
            init_method,
            output_layer_init_method,
            layer_number,
            module_name=self.module_name,
            attention_type=AttnType.self_attn,
            attn_mask_type=self_attn_mask_type,
        )
        self.hidden_dropout = parallel_transformer_layer_config.hidden_dropout
        self.bias_dropout_fusion = parallel_transformer_layer_config.bias_dropout_fusion
        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else None

        # Layernorm on the attention output
        self.post_attention_layernorm = LayerNorm(
            parallel_transformer_layer_config.hidden_size,
            eps=parallel_transformer_layer_config.layernorm_epsilon,
            no_persist_layer_norm=parallel_transformer_layer_config.no_persist_layer_norm,
            sequence_parallel=self.sequence_parallel,
        )

        if self.layer_type == LayerType.decoder:
            self.inter_attention = ParallelAttention(
                parallel_transformer_layer_config.parallel_attention_config,
                init_method,
                output_layer_init_method,
                layer_number,
                module_name=self.module_name,
                attention_type=AttnType.cross_attn,
            )
            # Layernorm on the attention output.
            self.post_inter_attention_layernorm = LayerNorm(
                parallel_transformer_layer_config.hidden_size,
                eps=parallel_transformer_layer_config.layernorm_epsilon,
                no_persist_layer_norm=parallel_transformer_layer_config.no_persist_layer_norm,
                sequence_parallel=self.sequence_parallel,
            )

        # MLP
        self.mlp = ParallelMLP(
            parallel_transformer_layer_config.parallel_mlp_config,
            init_method,
            output_layer_init_method,
            init_rng_key=f"{self.module_name}_layer_{layer_number}_mlp",
        )

        # Set bias+dropout+add fusion grad_enable execution handler.
        TORCH_MAJOR = int(torch.__version__.split(".")[0])
        TORCH_MINOR = int(torch.__version__.split(".")[1])
        use_nvfuser = TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10)
        self.bias_dropout_add_exec_handler = nullcontext if use_nvfuser else torch.enable_grad

    def forward(self, hidden_states, encoder_output=None, inference_params=None):
        # hidden_states: [s, b, h]

        # Layer norm at the beginning of the transformer layer.
        layernorm_output = self.input_layernorm(hidden_states)
        # Self attention.
        attention_output, attention_bias = self.self_attention(layernorm_output, inference_params=inference_params)

        # Residual connection.
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = hidden_states

        if self.drop_path is None:
            # jit scripting for a nn.module (with dropout) is not
            # trigerring the fusion kernel. For now, we use two
            # different nn.functional routines to account for varying
            # dropout semantics during training and inference phases.
            if self.bias_dropout_fusion:
                if self.training:
                    bias_dropout_add_func = bias_dropout_add_fused_train
                else:
                    bias_dropout_add_func = bias_dropout_add_fused_inference
            else:
                bias_dropout_add_func = get_bias_dropout_add(self.training)

            with self.bias_dropout_add_exec_handler():
                layernorm_input = bias_dropout_add_func(
                    attention_output, attention_bias.expand_as(residual), residual, self.hidden_dropout
                )
        else:
            out = torch.nn.functional.dropout(
                attention_output + attention_bias, p=self.hidden_dropout, training=self.training
            )
            layernorm_input = residual + self.drop_path(out)

        # Layer norm post the self attention.
        layernorm_output = self.post_attention_layernorm(layernorm_input)

        if self.layer_type == LayerType.decoder:
            attention_output, attention_bias = self.inter_attention(
                layernorm_output, enc_dec_attn_mask, encoder_output=encoder_output
            )
            # residual connection
            if self.apply_residual_connection_post_layernorm:
                residual = layernorm_output
            else:
                residual = layernorm_input

            with self.bias_dropout_add_exec_handler():
                layernorm_input = bias_dropout_add_func(
                    attention_output, attention_bias.expand_as(residual), residual, self.hidden_dropout
                )

            # Layer norm post the decoder attention
            layernorm_output = self.post_inter_attention_layernorm(layernorm_input)

        # MLP.
        mlp_output, mlp_bias = self.mlp(layernorm_output)

        # Second residual connection.
        if self.apply_residual_connection_post_layernorm:
            residual = layernorm_output
        else:
            residual = layernorm_input

        if self.drop_path is None:
            with self.bias_dropout_add_exec_handler():
                output = bias_dropout_add_func(mlp_output, mlp_bias.expand_as(residual), residual, self.hidden_dropout)

            # Jit compiled function creates 'view' tensor. This tensor
            # potentially gets saved in the MPU checkpoint function context,
            # which rejects view tensors. While making a viewless tensor here
            # won't result in memory savings (like the data loader, or
            # p2p_communication), it serves to document the origin of this
            # 'view' tensor.
            output = mpu.make_viewless_tensor(inp=output, requires_grad=output.requires_grad, keep_graph=True)

        else:
            out = torch.nn.functional.dropout(mlp_output + mlp_bias, p=self.hidden_dropout, training=self.training)
            output = residual + self.drop_path(out)

        return output


class NoopTransformerLayer(MegatronModule):
    """A single 'no-op' transformer layer.

    The sole purpose of this layer is for when a standalone embedding layer
    is used (i.e., args.standalone_embedding_stage == True). In this case,
    zero transformer layers are assigned when pipeline rank == 0. Additionally,
    when virtual pipeline rank >= 1, zero total model parameters are created
    (virtual rank 0 contains the input embedding). This results in the model's
    input and output tensors being the same, which causes an error when
    performing certain memory optimiations on the output tensor (e.g.,
    deallocating it). Thus, this layer disconnects the input from the output
    via a clone. Since ranks containing a no-op layer are generally under-
    utilized (both compute and memory), there's no worry of any performance
    degredation.
    """

    def __init__(self, layer_number):
        super().__init__()
        self.layer_number = layer_number

    def forward(self, hidden_states, encoder_output=None, inference_params=None):
        return hidden_states.clone()


@dataclasses.dataclass
class ParallelTransformerConfig(Config):
    bf16: bool = DISABLE_PREFIX()
    decoder_num_layers: int
    distribute_saved_activations: bool = DISABLE_PREFIX()
    encoder_num_layers: int
    fp32_residual_connection: bool = DISABLE_PREFIX()
    hidden_size: int
    layernorm_epsilon: float = DISABLE_PREFIX()
    model_type: ModelType = DISABLE_PREFIX()
    no_persist_layer_norm: bool = DISABLE_PREFIX()
    num_layers: int
    parallel_transformer_layer_config: ParallelTransformerLayerConfig
    pipeline_model_parallel_split_rank: int = DISABLE_PREFIX()
    recompute_granularity: str = DISABLE_PREFIX()
    recompute_method: str = DISABLE_PREFIX()
    recompute_num_layers: int = DISABLE_PREFIX()
    sequence_parallel: bool = DISABLE_PREFIX()
    standalone_embedding_stage: bool = DISABLE_PREFIX()
    transformer_pipeline_model_parallel_size: int = DISABLE_PREFIX()
    virtual_pipeline_model_parallel_size: int = DISABLE_PREFIX()


class ParallelTransformer(MegatronModule):
    """Transformer class."""

    def __init__(
        self,
        parallel_transformer_config: ParallelTransformerConfig,
        init_method,
        output_layer_init_method,
        module_name: str,
        layer_type=LayerType.encoder,
        self_attn_mask_type=AttnMaskType.padding,
        post_layer_norm=True,
        pre_process=True,
        post_process=True,
        drop_path_rate=0.0,
        num_layers=None,
    ):
        super(ParallelTransformer, self).__init__()

        self.module_name = (module_name,)
        self.layer_type = layer_type
        self.model_type = parallel_transformer_config.model_type
        self.bf16 = parallel_transformer_config.bf16
        self.fp32_residual_connection = parallel_transformer_config.fp32_residual_connection
        self.post_layer_norm = post_layer_norm
        self.pre_process = pre_process
        self.post_process = post_process
        self.input_tensor = None
        self.drop_path_rate = drop_path_rate

        # Store activation checkpointing flag.
        self.recompute_granularity = parallel_transformer_config.recompute_granularity
        self.recompute_method = parallel_transformer_config.recompute_method
        self.recompute_num_layers = parallel_transformer_config.recompute_num_layers
        self.sequence_parallel = parallel_transformer_config.sequence_parallel
        self.distribute_saved_activations = (
            parallel_transformer_config.distribute_saved_activations and not self.sequence_parallel
        )

        # Number of layers.
        if num_layers is None:
            self.num_layers = mpu.get_num_layers(
                pipeline_model_parallel_split_rank=parallel_transformer_config.pipeline_model_parallel_split_rank,
                standalone_embedding_stage=parallel_transformer_config.standalone_embedding_stage,
                transformer_pipeline_model_parallel_size=parallel_transformer_config.transformer_pipeline_model_parallel_size,
                encoder_num_layers=parallel_transformer_config.encoder_num_layers,
                decoder_num_layers=parallel_transformer_config.decoder_num_layers,
                num_layers=parallel_transformer_config.num_layers,
                is_encoder_and_decoder_model=parallel_transformer_config.model_type == ModelType.encoder_and_decoder,
                is_decoder=layer_type == LayerType.decoder,
            )
        else:
            self.num_layers = num_layers

        if layer_type == LayerType.encoder:
            drop_path_layers = parallel_transformer_config.num_layers
        else:
            drop_path_layers = (
                parallel_transformer_config.decoder_num_layers
                if parallel_transformer_config.decoder_num_layers is not None
                else parallel_transformer_config.num_layers
            )
        self.drop_path_rates = [rate.item() for rate in torch.linspace(0, self.drop_path_rate, drop_path_layers)]

        # Transformer layers.
        def build_layer(layer_number):
            return ParallelTransformerLayer(
                parallel_transformer_config.parallel_transformer_layer_config,
                init_method,
                output_layer_init_method,
                layer_number,
                self.module_name,
                layer_type=layer_type,
                self_attn_mask_type=self_attn_mask_type,
                drop_path_rate=self.drop_path_rates[layer_number - 1],
            )

        if parallel_transformer_config.virtual_pipeline_model_parallel_size is not None:
            assert (
                parallel_transformer_config.num_layers
                % parallel_transformer_config.virtual_pipeline_model_parallel_size
                == 0
            ), ("num_layers_per_stage must be divisible by " "virtual_pipeline_model_parallel_size")
            assert parallel_transformer_config.model_type != ModelType.encoder_and_decoder
            # Number of layers in each model chunk is the number of layers in the stage,
            # divided by the number of model chunks in a stage.
            self.num_layers = self.num_layers // parallel_transformer_config.virtual_pipeline_model_parallel_size
            # With 8 layers, 2 stages, and 4 model chunks, we want an assignment of
            # layers to stages like (each list is a model chunk):
            # Stage 0: [0]  [2]  [4]  [6]
            # Stage 1: [1]  [3]  [5]  [7]
            # With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of
            # layers to stages like (each list is a model chunk):
            # Stage 0: [0, 1]  [4, 5]
            # Stage 1: [2, 3]  [6, 7]
            offset = mpu.get_virtual_pipeline_model_parallel_rank() * (
                parallel_transformer_config.num_layers
                // parallel_transformer_config.virtual_pipeline_model_parallel_size
            ) + (mpu.get_pipeline_model_parallel_rank() * self.num_layers)
        else:
            # Each stage gets a contiguous set of layers.
            if (
                parallel_transformer_config.model_type == ModelType.encoder_and_decoder
                and mpu.get_pipeline_model_parallel_world_size() > 1
            ):
                pipeline_rank = mpu.get_pipeline_model_parallel_rank()
                if layer_type == LayerType.encoder:
                    offset = pipeline_rank * self.num_layers
                else:
                    num_ranks_in_enc = parallel_transformer_config.pipeline_model_parallel_split_rank
                    offset = (pipeline_rank - num_ranks_in_enc) * self.num_layers
            else:
                offset = mpu.get_pipeline_model_parallel_rank() * self.num_layers

        if self.num_layers == 0:
            # When a standalone embedding stage is used (e.g.,
            # args.standalone_embedding_stage == True), virtual pipeline ranks
            # on pipeline rank 0 will have zero transformer layers assigned to
            # them. This results in the model's input and output tensors to be
            # the same, which will cause failure for certain output tensor
            # optimizations (e.g., pipeline output deallocation). To remedy
            # this, we assign a 'no-op' layer on these ranks, which will
            # disconnect the input tensor from the output tensor.
            self.num_layers = 1
            self.layers = torch.nn.ModuleList([NoopTransformerLayer(1)])
        else:
            self.layers = torch.nn.ModuleList([build_layer(i + 1 + offset) for i in range(self.num_layers)])

        if self.post_process and self.post_layer_norm:
            # Final layer norm before output.
            self.final_layernorm = LayerNorm(
                parallel_transformer_config.hidden_size,
                eps=parallel_transformer_config.layernorm_epsilon,
                no_persist_layer_norm=parallel_transformer_config.no_persist_layer_norm,
                sequence_parallel=self.sequence_parallel,
            )

    def _get_layer(self, layer_number):
        return self.layers[layer_number]

    def _checkpointed_forward(self, hidden_states, encoder_output):
        """Forward method with activation checkpointing."""

        def custom(start, end):
            def custom_forward(*inputs):
                x_ = inputs[0]
                encoder_output = inputs[1]
                for index in range(start, end):
                    layer = self._get_layer(index)
                    x_ = layer(x_, encoder_output)
                return x_

            return custom_forward

        if self.recompute_method == "uniform":
            # Uniformly divide the total number of Transformer layers and checkpoint
            # the input activation of each divided chunk.
            # A method to further reduce memory usage reducing checkpoints.
            l = 0
            while l < self.num_layers:
                hidden_states = mpu.checkpoint(
                    custom(l, l + self.recompute_num_layers),
                    self.distribute_saved_activations,
                    hidden_states,
                    encoder_output,
                )
                l += self.recompute_num_layers

        elif self.recompute_method == "block":
            # Checkpoint the input activation of only a set number of individual
            # Transformer layers and skip the rest.
            # A method fully use the device memory removing redundant re-computation.
            for l in range(self.num_layers):
                if l < self.recompute_num_layers:
                    hidden_states = mpu.checkpoint(
                        custom(l, l + 1),
                        self.distribute_saved_activations,
                        hidden_states,
                        encoder_output,
                    )
                else:
                    hidden_states = custom(l, l + 1)(hidden_states, encoder_output)
        else:
            raise ValueError("Invalid activation recompute method.")

        return hidden_states

    def set_input_tensor(self, input_tensor):
        """Set input tensor to be used instead of forward()'s input.

        When doing pipeline parallelism the input from the previous
        stage comes from communication, not from the input, so the
        model's forward_step_func won't have it. This function is thus
        used by internal code to bypass the input provided by the
        forward_step_func"""
        self.input_tensor = input_tensor

    def forward(self, hidden_states, encoder_output=None, inference_params=None):
        # hidden_states: [s, b, h]

        # Checks.
        if inference_params:
            assert self.recompute_granularity is None, "inference does not work with activation checkpointing"

        if not self.pre_process:
            # See set_input_tensor()
            hidden_states = self.input_tensor

        # Viewless tensor.
        # - We only need to create a viewless tensor in the case of micro batch
        #   size (mbs) == 1, since in this case, 'hidden_states.transpose()'
        #   above creates a view tensor, and '.contiguous()' is a pass-through.
        #   For mbs >= 2, '.contiguous()' creates a new tensor, eliminating
        #   the need to make it viewless.
        #
        #   However, we don't explicitly check mbs == 1 here because
        #   make_viewless_tensor() has negligible overhead when its input
        #   is already viewless.
        #
        # - For the 'else' case above, calling make_viewless_tensor() here is
        #   likely redundant, since p2p_communication.py (likely originator)
        #   already creates viewless tensors. That said, make_viewless_tensor()
        #   is called here to be future-proof and corner-case-proof.
        hidden_states = mpu.make_viewless_tensor(
            hidden_states,
            requires_grad=True,
            keep_graph=True,
        )

        if self.sequence_parallel:
            rng_context = mpu.get_cuda_rng_tracker().fork()
        else:
            rng_context = nullcontext()

        with rng_context:
            # Forward pass.
            if self.recompute_granularity == "full":
                hidden_states = self._checkpointed_forward(hidden_states, encoder_output)
            else:
                for index in range(self.num_layers):
                    layer = self._get_layer(index)
                    hidden_states = layer(
                        hidden_states,
                        encoder_output=encoder_output,
                        inference_params=inference_params,
                    )

        # Final layer norm.
        if self.post_process and self.post_layer_norm:
            hidden_states = self.final_layernorm(hidden_states)

        return hidden_states

megatron/model/transformer.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilities for models."""

import math
import socket

import torch
from numpy.random import default_rng, SeedSequence
import multiprocessing
import concurrent.futures
import numpy as np
from typing import Sequence, List
import functools

from megatron import get_args

from megatron import mpu


def init_method_normal(sigma):
    """Init method based on N(0, sigma)."""

    def fill(rng: np.random.Generator, shape: Sequence[int]):
        n_threads = 16
        n_elements = np.prod(shape)
        values = np.empty(n_elements)
        assert n_elements % n_threads == 0, "Number of elements must be a multiple of number of threads!"
        step = np.ceil(n_elements / n_threads).astype(np.int_)

        # TODO(erich): hacky way to get new generator seeds, we should use spawn as soon as we have
        # numpy 1.25 and not this!!!
        seeds = rng.integers(0, 1 << 63, n_threads)

        _random_generators = [np.random.default_rng(s) for s in seeds]

        executor = concurrent.futures.ThreadPoolExecutor(n_threads)

        def _fill(generator, out, first, last):
            out[first:last] = generator.normal(loc=0.0, scale=sigma, size=step)

        futures = {}
        for i in range(n_threads):
            args = (_fill, _random_generators[i], values, i * step, (i + 1) * step)
            futures[executor.submit(*args)] = i
        concurrent.futures.wait(futures)
        return np.reshape(values, shape)

    return fill


# TODO(erich): combine these two functions to reduce code duplication
def scaled_init_method_normal(sigma, num_layers):
    """Init method based on N(0, sigma/sqrt(2*num_layers)."""
    std = sigma / math.sqrt(2.0 * num_layers)

    def fill(rng: np.random.Generator, shape: Sequence[int]):
        n_threads = 16
        n_elements = np.prod(shape)
        values = np.empty(n_elements)
        assert n_elements % n_threads == 0, "Number of elements must be a multiple of number of threads!"
        step = np.ceil(n_elements / n_threads).astype(np.int_)

        # TODO(erich): hacky way to get new generator seeds, we should use spawn as soon as we have
        # numpy 1.25 and not this!!!
        seeds = rng.integers(0, 1 << 63, n_threads)

        _random_generators = [np.random.default_rng(s) for s in seeds]
        executor = concurrent.futures.ThreadPoolExecutor(n_threads)

        def _fill(generator, out, first, last):
            out[first:last] = generator.normal(loc=0.0, scale=std, size=step)

        futures = {}
        for i in range(n_threads):
            args = (_fill, _random_generators[i], values, i * step, (i + 1) * step)
            futures[executor.submit(*args)] = i
        concurrent.futures.wait(futures)
        return np.reshape(values, shape)

    return fill


def check_shapes(tensor, expected_shapes=None, expected_ndim=None, expected_dtype=None):
    """Check that the passed-in `tensor` has the expected shape and ndim. Should an expected dim be None, it is not checked.

    Args:
        tensor: The tensor to shape-check.
        expected_shapes: The prefix of the shapes to expect in `tensor`.
        expected_ndim: The expected number of dimensions in `tensor`.
        expected_dtype: Expected dtype.
    """
    assert tensor is not None, "tensor is None"

    if expected_shapes is None and expected_ndim is None and expected_dtype is None:
        return

    if expected_dtype is not None:
        assert tensor.dtype == expected_dtype

    if expected_ndim is not None:
        assert tensor.ndim == expected_ndim, f"Unexpected ndims detected. Got: {tensor.ndim} expected: {expected_ndim}"

    expected_shapes = list(expected_shapes or [])

    if expected_ndim is not None:
        assert (
            len(expected_shapes) <= tensor.ndim
        ), f"Asking to check too many shapes. Got: {expected_shapes} expected: <= {tensor.ndim}"

    t_shapes = list(tensor.shape[: len(expected_shapes)])
    for i, e in enumerate(expected_shapes):
        if e is None:
            t_shapes[i] = None
    assert t_shapes == expected_shapes, f"Mismatch detected. Got: {t_shapes} expected: {expected_shapes}"


def get_linear_layer(rows, columns, init_method):
    """Simple linear layer with weight initialization."""
    layer = torch.nn.Linear(rows, columns)
    if get_args().perform_initialization:
        init_method(layer.weight)
    with torch.no_grad():
        layer.bias.zero_()
    return layer


@torch.jit.script
def gelu_impl(x):
    """OpenAI's gelu implementation."""
    return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))


def openai_gelu(x):
    return gelu_impl(x)


# This is actually Python equivalent of torch.nn.functional.gelu(), also with type hints for ONNX exporter
@torch.jit.script
def erf_gelu(x):
    return x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype) + torch.ones_like(x).to(dtype=x.dtype))


def print_named_parameters(model):
    """Print a summary of the parameters in a model."""
    prefix = ""
    if torch.distributed.is_initialized():
        # Print on only the first data parallel rank, but on all tensor/pipeline parallel ranks.
        should_print = mpu.get_data_parallel_rank() == 0
        if mpu.get_tensor_model_parallel_world_size() > 1:
            prefix = f"tensor-rank: {mpu.get_tensor_model_parallel_rank()} | "
        if mpu.get_pipeline_model_parallel_world_size() > 1:
            prefix = f"pipeline-rank: {mpu.get_pipeline_model_parallel_rank()} | "
    else:
        should_print = get_args().rank == 0

    if not should_print:
        return

    print(f"{prefix} > {type(model).__name__} parameters: ", flush=True)
    for name, param in model.named_parameters():
        if mpu.param_is_tensor_parallel_unique(param):
            print(f"{prefix}{name=}, {param.shape=}, norm={torch.norm(param.data.float()).item()}", flush=True)


def sync_data_parallel_replicated_parameters(a, b):
    pass


def sync_tensor_parallel_replicated_parameters(a, b):
    pass


def _sync_replicated_parameters(a, b, c, d, e):

    def _sync(p):
        pass


class ReplicationMismatchError(Exception):
    pass


class TensorParallelReplicationMismatchError(ReplicationMismatchError):
    pass


class DataParallelReplicationMismatchError(ReplicationMismatchError):
    pass


def validate_replicated_parameters(m, o):

    def collect_and_validate(a, b, c, d, prefix):
        pass


def _validate_replicated_parameters(
    *, a, b, c, d, prefix=""
):
    pass

megatron/model/utils.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Model parallel utility interface."""

from .cross_entropy import vocab_parallel_cross_entropy

from .initialize import is_unitialized
from .initialize import destroy_model_parallel
from .initialize import force_communicator_creation
from .initialize import get_data_parallel_group
from .initialize import get_data_parallel_rank
from .initialize import get_data_parallel_world_size
from .initialize import get_embedding_group
from .initialize import get_position_embedding_group
from .initialize import get_model_parallel_group
from .initialize import get_tensor_model_parallel_group
from .initialize import get_pipeline_model_parallel_group
from .initialize import get_tensor_model_parallel_rank, set_tensor_model_parallel_rank
from .initialize import get_pipeline_model_parallel_rank, set_pipeline_model_parallel_rank
from .initialize import is_pipeline_first_stage, is_pipeline_last_stage
from .initialize import is_rank_in_embedding_group
from .initialize import is_rank_in_position_embedding_group
from .initialize import is_pipeline_stage_before_split, is_pipeline_stage_after_split
from .initialize import is_pipeline_stage_at_split
from .initialize import get_num_layers
from .initialize import get_tensor_model_parallel_src_rank
from .initialize import get_data_parallel_src_rank
from .initialize import get_pipeline_model_parallel_first_rank
from .initialize import get_pipeline_model_parallel_last_rank
from .initialize import get_pipeline_model_parallel_next_rank
from .initialize import get_pipeline_model_parallel_prev_rank
from .initialize import get_tensor_model_parallel_world_size, set_tensor_model_parallel_world_size
from .initialize import get_pipeline_model_parallel_world_size, set_pipeline_model_parallel_world_size
from .initialize import get_virtual_pipeline_model_parallel_rank, set_virtual_pipeline_model_parallel_rank
from .initialize import initialize_model_parallel
from .initialize import model_parallel_is_initialized

from .layers import LinearWithGradAccumulationAndAsyncCommunication
from .layers import ColumnParallelLinear
from .layers import ParallelLinear
from .layers import RowParallelLinear
from .layers import VocabParallelEmbedding
from .layers import (
    set_tensor_model_parallel_attributes,
    set_defaults_if_not_set_tensor_model_parallel_attributes,
    copy_tensor_model_parallel_attributes,
    param_is_tensor_parallel_unique,
    param_is_tensor_parallel_replicated,
)

from .mappings import reduce_backward_from_tensor_model_parallel_region
from .mappings import reduce_forward_from_tensor_model_parallel_region
from .mappings import scatter_to_tensor_model_parallel_region
from .mappings import gather_from_tensor_model_parallel_region
from .mappings import scatter_to_sequence_parallel_region
from .mappings import gather_from_sequence_parallel_region
from .mappings import reduce_scatter_to_sequence_parallel_region

from .random import checkpoint
from .random import get_cuda_rng_tracker
from .random import model_parallel_cuda_manual_seed
from .random import gather_split_1d_tensor
from .random import split_tensor_into_1d_equal_chunks
from .random import make_viewless_tensor
from .random import assert_viewless_tensor
from .random import safely_set_viewless_tensor_data

from .utils import divide
from .utils import split_tensor_along_last_dim

megatron/mpu/__init__.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Communications utilities."""


from typing import Sequence, Union
import collections
import torch

from megatron import mpu
from megatron.model.utils import check_shapes


# TODO: use functions from megatron/p2p
def recv_from_prev_pipeline_rank_(recv_buffer=None):
    """Receive from previous pipeline stage and update the
    input buffer inplace."""
    if not mpu.is_pipeline_first_stage():
        assert recv_buffer is not None
        recv_prev_op = torch.distributed.P2POp(
            torch.distributed.irecv, recv_buffer, mpu.get_pipeline_model_parallel_prev_rank()
        )
        reqs = torch.distributed.batch_isend_irecv([recv_prev_op])
        for req in reqs:
            req.wait()
        # To protect against race condition when using batch_isend_irecv().
        torch.cuda.synchronize()


# TODO: use functions from megatron/p2p
def send_to_next_pipeline_rank(tensor=None):
    """Send output to the next pipeline stage."""
    if not mpu.is_pipeline_last_stage():
        assert tensor is not None
        send_next_op = torch.distributed.P2POp(
            torch.distributed.isend, tensor, mpu.get_pipeline_model_parallel_next_rank()
        )
        reqs = torch.distributed.batch_isend_irecv([send_next_op])
        for req in reqs:
            req.wait()
        # To protect against race condition when using batch_isend_irecv().
        torch.cuda.synchronize()


def _is_cuda(tensor):
    """Check if a tensor is not none and is cuda."""
    assert tensor is not None
    assert tensor.is_cuda


def _is_cuda_contiguous(tensor):
    """Check if a tensor is not none, is cuda, and is contiguous."""
    _is_cuda(tensor)
    assert tensor.is_contiguous()


def _check_shapes(tensor: torch.Tensor, size: Union[int, Sequence[int]], dtype: torch.dtype):
    """Check that the sending and receiver tensors will be compatible."""
    # This logic is required because torch.empty will promote a size of int to [int] and there are lots
    # of callers that rely on this convention.
    shape = size if isinstance(size, collections.abc.Sequence) else [size]

    check_shapes(tensor, expected_shapes=shape, expected_ndim=len(shape), expected_dtype=dtype)


def broadcast_from_last_pipeline_stage(size, dtype, tensor=None):
    """Broadcast a tensor from last pipeline stage to all ranks."""

    is_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=True)
    # If first stage and last state are the same, then there is no
    # pipeline parallelism and no need to communicate.
    if mpu.is_pipeline_first_stage() and is_last_stage:
        return tensor

    if is_last_stage:
        _is_cuda_contiguous(tensor)
        _check_shapes(tensor, size, dtype)
    else:
        tensor = torch.empty(size, dtype=dtype, device=torch.cuda.current_device())
    # Get the group and corresponding source rank.
    src = mpu.get_pipeline_model_parallel_last_rank()
    group = mpu.get_pipeline_model_parallel_group()
    torch.distributed.broadcast(tensor, src, group)

    return tensor


def broadcast_from_last_to_first_pipeline_stage(size, dtype, tensor=None):
    """Broadcast tensor values from last stage into the first stage."""

    is_last_stage = mpu.is_pipeline_last_stage()
    is_first_stage = mpu.is_pipeline_first_stage()
    # If first stage and last state are the same, then there is no
    # pipeline parallelism and no need to communicate.
    if is_first_stage and is_last_stage:
        return tensor
    # Only first and last stage pipeline stages need to be involved.
    if is_last_stage or is_first_stage:
        if is_last_stage:
            _is_cuda_contiguous(tensor)
            _check_shapes(tensor, size, dtype)
        else:
            tensor = torch.empty(size, dtype=dtype, device=torch.cuda.current_device())
        src = mpu.get_pipeline_model_parallel_last_rank()
        group = mpu.get_embedding_group()
        # Broadcast from last stage into the first stage.
        torch.distributed.broadcast(tensor, src, group)
    else:
        tensor = None

    return tensor


def copy_from_last_to_first_pipeline_stage(size, dtype, tensor=None):
    """Copy tensor values from last stage into the first stage.
    Note that the input tensor is updated in place."""

    is_last_stage = mpu.is_pipeline_last_stage()
    is_first_stage = mpu.is_pipeline_first_stage()
    # If first stage and last state are the same, then there is no
    # pipeline parallelism and no need to communicate.
    if is_first_stage and is_last_stage:
        return
    # Only first and last stage pipeline stages need to be involved.
    if is_last_stage or is_first_stage:
        _is_cuda(tensor)
        is_contiguous = tensor.is_contiguous()
        src = mpu.get_pipeline_model_parallel_last_rank()
        group = mpu.get_embedding_group()
        if is_contiguous:
            tensor_ = tensor
        else:
            if is_last_stage:
                tensor_ = tensor.contiguous()
            else:
                tensor_ = torch.empty(size, dtype=dtype, device=torch.cuda.current_device())
        # Broadcast from last stage into the first stage.
        torch.distributed.broadcast(tensor_, src, group)
        # Update the first stage tensor
        if is_first_stage and not is_contiguous:
            tensor[...] = tensor_


def broadcast_tensor(size, dtype, tensor=None, rank=0, device=0):
    """Given size and type of a tensor on all ranks and the tensor value
    only on a specific rank, broadcast from that rank to all other ranks.
    Args:
        size: size of the tensor
        dtype: type of the tensor
        tensor: tensor to be broadcasted
        rank: primary rank for broadcasting
        device: device of the tensor. If not set to None, then we use cuda.current_device().
            Default is 0, since we use cuda.current_device() to get the device.
    """
    if device is not None:
        device = torch.cuda.current_device()
    if torch.distributed.get_rank() == rank:
        if device is not None:
            _is_cuda_contiguous(tensor)
        _check_shapes(tensor, size, dtype)
    else:
        tensor = torch.empty(size, dtype=dtype, device=device)
    torch.distributed.broadcast(tensor, rank)
    return tensor


def broadcast_list(size, dtype, list_values=None, rank=0, device=0):
    """Broadcast a list of values with a given type.
    Args:
        size: size of the list
        dtype: dtype of the list
        list_values: list of values to be broadcasted
        rank: primary rank for broadcasting
        device: device of the tensor. If not set to None, then we use cuda.current_device().
            Default is 0, since we use cuda.current_device() to get the device.
    """
    tensor = None
    if device is not None:
        device = torch.cuda.current_device()
    if torch.distributed.get_rank() == rank:
        tensor = torch.tensor(list_values, dtype=dtype, device=device)
    return broadcast_tensor(size, dtype, tensor=tensor, rank=rank, device=device)


def broadcast_int_list(size, int_list=None, rank=0, device=0):
    """Broadcast a list of interger values.
    Args:
        size: size of the list
        int_list: list of values to be broadcasted
        rank: primary rank for broadcasting
        device: device of the tensor. If not set to None, then we use cuda.current_device().
            Default is 0, since we use cuda.current_device() to get the device.
    """
    if device is not None:
        device = torch.cuda.current_device()
    return broadcast_list(size, torch.int64, list_values=int_list, rank=rank, device=device)


def broadcast_float_list(size, float_list=None, rank=0, device=0):
    """Broadcast a list of float values.
    Args:
        size: size of the list
        float_list: list of values to be broadcasted
        rank: primary rank for broadcasting
        device: device of the tensor. If not set to None, then we use cuda.current_device().
            Default is 0, since we use cuda.current_device() to get the device.
    """
    if device is not None:
        device = torch.cuda.current_device()
    return broadcast_list(size, torch.float32, list_values=float_list, rank=rank, device=device)

megatron/mpu/communication.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import torch

from .initialize import (
    get_tensor_model_parallel_group,
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
)
from .utils import VocabUtility


class _VocabParallelCrossEntropy(torch.autograd.Function):
    @staticmethod
    def forward(ctx, vocab_parallel_logits, target):

        # Maximum value along vocab dimension across all GPUs.
        logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
        torch.distributed.all_reduce(
            logits_max, op=torch.distributed.ReduceOp.MAX, group=get_tensor_model_parallel_group()
        )
        # Subtract the maximum value.
        vocab_parallel_logits.sub_(logits_max.unsqueeze(dim=-1))

        # Get the partition's vocab indecies
        get_vocab_range = VocabUtility.vocab_range_from_per_partition_vocab_size
        partition_vocab_size = vocab_parallel_logits.size()[-1]
        rank = get_tensor_model_parallel_rank()
        world_size = get_tensor_model_parallel_world_size()
        vocab_start_index, vocab_end_index = get_vocab_range(partition_vocab_size, rank, world_size)

        # Create a mask of valid vocab ids (1 means it needs to be masked).
        target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
        masked_target = target.clone() - vocab_start_index
        masked_target[target_mask] = 0

        # Get predicted-logits = logits[target].
        # For Simplicity, we convert logits to a 2-D tensor with size
        # [*, partition-vocab-size] and target to a 1-D tensor of size [*].
        logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
        masked_target_1d = masked_target.view(-1)
        arange_1d = torch.arange(start=0, end=logits_2d.size()[0], device=logits_2d.device)
        predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
        predicted_logits_1d = predicted_logits_1d.clone().contiguous()
        predicted_logits = predicted_logits_1d.view_as(target)
        predicted_logits[target_mask] = 0.0
        # All reduce is needed to get the chunks from other GPUs.
        torch.distributed.all_reduce(
            predicted_logits, op=torch.distributed.ReduceOp.SUM, group=get_tensor_model_parallel_group()
        )

        # Sum of exponential of logits along vocab dimension across all GPUs.
        exp_logits = vocab_parallel_logits
        torch.exp(vocab_parallel_logits, out=exp_logits)
        sum_exp_logits = exp_logits.sum(dim=-1)
        torch.distributed.all_reduce(
            sum_exp_logits, op=torch.distributed.ReduceOp.SUM, group=get_tensor_model_parallel_group()
        )

        # Loss = log(sum(exp(logits))) - predicted-logit.
        loss = torch.log(sum_exp_logits) - predicted_logits

        # Store softmax, target-mask and masked-target for backward pass.
        exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
        ctx.save_for_backward(exp_logits, target_mask, masked_target_1d)

        return loss

    @staticmethod
    def backward(ctx, grad_output):

        # Retreive tensors from the forward path.
        softmax, target_mask, masked_target_1d = ctx.saved_tensors

        # All the inputs have softmax as thier gradient.
        grad_input = softmax
        # For simplicity, work with the 2D gradient.
        partition_vocab_size = softmax.size()[-1]
        grad_2d = grad_input.view(-1, partition_vocab_size)

        # Add the gradient from matching classes.
        arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device)
        grad_2d[arange_1d, masked_target_1d] -= 1.0 - target_mask.view(-1).float()

        # Finally elementwise multiplication with the output gradients.
        grad_input.mul_(grad_output.unsqueeze(dim=-1))

        return grad_input, None


def vocab_parallel_cross_entropy(vocab_parallel_logits, target):
    """Helper function for the cross entropy."""
    return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target)

megatron/mpu/cross_entropy.py
persimmon-ai-labs/adept-inference
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Taken from https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/losses/cross_entropy_parallel.py
# But we import from Megatron instead of from Apex
# HazyResearch is licensed under BSD 3.0

# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/cross_entropy.py
# But we make it much faster: we compute the local loss and the LSE, and by exchanging the LSE and
# the losses we can get the global loss. There's no need to do it step by step
# (compute local max, exchange, compute exp, compute local sum, exchange, etc.)
import torch
import torch.nn as nn

import xentropy_cuda_lib

from .initialize import get_tensor_model_parallel_group
from .initialize import get_tensor_model_parallel_rank
from .initialize import get_tensor_model_parallel_world_size
from .utils import VocabUtility

# `all_gather_into_tensor` and `reduce_scatter_tensor` are new placeholders for
# `_all_gather_base` and `_reduce_scatter_base`. They require the most recent
# version of PyTorch. The following 4 lines are for backward compatibility with
# older PyTorch.
if "all_gather_into_tensor" not in dir(torch.distributed):
    torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base
if "reduce_scatter_tensor" not in dir(torch.distributed):
    torch.distributed.reduce_scatter_tensor = torch.distributed._reduce_scatter_base


class SoftmaxCrossEntropyLossParallelFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, logits_parallel, labels, smoothing=0.0, ignored_index=-100, inplace_backward=False):
        """
        logits_parallel: (batch, vocab_size / world_size)
        labels: (batch,)
        """
        batch, partition_vocab_size = logits_parallel.shape
        assert labels.shape == (batch,)
        rank = get_tensor_model_parallel_rank()
        world_size = get_tensor_model_parallel_world_size()

        if world_size == 1:
            losses, lse = xentropy_cuda_lib.forward(logits_parallel, labels, smoothing)
            losses.masked_fill_(labels == ignored_index, 0)
            labels_local = labels
        else:
            vocab_start_index, vocab_end_index = VocabUtility.vocab_range_from_per_partition_vocab_size(
                partition_vocab_size, get_tensor_model_parallel_rank(), get_tensor_model_parallel_world_size()
            )

            # Create a mask of valid vocab ids (1 means it needs to be masked).
            labels_mask = (labels < vocab_start_index) | (labels >= vocab_end_index)
            ignored_mask = labels == ignored_index
            labels_local = torch.where(ignored_mask, labels, labels - vocab_start_index)
            masked_labels = labels_local.clone()
            masked_labels[labels_mask] = ignored_index

            losses, lse_local = xentropy_cuda_lib.forward(logits_parallel, masked_labels, smoothing)
            assert lse_local.shape == (batch,)
            assert losses.shape == (batch,)
            losses.masked_fill_(masked_labels == ignored_index, 0)

            lse_allgather = torch.empty(world_size, batch, dtype=lse_local.dtype, device=lse_local.device)
            handle_lse = torch.distributed.all_gather_into_tensor(
                lse_allgather, lse_local.contiguous(), group=get_tensor_model_parallel_group(), async_op=True
            )
            handle_losses = torch.distributed.all_reduce(
                losses, op=torch.distributed.ReduceOp.SUM, group=get_tensor_model_parallel_group(), async_op=True
            )
            handle_lse.wait()
            lse = torch.logsumexp(lse_allgather, dim=0)
            # The losses are going to be lse_local - predicted_logit, we just have to subtract
            # the lse_local and add the lse (global).
            rank_per_sample = torch.div(labels, partition_vocab_size, rounding_mode="floor")
            lse_local = lse_allgather[rank_per_sample, torch.arange(batch, device=lse_allgather.device)]

            handle_losses.wait()
            losses += lse - lse_local
            losses.masked_fill_(ignored_mask, 0)

        ctx.save_for_backward(logits_parallel, lse, labels_local)
        ctx.smoothing = smoothing
        ctx.ignored_index = ignored_index
        ctx.inplace_backward = inplace_backward
        return losses

    @staticmethod
    def backward(ctx, grad_loss):
        logits_parallel, lse, labels = ctx.saved_tensors
        if not grad_loss.is_contiguous():
            grad_loss = grad_loss.contiguous()
        grad_loss.masked_fill_(labels == ctx.ignored_index, 0)
        grad_logits = xentropy_cuda_lib.backward(
            grad_loss, logits_parallel, lse, labels, ctx.smoothing, ctx.inplace_backward
        )
        return grad_logits, None, None, None, None, None


class CrossEntropyLossParallel(nn.Module):
    def __init__(self, ignore_index=-100, reduction="mean", label_smoothing=0.0, inplace_backward=False):
        super().__init__()
        if reduction not in ["mean", "none"]:
            raise NotImplementedError("Only support reduction = 'mean' or 'none'")
        self.ignore_index = ignore_index
        self.reduction = reduction
        self.label_smoothing = label_smoothing
        self.inplace_backward = inplace_backward

    def forward(self, input, target):
        assert input.is_cuda and target.is_cuda
        # SoftmaxCrossEntropyLoss implicitly casts to float
        loss = SoftmaxCrossEntropyLossParallelFn.apply(
            input, target, self.label_smoothing, self.ignore_index, self.inplace_backward
        )
        if self.reduction == "mean":
            return loss.sum() / (target != self.ignore_index).sum()
        else:
            return loss

megatron/mpu/cross_entropy_parallel.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""Model and data parallel groups."""

import os
import sys
import torch
import warnings

from .utils import ensure_divisibility


# Intra-layer model parallel group that the current rank belongs to.
_TENSOR_MODEL_PARALLEL_GROUP = None
# Inter-layer model parallel group that the current rank belongs to.
_PIPELINE_MODEL_PARALLEL_GROUP = None
# Model parallel group (both intra- and pipeline) that the current rank belongs to.
_MODEL_PARALLEL_GROUP = None
# Embedding group.
_EMBEDDING_GROUP = None
# Position embedding group.
_POSITION_EMBEDDING_GROUP = None
# Data parallel group that the current rank belongs to.
_DATA_PARALLEL_GROUP = None

_VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = None
_VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
_PIPELINE_MODEL_PARALLEL_SPLIT_RANK = None

# These values enable us to change the mpu sizes on the fly.
_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None
_MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
_MPU_TENSOR_MODEL_PARALLEL_RANK = None
_MPU_PIPELINE_MODEL_PARALLEL_RANK = None

# A list of ranks that have a copy of the embedding.
_EMBEDDING_GLOBAL_RANKS = None

# A list of ranks that have a copy of the position embedding.
_POSITION_EMBEDDING_GLOBAL_RANKS = None

# A list of global ranks for each pipeline group to ease calculation of the source
# rank when broadcasting from the first or last pipeline stage.
_PIPELINE_GLOBAL_RANKS = None

# A list of global ranks for each data parallel group to ease calculation of the source
# rank when broadcasting weights from src to all other data parallel ranks
_DATA_PARALLEL_GLOBAL_RANKS = None


def is_unitialized():
    """Useful for code segments that may be accessed with or without mpu initialization"""
    return _DATA_PARALLEL_GROUP is None


def initialize_model_parallel(
    tensor_model_parallel_size_=1,
    pipeline_model_parallel_size_=1,
    virtual_pipeline_model_parallel_size_=None,
    pipeline_model_parallel_split_rank_=None,
    default_backend=None,
    p2p_backend=None,
    cross_block_networking=False,
    share_word_embeddings=False,
    force_socket_networking=False,
):
    """
    Initialize model data parallel groups.

    Arguments:
        tensor_model_parallel_size: number of GPUs used for tensor model parallelism.
        pipeline_model_parallel_size: number of GPUs used for pipeline model parallelism.
        virtual_pipeline_model_parallel_size: number of virtual stages (interleaved
                                              pipeline).
        pipeline_model_parallel_split_rank: for models with both encoder and decoder,
                                            rank in pipeline with split point.
        default_backend: Backend of process groups except for pipeline parallel ones.
            If :obj:`None`, the backend specified in `torch.distributed.init_process_group` will be used.
        p2p_backend: Backend of process groups for pipeline model parallel.
            If :obj:`None`, the backend specified in `torch.distributed.init_process_group` will be used.
        cross_block_networking: Use Socket networking where necessary to ensure cross block communication
            doesn't use IB networking which doesn't span blocks.
        share_word_embeddings: True if the input and output embeddings are to be kept in sync.
        force_socket_networking: Use Socket networking for all communication.


    Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
    use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
    the model pipeline. The present function will
    create 8 tensor model-parallel groups, 4 pipeline model-parallel groups
    and 8 data-parallel groups as:
        8 data_parallel groups:
            [g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]
        8 tensor model-parallel groups:
            [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15]
        4 pipeline model-parallel groups:
            [g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15]
    Note that for efficiency, the caller should make sure adjacent ranks
    are on the same DGX box. For example if we are using 2 DGX-1 boxes
    with a total of 16 GPUs, rank 0 to 7 belong to the first box and
    ranks 8 to 15 belong to the second box.
    """
    if torch.distributed.get_rank() == 0:
        print("> initializing tensor model parallel with size {}".format(tensor_model_parallel_size_))
        print("> initializing pipeline model parallel with size {}".format(pipeline_model_parallel_size_))
    # Get world size and rank. Ensure some consistencies.
    assert torch.distributed.is_initialized()
    assert default_backend is None or default_backend in ("nccl", "ucc")
    assert p2p_backend is None or p2p_backend in ("nccl", "ucc")
    if "ucc" in (default_backend, p2p_backend):
        check_torch_ucc_availability()
    if default_backend == "ucc":
        warnings.warn("The UCC's functionality as `default_backend` is not well verified", ExperimentalWarning)

    world_size = torch.distributed.get_world_size()
    tensor_model_parallel_size = min(tensor_model_parallel_size_, world_size)
    pipeline_model_parallel_size = min(pipeline_model_parallel_size_, world_size)
    ensure_divisibility(world_size, tensor_model_parallel_size * pipeline_model_parallel_size)
    data_parallel_size = world_size // (tensor_model_parallel_size * pipeline_model_parallel_size)

    num_tensor_model_parallel_groups = world_size // tensor_model_parallel_size
    num_pipeline_model_parallel_groups = world_size // pipeline_model_parallel_size
    num_data_parallel_groups = world_size // data_parallel_size

    if virtual_pipeline_model_parallel_size_ is not None:
        global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
        global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
        _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = 0
        _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = virtual_pipeline_model_parallel_size_

    if pipeline_model_parallel_split_rank_ is not None:
        global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
        _PIPELINE_MODEL_PARALLEL_SPLIT_RANK = pipeline_model_parallel_split_rank_

    rank = torch.distributed.get_rank()

    # Build the data-parallel groups.
    global _DATA_PARALLEL_GROUP
    global _DATA_PARALLEL_GLOBAL_RANKS
    assert _DATA_PARALLEL_GROUP is None, "data parallel group is already initialized"
    all_data_parallel_group_ranks = []
    for i in range(pipeline_model_parallel_size):
        start_rank = i * num_pipeline_model_parallel_groups
        end_rank = (i + 1) * num_pipeline_model_parallel_groups
        for j in range(tensor_model_parallel_size):
            ranks = range(start_rank + j, end_rank, tensor_model_parallel_size)
            all_data_parallel_group_ranks.append(list(ranks))
            group = torch.distributed.new_group(ranks, backend=default_backend)
            if rank in ranks:
                _DATA_PARALLEL_GROUP = group
                _DATA_PARALLEL_GLOBAL_RANKS = ranks

    if force_socket_networking:
        os.environ["NCCL_NET"] = "Socket"
    else:
        os.environ["NCCL_NET"] = "IB"
    force_communicator_creation(
        rank=get_data_parallel_rank(),
        world_size=get_data_parallel_world_size(),
        all_reduce=True,
        all_gather=True,
        group=get_data_parallel_group(),
    )

    # Build the model-parallel groups.
    global _MODEL_PARALLEL_GROUP
    assert _MODEL_PARALLEL_GROUP is None, "model parallel group is already initialized"
    for i in range(data_parallel_size):
        ranks = [data_parallel_group_ranks[i] for data_parallel_group_ranks in all_data_parallel_group_ranks]
        group = torch.distributed.new_group(ranks, backend=default_backend)
        if rank in ranks:
            _MODEL_PARALLEL_GROUP = group

    if cross_block_networking or force_socket_networking:
        os.environ["NCCL_NET"] = "Socket"
    else:
        os.environ["NCCL_NET"] = "IB"
    force_communicator_creation(all_reduce=True, group=get_model_parallel_group())

    # Build the tensor model-parallel groups.
    global _TENSOR_MODEL_PARALLEL_GROUP
    assert _TENSOR_MODEL_PARALLEL_GROUP is None, "tensor model parallel group is already initialized"
    for i in range(num_tensor_model_parallel_groups):
        ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size)
        group = torch.distributed.new_group(ranks, backend=default_backend)
        if rank in ranks:
            # torch.distributed.barrier(group)
            _TENSOR_MODEL_PARALLEL_GROUP = group

    if force_socket_networking:
        os.environ["NCCL_NET"] = "Socket"
    else:
        os.environ["NCCL_NET"] = "IB"
    force_communicator_creation(
        rank=get_tensor_model_parallel_rank(),
        src_rank=get_tensor_model_parallel_src_rank(),
        world_size=get_tensor_model_parallel_world_size(),
        all_reduce=True,
        all_gather=True,
        broadcast=True,
        group=get_tensor_model_parallel_group(),
    )

    # Build the pipeline model-parallel groups and embedding groups
    # (first and last rank in each pipeline model-parallel group).
    global _PIPELINE_MODEL_PARALLEL_GROUP
    global _PIPELINE_GLOBAL_RANKS
    assert _PIPELINE_MODEL_PARALLEL_GROUP is None, "pipeline model parallel group is already initialized"
    global _EMBEDDING_GROUP
    global _EMBEDDING_GLOBAL_RANKS
    assert _EMBEDDING_GROUP is None, "embedding group is already initialized"
    global _POSITION_EMBEDDING_GROUP
    global _POSITION_EMBEDDING_GLOBAL_RANKS
    assert _POSITION_EMBEDDING_GROUP is None, "position embedding group is already initialized"
    for i in range(num_pipeline_model_parallel_groups):
        ranks = range(i, world_size, num_pipeline_model_parallel_groups)
        group = torch.distributed.new_group(ranks, backend=p2p_backend)
        if rank in ranks:
            _PIPELINE_MODEL_PARALLEL_GROUP = group
            _PIPELINE_GLOBAL_RANKS = ranks
        # Setup embedding group (to exchange gradients between
        # first and last stages).
        if len(ranks) > 1:
            embedding_ranks = [ranks[0], ranks[-1]]
            position_embedding_ranks = [ranks[0]]
            if pipeline_model_parallel_split_rank_ is not None:
                if ranks[pipeline_model_parallel_split_rank_] not in embedding_ranks:
                    embedding_ranks = [ranks[0], ranks[pipeline_model_parallel_split_rank_], ranks[-1]]
                if ranks[pipeline_model_parallel_split_rank_] not in position_embedding_ranks:
                    position_embedding_ranks = [ranks[0], ranks[pipeline_model_parallel_split_rank_]]
        else:
            embedding_ranks = ranks
            position_embedding_ranks = ranks

        group = torch.distributed.new_group(embedding_ranks, backend=default_backend)
        if rank in embedding_ranks:
            _EMBEDDING_GROUP = group
        if rank in ranks:
            _EMBEDDING_GLOBAL_RANKS = embedding_ranks

        group = torch.distributed.new_group(position_embedding_ranks, backend=default_backend)
        if rank in position_embedding_ranks:
            _POSITION_EMBEDDING_GROUP = group
        if rank in ranks:
            _POSITION_EMBEDDING_GLOBAL_RANKS = position_embedding_ranks

    if cross_block_networking or force_socket_networking:
        os.environ["NCCL_NET"] = "Socket"
    else:
        os.environ["NCCL_NET"] = "IB"
    force_pipeline_communicator_creation(virtual_pipeline_model_parallel_size_ is not None)

    if share_word_embeddings and is_rank_in_embedding_group(ignore_virtual=True):
        force_communicator_creation(all_reduce=True, group=get_embedding_group())

    if is_rank_in_position_embedding_group() and pipeline_model_parallel_split_rank_ is not None:
        force_communicator_creation(all_reduce=True, group=get_position_embedding_group())
    if force_socket_networking:
        os.environ["NCCL_NET"] = "Socket"
    else:
        os.environ["NCCL_NET"] = "IB"


# Used to warn when the UCC is specified.
class ExperimentalWarning(Warning):
    pass


def check_torch_ucc_availability():
    try:
        import torch_ucc  # NOQA
    except ImportError:
        raise ImportError(
            "UCC backend requires [torch_ucc](https://github.com/facebookresearch/torch_ucc) but not found"
        )


def model_parallel_is_initialized():
    """Check if model and data parallel groups are initialized."""
    if _TENSOR_MODEL_PARALLEL_GROUP is None or _PIPELINE_MODEL_PARALLEL_GROUP is None or _DATA_PARALLEL_GROUP is None:
        return False
    return True


def get_model_parallel_group():
    """Get the model parallel group the caller rank belongs to."""
    assert _MODEL_PARALLEL_GROUP is not None, "model parallel group is not initialized"
    return _MODEL_PARALLEL_GROUP


def get_tensor_model_parallel_group():
    """Get the tensor model parallel group the caller rank belongs to."""
    assert _TENSOR_MODEL_PARALLEL_GROUP is not None, "intra_layer_model parallel group is not initialized"
    return _TENSOR_MODEL_PARALLEL_GROUP


def get_pipeline_model_parallel_group():
    """Get the pipeline model parallel group the caller rank belongs to."""
    assert _PIPELINE_MODEL_PARALLEL_GROUP is not None, "pipeline_model parallel group is not initialized"
    return _PIPELINE_MODEL_PARALLEL_GROUP


def get_data_parallel_group():
    """Get the data parallel group the caller rank belongs to."""
    assert _DATA_PARALLEL_GROUP is not None, "data parallel group is not initialized"
    return _DATA_PARALLEL_GROUP


def get_embedding_group():
    """Get the embedding group the caller rank belongs to."""
    assert _EMBEDDING_GROUP is not None, "embedding group is not initialized"
    return _EMBEDDING_GROUP


def get_position_embedding_group():
    """Get the position embedding group the caller rank belongs to."""
    assert _POSITION_EMBEDDING_GROUP is not None, "position embedding group is not initialized"
    return _POSITION_EMBEDDING_GROUP


def set_tensor_model_parallel_world_size(world_size):
    """Set the tensor model parallel size"""
    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
    _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size


def set_pipeline_model_parallel_world_size(world_size):
    """Set the pipeline model parallel size"""
    global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
    _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = world_size


def get_tensor_model_parallel_world_size():
    """Return world size for the tensor model parallel group."""
    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
    if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None:
        return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
    return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())


def get_pipeline_model_parallel_world_size():
    """Return world size for the pipeline model parallel group."""
    global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
    if _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE is not None:
        return _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
    return torch.distributed.get_world_size(group=get_pipeline_model_parallel_group())


def set_tensor_model_parallel_rank(rank):
    """Set tensor model parallel rank."""
    global _MPU_TENSOR_MODEL_PARALLEL_RANK
    _MPU_TENSOR_MODEL_PARALLEL_RANK = rank


def set_pipeline_model_parallel_rank(rank):
    """Set pipeline model parallel rank."""
    global _MPU_PIPELINE_MODEL_PARALLEL_RANK
    _MPU_PIPELINE_MODEL_PARALLEL_RANK = rank


def get_tensor_model_parallel_rank():
    """Return my rank for the tensor model parallel group."""
    global _MPU_TENSOR_MODEL_PARALLEL_RANK
    if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None:
        return _MPU_TENSOR_MODEL_PARALLEL_RANK
    return torch.distributed.get_rank(group=get_tensor_model_parallel_group())


def get_pipeline_model_parallel_rank():
    """Return my rank for the pipeline model parallel group."""
    global _MPU_PIPELINE_MODEL_PARALLEL_RANK
    if _MPU_PIPELINE_MODEL_PARALLEL_RANK is not None:
        return _MPU_PIPELINE_MODEL_PARALLEL_RANK
    return torch.distributed.get_rank(group=get_pipeline_model_parallel_group())


def get_num_layers(
    pipeline_model_parallel_split_rank: int,
    standalone_embedding_stage: bool,
    transformer_pipeline_model_parallel_size: int,
    encoder_num_layers: int,
    decoder_num_layers: int,
    num_layers: int,
    is_encoder_and_decoder_model: bool,
    is_decoder: bool = False,
):
    """Compute the number of transformer layers resident on the current rank."""
    if get_pipeline_model_parallel_world_size() > 1:
        if is_encoder_and_decoder_model:
            assert pipeline_model_parallel_split_rank is not None

            # When a standalone embedding stage is used, a rank is taken from
            # the encoder's ranks, to be used for the encoder's embedding
            # layer. This way, the rank referenced by the 'split rank' remains
            # the same whether or not a standalone embedding stage is used.
            num_ranks_in_encoder = (
                pipeline_model_parallel_split_rank - 1
                if standalone_embedding_stage
                else pipeline_model_parallel_split_rank
            )
            num_ranks_in_decoder = transformer_pipeline_model_parallel_size - num_ranks_in_encoder
            assert (
                encoder_num_layers % num_ranks_in_encoder == 0
            ), "encoder_num_layers (%d) must be divisible by number of ranks given to encoder (%d)" % (
                encoder_num_layers,
                num_ranks_in_encoder,
            )
            assert (
                decoder_num_layers % num_ranks_in_decoder == 0
            ), "decoder_num_layers (%d) must be divisible by number of ranks given to decoder (%d)" % (
                decoder_num_layers,
                num_ranks_in_decoder,
            )
            if is_pipeline_stage_before_split():
                num_layers = (
                    0
                    if standalone_embedding_stage and get_pipeline_model_parallel_rank() == 0
                    else encoder_num_layers // num_ranks_in_encoder
                )
            else:
                num_layers = decoder_num_layers // num_ranks_in_decoder
        else:
            assert num_layers == encoder_num_layers
            assert (
                num_layers % transformer_pipeline_model_parallel_size == 0
            ), "num_layers must be divisible by transformer_pipeline_model_parallel_size"

            # When a standalone embedding stage is used, all transformer layers
            # are divided among pipeline rank >= 1, while on pipeline rank 0,
            # ranks either contain the input embedding layer (virtual pp rank 0),
            # or no layers at all (virtual pp rank >= 1).
            num_layers = (
                0
                if standalone_embedding_stage and get_pipeline_model_parallel_rank() == 0
                else num_layers // transformer_pipeline_model_parallel_size
            )
    else:
        if not is_decoder:
            num_layers = encoder_num_layers
        else:
            num_layers = decoder_num_layers
    return num_layers


def is_pipeline_first_stage(ignore_virtual=False):
    """Return True if in the first pipeline model-parallel stage, False otherwise."""
    if not ignore_virtual:
        if (
            get_virtual_pipeline_model_parallel_world_size() is not None
            and get_virtual_pipeline_model_parallel_rank() != 0
        ):
            return False
    return get_pipeline_model_parallel_rank() == 0


def is_pipeline_last_stage(ignore_virtual=False):
    """Return True if in the last pipeline model-parallel stage, False otherwise."""
    if not ignore_virtual:
        virtual_pipeline_model_parallel_world_size = get_virtual_pipeline_model_parallel_world_size()
        if virtual_pipeline_model_parallel_world_size is not None and get_virtual_pipeline_model_parallel_rank() != (
            virtual_pipeline_model_parallel_world_size - 1
        ):
            return False
    return get_pipeline_model_parallel_rank() == (get_pipeline_model_parallel_world_size() - 1)


def is_rank_in_embedding_group(ignore_virtual=False):
    """Return true if current rank is in embedding group, False otherwise."""
    rank = torch.distributed.get_rank()
    global _EMBEDDING_GLOBAL_RANKS
    if ignore_virtual:
        return rank in _EMBEDDING_GLOBAL_RANKS
    if rank in _EMBEDDING_GLOBAL_RANKS:
        if rank == _EMBEDDING_GLOBAL_RANKS[0]:
            return is_pipeline_first_stage(ignore_virtual=False)
        elif rank == _EMBEDDING_GLOBAL_RANKS[-1]:
            return is_pipeline_last_stage(ignore_virtual=False)
        else:
            return True
    return False


def is_rank_in_position_embedding_group():
    """Return true if current rank is in position embedding group, False otherwise."""
    rank = torch.distributed.get_rank()
    global _POSITION_EMBEDDING_GLOBAL_RANKS
    return rank in _POSITION_EMBEDDING_GLOBAL_RANKS


def is_pipeline_stage_before_split(rank=None):
    """Return True if pipeline stage executes encoder block for a model
    with both encoder and decoder."""
    if get_pipeline_model_parallel_world_size() == 1:
        return True
    if rank is None:
        rank = get_pipeline_model_parallel_rank()
    global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
    if _PIPELINE_MODEL_PARALLEL_SPLIT_RANK is None:
        return True
    if rank < _PIPELINE_MODEL_PARALLEL_SPLIT_RANK:
        return True
    return False


def is_pipeline_stage_after_split(rank=None):
    """Return True if pipeline stage executes decoder block for a model
    with both encoder and decoder."""
    if get_pipeline_model_parallel_world_size() == 1:
        return True
    if rank is None:
        rank = get_pipeline_model_parallel_rank()
    global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
    if _PIPELINE_MODEL_PARALLEL_SPLIT_RANK is None:
        return True
    if rank >= _PIPELINE_MODEL_PARALLEL_SPLIT_RANK:
        return True
    return False


def is_pipeline_stage_at_split():
    """Return true if pipeline stage executes decoder block and next
    stage executes encoder block for a model with both encoder and
    decoder."""
    rank = get_pipeline_model_parallel_rank()
    return is_pipeline_stage_before_split(rank) and is_pipeline_stage_after_split(rank + 1)


def get_virtual_pipeline_model_parallel_rank():
    """Return the virtual pipeline-parallel rank."""
    global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
    return _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK


def set_virtual_pipeline_model_parallel_rank(rank):
    """Set the virtual pipeline-parallel rank."""
    global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
    _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = rank


def get_virtual_pipeline_model_parallel_world_size():
    """Return the virtual pipeline-parallel world size."""
    global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
    return _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE


def get_tensor_model_parallel_src_rank():
    """Calculate the global rank corresponding to the first local rank
    in the tensor model parallel group."""
    global_rank = torch.distributed.get_rank()
    local_world_size = get_tensor_model_parallel_world_size()
    return (global_rank // local_world_size) * local_world_size


def get_data_parallel_src_rank():
    """Calculate the global rank corresponding to the first local rank
    in the data parallel group."""
    assert _DATA_PARALLEL_GLOBAL_RANKS is not None, "Data parallel group is not initialized"
    return _DATA_PARALLEL_GLOBAL_RANKS[0]


def get_pipeline_model_parallel_first_rank():
    assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
    return _PIPELINE_GLOBAL_RANKS[0]


def get_pipeline_model_parallel_last_rank():
    assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
    last_rank_local = get_pipeline_model_parallel_world_size() - 1
    return _PIPELINE_GLOBAL_RANKS[last_rank_local]


def get_pipeline_model_parallel_next_rank():
    assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
    rank_in_pipeline = get_pipeline_model_parallel_rank()
    world_size = get_pipeline_model_parallel_world_size()
    return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline + 1) % world_size]


def get_pipeline_model_parallel_prev_rank():
    assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
    rank_in_pipeline = get_pipeline_model_parallel_rank()
    world_size = get_pipeline_model_parallel_world_size()
    return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline - 1) % world_size]


def get_data_parallel_world_size():
    """Return world size for the data parallel group."""
    return torch.distributed.get_world_size(group=get_data_parallel_group())


def get_data_parallel_rank():
    """Return my rank for the data parallel group."""
    return torch.distributed.get_rank(group=get_data_parallel_group())


def force_communicator_creation(
    rank=None,
    src_rank=None,
    world_size=None,
    group=None,
    all_reduce=False,
    all_gather=False,
    barrier=False,
    broadcast=False,
):
    """Will do trivial ops of the requested type so that pytorch is forced to create
    a communicator.  This is useful because we control NCCL parameters related
    to each communicator through changing environment variables.  It is impossible
    to make sure the right env variables are set whereever the first use happens.

    By using this function, we can make sure to set the right env variables right before
    calling it."""
    if barrier:
        torch.distributed.barrier(group=group)
    if all_reduce:
        one_tensor = torch.cuda.FloatTensor([1.0])
        torch.distributed.all_reduce(one_tensor, op=torch.distributed.ReduceOp.SUM, group=group)
    if all_gather:
        assert (
            rank is not None and world_size is not None
        ), "Must supply rank and world_size for all_gather initialization"
        tensor_list = [torch.empty_like(one_tensor) for _ in range(world_size)]
        tensor_list[rank] = one_tensor
        torch.distributed.all_gather(tensor_list, one_tensor, group=group)
    if broadcast:
        one_tensor = torch.cuda.FloatTensor([1.0])
        torch.distributed.broadcast(one_tensor, src_rank, group=group)


def force_pipeline_communicator_creation(ignore_virtual=False):
    """Will do trivial ops for each of the pipeline communications to make sure that
    pytorch is forced to create the require communicators. See docstring for
    force_communicator_creation to understand how this is useful."""
    if get_pipeline_model_parallel_world_size() == 1:
        return
    recv_prev_tensor = torch.cuda.FloatTensor([1.0])
    recv_next_tensor = torch.cuda.FloatTensor([1.0])
    send_prev_tensor = torch.cuda.FloatTensor([1.0])
    send_next_tensor = torch.cuda.FloatTensor([1.0])
    ops = []
    if not is_pipeline_first_stage(ignore_virtual):
        recv_prev_op = torch.distributed.P2POp(
            torch.distributed.irecv,
            recv_prev_tensor,
            get_pipeline_model_parallel_prev_rank(),
            group=get_pipeline_model_parallel_group(),
        )
        send_prev_op = torch.distributed.P2POp(
            torch.distributed.isend,
            send_prev_tensor,
            get_pipeline_model_parallel_prev_rank(),
            group=get_pipeline_model_parallel_group(),
        )
        ops.append(recv_prev_op)
        ops.append(send_prev_op)
    if not is_pipeline_last_stage(ignore_virtual):
        send_next_op = torch.distributed.P2POp(
            torch.distributed.isend,
            send_next_tensor,
            get_pipeline_model_parallel_next_rank(),
            group=get_pipeline_model_parallel_group(),
        )
        recv_next_op = torch.distributed.P2POp(
            torch.distributed.irecv,
            recv_next_tensor,
            get_pipeline_model_parallel_next_rank(),
            group=get_pipeline_model_parallel_group(),
        )
        ops.append(send_next_op)
        ops.append(recv_next_op)

    reqs = torch.distributed.batch_isend_irecv(ops)
    for req in reqs:
        req.wait()


def destroy_model_parallel():
    """Set the groups to none."""
    global _MODEL_PARALLEL_GROUP
    _MODEL_PARALLEL_GROUP = None
    global _TENSOR_MODEL_PARALLEL_GROUP
    _TENSOR_MODEL_PARALLEL_GROUP = None
    global _PIPELINE_MODEL_PARALLEL_GROUP
    _PIPELINE_MODEL_PARALLEL_GROUP = None
    global _DATA_PARALLEL_GROUP
    _DATA_PARALLEL_GROUP = None
    global _EMBEDDING_GROUP
    _EMBEDDING_GROUP = None
    global _POSITION_EMBEDDING_GROUP
    _POSITION_EMBEDDING_GROUP = None

megatron/mpu/initialize.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# Parts of the code here are adapted from PyTorch
# repo: https://github.com/pytorch/pytorch


import numpy as np
import hashlib

import torch
import torch.nn.functional as F
from torch.nn.parameter import Parameter

from megatron import get_args, get_global_memory_buffer

from .initialize import (
    get_tensor_model_parallel_group,
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
)
from .mappings import (
    reduce_backward_from_tensor_model_parallel_region,
    gather_from_tensor_model_parallel_region,
    reduce_forward_from_tensor_model_parallel_region,
    reduce_scatter_to_sequence_parallel_region,
    scatter_to_tensor_model_parallel_region,
)
from .utils import VocabUtility, divide

# tensor_model_parallel_unique: True if the parameter weights are unique on different tensor parallel shards.
#  Examples of the True case would be dense layers that are sharded across tensor parallel ranks.
#  Examples of the False case would be LayerNorm layers that are just replicated across tensor parallel ranks.
# reduce_tensor_model_parallel_gradients: True if the gradients should be reduced across tensor parallel ranks.
#  This happens in optimizer.py.
#  Examples of the True case would be shared QK LayerNorms that are replicated across tensor parallel ranks, but
#  operate on different activations and need to be kept in sync.
_MODEL_PARALLEL_ATTRIBUTE_DEFAULTS = {
    "tensor_model_parallel_unique": False,
    "reduce_tensor_model_parallel_gradients": False,
}


def param_is_tensor_parallel_unique(param):
    """Return true if the parameter is sharded across tensor parallel ranks or is on rank 0.

    Typically used to determine if a given parameter's weights should count toward something like
    the total params norm. The goal is to avoid double counting parameters that have exact copies
    on multiple tensor parallel ranks.
    """
    return getattr(param, "tensor_model_parallel_unique", False) or get_tensor_model_parallel_rank() == 0


def param_is_tensor_parallel_replicated(param):
    """Retruns true if the parameter should be the same across all tensor parallel ranks."""
    return not getattr(param, "tensor_model_parallel_unique", False)


def set_tensor_model_parallel_attributes(tensor, *, is_tensor_parallel_unique, reduce_tensor_parallel_grads):
    if reduce_tensor_parallel_grads and is_tensor_parallel_unique:
        raise ValueError(f"Cannot set both is_tensor_parallel_unique and reduce_tensor_parallel_grads to True")
    # Make sure the attributes are not set.
    for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
        assert not hasattr(tensor, attribute)
    # Set the attributes.
    setattr(tensor, "tensor_model_parallel_unique", is_tensor_parallel_unique)
    setattr(tensor, "reduce_tensor_model_parallel_gradients", reduce_tensor_parallel_grads)
    # Make sure the attributes are all set.
    for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
        assert hasattr(tensor, attribute)


def set_defaults_if_not_set_tensor_model_parallel_attributes(tensor):
    def maybe_set(attribute, value):
        if not hasattr(tensor, attribute):
            setattr(tensor, attribute, value)

    for attribute, default in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS.items():
        maybe_set(attribute, default)


def copy_tensor_model_parallel_attributes(destination_tensor, source_tensor):
    def maybe_copy(attribute):
        if hasattr(source_tensor, attribute):
            setattr(destination_tensor, attribute, getattr(source_tensor, attribute))

    for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
        maybe_copy(attribute)


def _get_parameter_placement() -> torch.device:
    """Figure out where to place a parameter."""
    args = get_args()
    if args.place_on_cpu:
        return torch.device("cpu")
    else:
        return torch.cuda.current_device()


USED_RNG_KEYS = set()


def get_rng_for_key(rng_key: str) -> np.random.Generator:
    """Get a numpy random number generator for a given key."""
    if rng_key in USED_RNG_KEYS:
        raise ValueError(f"RNG key {rng_key} already used")
    USED_RNG_KEYS.add(rng_key)

    args = get_args()
    msg = hashlib.sha256(str(args.seed).encode("utf-8"))
    msg.update(rng_key.encode("utf-8"))
    seed = list(msg.digest())

    return np.random.default_rng(seed=seed)


def _initialize_affine_weight(
    weight,
    output_size,
    input_size,
    per_partition_size,
    partition_dim,
    init_method,
    rng_key: str,
    perform_initialization: bool,
    return_master_weight=False,
):
    """Initialize affine weight for model parallel.

    Build the master weight on CPU using numpy. This saves GPU memory and is fully deterministic.
    Then select the relevant chunk (for tensor parallelism) and move to GPU.

    Arguments:
        weight: weight tensor to initialize.
        output_size: first dimension of weight tensor.
        input_size: second dimension of weight tensor.
        per_partition_size: size of the partitioned dimension.
        partition_dim: dimension to partition along.
        init_method: weight initialization function.
        rng_key: a unique string for the weight to be initialized. Should be the same regardless of tensor or
            pipeline parallelism.
        perform_initialization: if False, skip initialization.
        return_master_weight: if True, return the master weight.
    """

    set_tensor_model_parallel_attributes(
        tensor=weight, is_tensor_parallel_unique=True, reduce_tensor_parallel_grads=False
    )

    if not perform_initialization:
        return None

    # Initialize master weight on CPU to save GPU memory and use fully deterministic numpy.
    rng = get_rng_for_key(rng_key)
    master_weight = torch.tensor(
        init_method(rng=rng, shape=[output_size, input_size]), dtype=torch.float, requires_grad=False
    )
    args = get_args()

    # Select the relevant chunk and move to GPU.
    weight_list = torch.split(master_weight, per_partition_size, dim=partition_dim)
    rank = get_tensor_model_parallel_rank()
    world_size = get_tensor_model_parallel_world_size()
    my_weight_list = weight_list[rank::world_size]

    with torch.no_grad():
        combined_weight = torch.cat(my_weight_list, dim=partition_dim)
        combined_weight = combined_weight.to(dtype=args.params_dtype)
        weight.copy_(combined_weight)
    if return_master_weight:
        return master_weight
    return None


class VocabParallelEmbedding(torch.nn.Module):
    """Embedding parallelized in the vocabulary dimension.

    This is mainly adapted from torch.nn.Embedding and all the default
    values are kept.
    Arguments:
        num_embeddings: vocabulary size.
        embedding_dim: size of hidden state.
        init_method: method to initialize weights.
    """

    def __init__(self, num_embeddings, embedding_dim, init_method, init_rng_key: str):
        super(VocabParallelEmbedding, self).__init__()
        # Keep the input dimensions.
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        # Set the defaults for compatibility.
        self.padding_idx = None
        self.max_norm = None
        self.norm_type = 2.0
        self.scale_grad_by_freq = False
        self.sparse = False
        self._weight = None
        self.tensor_model_parallel_size = get_tensor_model_parallel_world_size()
        # Divide the weight matrix along the vocabulary dimension.
        self.vocab_start_index, self.vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
            self.num_embeddings, get_tensor_model_parallel_rank(), self.tensor_model_parallel_size
        )
        self.num_embeddings_per_partition = self.vocab_end_index - self.vocab_start_index

        # Allocate weights and initialize.
        args = get_args()
        self.weight = Parameter(
            torch.empty(
                self.num_embeddings_per_partition,
                self.embedding_dim,
                device=_get_parameter_placement(),
                dtype=args.params_dtype,
            )
        )
        _initialize_affine_weight(
            weight=self.weight,
            output_size=self.num_embeddings,
            input_size=self.embedding_dim,
            per_partition_size=self.num_embeddings_per_partition,
            partition_dim=0,
            init_method=init_method,
            rng_key=init_rng_key,
            perform_initialization=args.perform_initialization,
        )

    def forward(self, input_):
        if self.tensor_model_parallel_size > 1:
            # Build the mask.
            input_mask = (input_ < self.vocab_start_index) | (input_ >= self.vocab_end_index)
            # Mask the input.
            masked_input = input_.clone() - self.vocab_start_index
            masked_input[input_mask] = 0
        else:
            masked_input = input_
            # Get the embeddings.
        output_parallel = F.embedding(
            masked_input,
            self.weight,
            self.padding_idx,
            self.max_norm,
            self.norm_type,
            self.scale_grad_by_freq,
            self.sparse,
        )
        # Mask the output embedding.
        if self.tensor_model_parallel_size > 1:
            output_parallel[input_mask, :] = 0.0
        # Reduce across all the model parallel GPUs.
        output = reduce_forward_from_tensor_model_parallel_region(output_parallel)
        return output


class LinearWithGradAccumulationAndAsyncCommunication(torch.autograd.Function):
    """
    Linear layer execution with asynchronous communication and gradient accumulation
    fusion in backprop.
    """

    @staticmethod
    def forward(ctx, input, weight, bias, gradient_accumulation_fusion, async_grad_allreduce, sequence_parallel):
        ctx.save_for_backward(input, weight)
        ctx.use_bias = bias is not None
        ctx.gradient_accumulation_fusion = gradient_accumulation_fusion
        ctx.async_grad_allreduce = async_grad_allreduce
        ctx.sequence_parallel = sequence_parallel

        if sequence_parallel:
            world_size = get_tensor_model_parallel_world_size()
            dim_size = list(input.size())
            dim_size[0] = dim_size[0] * world_size

            all_gather_buffer = get_global_memory_buffer().get_tensor(dim_size, input.dtype, "mpu")
            torch.distributed._all_gather_base(all_gather_buffer, input, group=get_tensor_model_parallel_group())
            total_input = all_gather_buffer
        else:
            total_input = input

        output = torch.matmul(total_input, weight.t())
        if bias is not None:
            output = output + bias
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, weight = ctx.saved_tensors
        use_bias = ctx.use_bias

        if ctx.sequence_parallel:
            world_size = get_tensor_model_parallel_world_size()
            dim_size = list(input.size())
            dim_size[0] = dim_size[0] * world_size

            all_gather_buffer = get_global_memory_buffer().get_tensor(dim_size, input.dtype, "mpu")
            handle = torch.distributed._all_gather_base(
                all_gather_buffer, input, group=get_tensor_model_parallel_group(), async_op=True
            )

            # Delay the start of intput gradient computation shortly (3us) to have
            # gather scheduled first and have GPU resources allocated
            _ = torch.empty(1, device=grad_output.device) + 1
            total_input = all_gather_buffer
        else:
            total_input = input
        grad_input = grad_output.matmul(weight)

        if ctx.sequence_parallel:
            handle.wait()

        # Convert the tensor shapes to 2D for execution compatibility
        # Necessary when using cross-attention and the flash attention module.
        grad_output = grad_output.contiguous()
        grad_output = grad_output.view(grad_output.shape[0] * grad_output.shape[1], grad_output.shape[2])
        total_input = total_input.view(total_input.shape[0] * total_input.shape[1], total_input.shape[2])

        if ctx.async_grad_allreduce:
            # Asynchronous all-reduce
            handle = torch.distributed.all_reduce(grad_input, group=get_tensor_model_parallel_group(), async_op=True)
            # Delay the start of weight gradient computation shortly (3us) to have
            # all-reduce scheduled first and have GPU resources allocated
            _ = torch.empty(1, device=grad_output.device) + 1

        if ctx.sequence_parallel:
            assert not ctx.async_grad_allreduce
            dim_size = list(input.size())
            sub_grad_input = torch.empty(
                dim_size, dtype=input.dtype, device=torch.cuda.current_device(), requires_grad=False
            )
            # reduce_scatter
            handle = torch.distributed._reduce_scatter_base(
                sub_grad_input, grad_input, group=get_tensor_model_parallel_group(), async_op=True
            )
            # Delay the start of weight gradient computation shortly (3us) to have
            # reduce scatter scheduled first and have GPU resources allocated
            _ = torch.empty(1, device=grad_output.device) + 1

        if ctx.gradient_accumulation_fusion:
            from megatron_fused_kernels import fused_dense_cuda

            fused_dense_cuda.wgrad_gemm_accum_fp32(total_input, grad_output, weight.main_grad)
            grad_weight = None
        else:
            grad_weight = grad_output.t().matmul(total_input)
        grad_bias = grad_output.sum(dim=0) if use_bias else None

        if ctx.sequence_parallel:
            handle.wait()
            return sub_grad_input, grad_weight, grad_bias, None, None, None

        if ctx.async_grad_allreduce:
            handle.wait()

        return grad_input, grad_weight, grad_bias, None, None, None


class ColumnParallelLinear(torch.nn.Module):
    """Linear layer with column parallelism.

    The linear layer is defined as Y = XA + b. A is parallelized along
    its second dimension as A = [A_1, ..., A_p].

    Arguments:
        input_size: first dimension of matrix A.
        output_size: second dimension of matrix A.
        bias: If true, add bias
        gather_output: If true, call all-gather on output and make Y available
                       to all GPUs, otherwise, every GPU will have its output
                       which is Y_i = XA_i
        init_method: method to initialize weights. Note that bias is always set
                     to zero.
        keep_master_weight_for_test: This was added for testing and should be
                                     set to False. It returns the master weights
                                     used for initialization.
        skip_bias_add: This was added to enable performance optimations where bias
                       can be fused with other elementwise operations. we skip
                       adding bias but instead return it.
    """

    def __init__(
        self,
        input_size,
        output_size,
        *,
        init_method,
        init_rng_key: str,
        sequence_parallel: bool,
        bias=True,
        gather_output=True,
        keep_master_weight_for_test=False,
        skip_bias_add=False,
    ):
        super(ColumnParallelLinear, self).__init__()

        # Keep input parameters
        self.input_size = input_size
        self.output_size = output_size
        self.gather_output = gather_output
        # Divide the weight matrix along the last dimension.
        world_size = get_tensor_model_parallel_world_size()
        self.output_size_per_partition = divide(output_size, world_size)
        self.skip_bias_add = skip_bias_add

        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        # Initialize weight.
        args = get_args()
        self.weight = Parameter(
            torch.empty(
                self.output_size_per_partition,
                self.input_size,
                device=_get_parameter_placement(),
                dtype=args.params_dtype,
            )
        )
        _initialize_affine_weight(
            weight=self.weight,
            output_size=self.output_size,
            input_size=self.input_size,
            per_partition_size=self.output_size_per_partition,
            partition_dim=0,
            init_method=init_method,
            rng_key=init_rng_key,
            perform_initialization=args.perform_initialization,
            return_master_weight=keep_master_weight_for_test,
        )
        if bias:
            self.bias = Parameter(
                torch.empty(self.output_size_per_partition, device=_get_parameter_placement(), dtype=args.params_dtype)
            )
            set_tensor_model_parallel_attributes(
                self.bias, is_tensor_parallel_unique=True, reduce_tensor_parallel_grads=False
            )
            # Always initialize bias to zero.
            with torch.no_grad():
                self.bias.zero_()
        else:
            self.register_parameter("bias", None)
        self.async_tensor_model_parallel_allreduce = args.async_tensor_model_parallel_allreduce and world_size > 1
        self.sequence_parallel = sequence_parallel
        assert not self.async_tensor_model_parallel_allreduce or not self.sequence_parallel
        self.gradient_accumulation_fusion = args.gradient_accumulation_fusion

    def forward(self, input_):
        bias = self.bias if not self.skip_bias_add else None

        if self.async_tensor_model_parallel_allreduce or self.sequence_parallel:
            input_parallel = input_
        else:
            input_parallel = reduce_backward_from_tensor_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = LinearWithGradAccumulationAndAsyncCommunication.apply(
            input_parallel,
            self.weight,
            bias,
            self.gradient_accumulation_fusion,
            self.async_tensor_model_parallel_allreduce,
            self.sequence_parallel,
        )
        if self.gather_output:
            assert not self.sequence_parallel
            output = gather_from_tensor_model_parallel_region(output_parallel)
        else:
            output = output_parallel
        output_bias = self.bias if self.skip_bias_add else None
        return output, output_bias


class ParallelLinear(ColumnParallelLinear):
    """Drop-in replacement for torch.nn.Linear, with tensor parallelism support."""

    def __init__(
        self,
        input_size,
        output_size,
        *,
        init_method,
        init_rng_key: str,
        sequence_parallel: bool,
        bias=True,
    ):
        super(ParallelLinear, self).__init__(
            input_size,
            output_size,
            bias=bias,
            init_method=init_method,
            init_rng_key=init_rng_key,
            gather_output=True,
            sequence_parallel=sequence_parallel,
        )

    def forward(self, input_):
        # normal linear layer doesn't return tuples, make this behave the same way
        output, _ = super().forward(input_)
        return output


class RowParallelLinear(torch.nn.Module):
    """Linear layer with row parallelism.

    The linear layer is defined as Y = XA + b. A is parallelized along
    its first dimension and X along its second dimension as:
               -   -
              | A_1 |
              | .   |
          A = | .   |        X = [X_1, ..., X_p]
              | .   |
              | A_p |
               -   -
    Arguments:
        input_size: first dimension of matrix A.
        output_size: second dimension of matrix A.
        bias: If true, add bias. Note that bias is not parallelized.
        input_is_parallel: If true, we assume that the input is already
                           split across the GPUs and we do not split
                           again.
        init_method: method to initialize weights. Note that bias is always set
                     to zero.
        keep_master_weight_for_test: This was added for testing and should be
                                     set to False. It returns the master weights
                                     used for initialization.
        skip_bias_add: This was added to enable performance optimization where bias
                       can be fused with other elementwise operations. We skip
                       adding bias but instead return it.
    """

    def __init__(
        self,
        input_size,
        output_size,
        *,
        init_method,
        init_rng_key: str,
        sequence_parallel: bool,
        bias=True,
        input_is_parallel=False,
        keep_master_weight_for_test=False,
        skip_bias_add=False,
    ):
        super(RowParallelLinear, self).__init__()

        # Keep input parameters
        self.input_size = input_size
        self.output_size = output_size
        self.input_is_parallel = input_is_parallel
        # Divide the weight matrix along the last dimension.
        world_size = get_tensor_model_parallel_world_size()
        self.input_size_per_partition = divide(input_size, world_size)
        self.skip_bias_add = skip_bias_add

        # Parameters.
        # Note: torch.nn.functional.linear performs XA^T + b and as a result
        # we allocate the transpose.
        # Initialize weight.
        args = get_args()
        self.weight = Parameter(
            torch.empty(
                self.output_size,
                self.input_size_per_partition,
                device=_get_parameter_placement(),
                dtype=args.params_dtype,
            )
        )
        _initialize_affine_weight(
            weight=self.weight,
            output_size=self.output_size,
            input_size=self.input_size,
            per_partition_size=self.input_size_per_partition,
            partition_dim=1,
            init_method=init_method,
            rng_key=init_rng_key,
            perform_initialization=args.perform_initialization,
            return_master_weight=keep_master_weight_for_test,
        )
        if bias:
            self.bias = Parameter(
                torch.empty(self.output_size, device=_get_parameter_placement(), dtype=args.params_dtype)
            )
            setattr(self.bias, "sequence_parallel", sequence_parallel)

            # Always initialize bias to zero.
            with torch.no_grad():
                self.bias.zero_()
        else:
            self.register_parameter("bias", None)
        self.sequence_parallel = sequence_parallel
        self.gradient_accumulation_fusion = args.gradient_accumulation_fusion

    def forward(self, input_):
        # Set up backprop all-reduce.
        if self.input_is_parallel:
            input_parallel = input_
        else:
            assert not self.sequence_parallel
            input_parallel = scatter_to_tensor_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = LinearWithGradAccumulationAndAsyncCommunication.apply(
            input_parallel, self.weight, None, self.gradient_accumulation_fusion, None, None
        )
        # All-reduce across all the partitions.
        if self.sequence_parallel:
            output_ = reduce_scatter_to_sequence_parallel_region(output_parallel)
        else:
            output_ = reduce_forward_from_tensor_model_parallel_region(output_parallel)
        if not self.skip_bias_add:
            output = output_ + self.bias if self.bias is not None else output_
            output_bias = None
        else:
            output = output_
            output_bias = self.bias
        return output, output_bias


def argmax_parallel(input_: torch.Tensor) -> torch.Tensor:
    """Argmax that supports tensor parallel inputs. Assumes dim=-1. Reverts to local argmax when not using parallelism."""
    local_max, local_index = torch.max(input_, dim=-1, keepdim=True)
    local_index += input_.shape[-1] * get_tensor_model_parallel_rank()

    gathered_max = gather_from_tensor_model_parallel_region(local_max)
    gathered_index = gather_from_tensor_model_parallel_region(local_index)

    global_argmax = torch.take_along_dim(
        gathered_index, torch.argmax(gathered_max, dim=-1, keepdim=True), dim=-1
    ).squeeze(-1)
    assert global_argmax.shape == input_.shape[:-1]
    return global_argmax

megatron/mpu/layers.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from .initialize import (
    get_tensor_model_parallel_group,
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
)
from .utils import split_tensor_along_last_dim


def _reduce(input_):
    """All-reduce the input tensor across model parallel group."""

    # Bypass the function if we are using only 1 GPU.
    if get_tensor_model_parallel_world_size() == 1:
        return input_

    # All-reduce, requires contiguous tensor.
    input_ = input_.contiguous()
    torch.distributed.all_reduce(input_, group=get_tensor_model_parallel_group())

    return input_


def _split_along_last_dim(input_):
    """Split the tensor along its last dimension and keep the
    corresponding slice."""

    world_size = get_tensor_model_parallel_world_size()
    # Bypass the function if we are using only 1 GPU.
    if world_size == 1:
        return input_

    # Split along last dimension.
    input_list = split_tensor_along_last_dim(input_, world_size)

    # Note: torch.split does not create contiguous tensors by default.
    rank = get_tensor_model_parallel_rank()
    output = input_list[rank].contiguous()

    return output


def _split_along_first_dim(input_):
    """Split the tensor along its first dimension and keep the
    corresponding slice."""

    world_size = get_tensor_model_parallel_world_size()
    # Bypass the function if we are using only 1 GPU.
    if world_size == 1:
        return input_

    # Split along first dimension.
    dim_size = input_.size()[0]
    assert (
        dim_size % world_size == 0
    ), f"First dimension of the tensor should be divisible by tensor parallel size. {dim_size} % {world_size} != 0"
    local_dim_size = dim_size // world_size
    rank = get_tensor_model_parallel_rank()
    dim_offset = rank * local_dim_size

    output = input_[dim_offset : dim_offset + local_dim_size].contiguous()

    return output


def _gather_along_last_dim(input_):
    """Gather tensors and concatinate along the last dimension."""

    world_size = get_tensor_model_parallel_world_size()
    # Bypass the function if we are using only 1 GPU.
    if world_size == 1:
        return input_

    # Size and dimension.
    last_dim = input_.dim() - 1
    rank = get_tensor_model_parallel_rank()

    tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
    tensor_list[rank] = input_
    torch.distributed.all_gather(tensor_list, input_, group=get_tensor_model_parallel_group())

    # Note: torch.cat already creates a contiguous tensor.
    output = torch.cat(tensor_list, dim=last_dim).contiguous()

    return output


def _gather_along_first_dim(input_):
    """Gather tensors and concatinate along the first dimension."""

    world_size = get_tensor_model_parallel_world_size()
    # Bypass the function if we are using only 1 GPU.
    if world_size == 1:
        return input_

    dim_size = list(input_.size())
    dim_size[0] = dim_size[0] * world_size

    output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
    torch.distributed._all_gather_base(output, input_.contiguous(), group=get_tensor_model_parallel_group())

    return output


def _reduce_scatter_along_first_dim(input_):
    """Reduce-scatter the input tensor across model parallel group."""
    world_size = get_tensor_model_parallel_world_size()
    # Bypass the function if we are using only 1 GPU.
    if world_size == 1:
        return input_

    dim_size = list(input_.size())
    assert dim_size[0] % world_size == 0, "First dimension of the tensor should be divisible by tensor parallel size"

    dim_size[0] = dim_size[0] // world_size

    output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
    torch.distributed._reduce_scatter_base(output, input_.contiguous(), group=get_tensor_model_parallel_group())
    return output


class _ReduceBackwardFromTensorModelParallelRegion(torch.autograd.Function):
    """All-reduce the backward pass input from the model parallel region."""

    @staticmethod
    def symbolic(graph, input_):
        return input_

    @staticmethod
    def forward(ctx, input_):
        return input_

    @staticmethod
    def backward(ctx, grad_output):
        return _reduce(grad_output)


class _ReduceForwardFromTensorModelParallelRegion(torch.autograd.Function):
    """All-reduce the input from the model parallel region."""

    @staticmethod
    def symbolic(graph, input_):
        return _reduce(input_)

    @staticmethod
    def forward(ctx, input_):
        return _reduce(input_)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output


class _ScatterToModelParallelRegion(torch.autograd.Function):
    """Split the input and keep only the corresponding chuck to the rank."""

    @staticmethod
    def symbolic(graph, input_):
        return _split_along_last_dim(input_)

    @staticmethod
    def forward(ctx, input_):
        return _split_along_last_dim(input_)

    @staticmethod
    def backward(ctx, grad_output):
        return _gather_along_last_dim(grad_output)


class _GatherFromModelParallelRegion(torch.autograd.Function):
    """Gather the input from model parallel region and concatinate."""

    @staticmethod
    def symbolic(graph, input_):
        return _gather_along_last_dim(input_)

    @staticmethod
    def forward(ctx, input_):
        return _gather_along_last_dim(input_)

    @staticmethod
    def backward(ctx, grad_output):
        return _split_along_last_dim(grad_output)


class _ScatterToSequenceParallelRegion(torch.autograd.Function):
    """Split the input and keep only the corresponding chuck to the rank."""

    @staticmethod
    def symbolic(graph, input_):
        return _split_along_first_dim(input_)

    @staticmethod
    def forward(ctx, input_):
        return _split_along_first_dim(input_)

    @staticmethod
    def backward(ctx, grad_output):
        return _gather_along_first_dim(grad_output)


class _GatherFromSequenceParallelRegion(torch.autograd.Function):
    """Gather the input from sequence parallel region and concatinate."""

    @staticmethod
    def symbolic(graph, input_, tensor_parallel_output_grad=True):
        return _gather_along_first_dim(input_)

    @staticmethod
    def forward(ctx, input_, tensor_parallel_output_grad=True):
        ctx.tensor_parallel_output_grad = tensor_parallel_output_grad
        return _gather_along_first_dim(input_)

    @staticmethod
    def backward(ctx, grad_output):
        tensor_parallel_output_grad = ctx.tensor_parallel_output_grad

        # If the computation graph after the gather operation is
        # in the tensor parallel mode, output gradients need to reduce
        # scattered and whereas if the computation is duplicated,
        # output gradients need to be scattered.
        if tensor_parallel_output_grad:
            return _reduce_scatter_along_first_dim(grad_output), None
        else:
            return _split_along_first_dim(grad_output), None


class _ReduceScatterToSequenceParallelRegion(torch.autograd.Function):
    """Reduce scatter the input from the model parallel region."""

    @staticmethod
    def symbolic(graph, input_):
        return _reduce_scatter_along_first_dim(input_)

    @staticmethod
    def forward(ctx, input_):
        return _reduce_scatter_along_first_dim(input_)

    @staticmethod
    def backward(ctx, grad_output):
        return _gather_along_first_dim(grad_output)


# -----------------
# Helper functions.
# -----------------


def reduce_backward_from_tensor_model_parallel_region(input_):
    return _ReduceBackwardFromTensorModelParallelRegion.apply(input_)


def reduce_forward_from_tensor_model_parallel_region(input_):
    return _ReduceForwardFromTensorModelParallelRegion.apply(input_)


def scatter_to_tensor_model_parallel_region(input_):
    return _ScatterToModelParallelRegion.apply(input_)


def gather_from_tensor_model_parallel_region(input_):
    return _GatherFromModelParallelRegion.apply(input_)


def scatter_to_sequence_parallel_region(input_):
    return _ScatterToSequenceParallelRegion.apply(input_)


def gather_from_sequence_parallel_region(input_, tensor_parallel_output_grad=True):
    return _GatherFromSequenceParallelRegion.apply(input_, tensor_parallel_output_grad)


def reduce_scatter_to_sequence_parallel_region(input_):
    return _ReduceScatterToSequenceParallelRegion.apply(input_)

megatron/mpu/mappings.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# Parts of the code here are adapted from PyTorch
# repo: https://github.com/pytorch/pytorch

import contextlib

import torch
from torch import _C
from torch.cuda import _lazy_call, device as device_ctx_manager
from torch.utils.checkpoint import detach_variable

from megatron.memory import allocate_mem_buff

from .initialize import get_data_parallel_rank
from .initialize import get_tensor_model_parallel_group
from .initialize import get_tensor_model_parallel_rank
from .initialize import get_tensor_model_parallel_world_size


# Default name for the model parallel rng tracker.
_MODEL_PARALLEL_RNG_TRACKER_NAME = "model-parallel-rng"


def _set_cuda_rng_state(new_state, device=-1):
    """Sets the random number generator state of the current GPU.

    Argumentss:
        new_state (torch.ByteTensor): The desired state
    This function is adapted from PyTorch repo (torch.cuda.set_rng_state)
    with a single change: the input state is not cloned. Cloning caused
    major performance issues for +4 GPU cases.
    """
    if hasattr(_C, "_cuda_setRNGState") and callable(_C._cuda_setRNGState):
        # older PyTorch
        def cb():
            with device_ctx_manager(device):
                _C._cuda_setRNGState(new_state)

    else:
        # newer PyTorch
        if device == -1:
            device = torch.device("cuda")
        elif isinstance(device, str):
            device = torch.device(device)
        elif isinstance(device, int):
            device = torch.device("cuda", device)

        def cb():
            idx = device.index
            if idx is None:
                idx = torch.cuda.current_device()
            default_generator = torch.cuda.default_generators[idx]
            default_generator.set_state(new_state)

    _lazy_call(cb)


def split_tensor_into_1d_equal_chunks(tensor, new_buffer=False):
    """Break a tensor into equal 1D chunks."""
    partition_size = torch.numel(tensor) // get_tensor_model_parallel_world_size()
    start_index = partition_size * get_tensor_model_parallel_rank()
    end_index = start_index + partition_size
    if new_buffer:
        data = torch.empty(partition_size, dtype=tensor.dtype, device=torch.cuda.current_device(), requires_grad=False)
        data.copy_(tensor.view(-1)[start_index:end_index])
    else:
        data = tensor.view(-1)[start_index:end_index]
    return data


def gather_split_1d_tensor(tensor):
    """Opposite of above function, gather values from model parallel ranks."""
    numel_gathered = torch.numel(tensor) * get_tensor_model_parallel_world_size()
    gathered = torch.empty(numel_gathered, dtype=tensor.dtype, device=torch.cuda.current_device(), requires_grad=False)
    # TODO: This API is experimental in pytorch (as of Feb 2022) and
    # this might break in future pytorch releases. We chose this API
    # as opposed to torch.distributed.all_gather for efficiency reasons.
    # This API calls directly NCCL all-gather versus the former does
    # internal copies and can potentially cause slow down.
    torch.distributed._all_gather_base(gathered, tensor, group=get_tensor_model_parallel_group())
    return gathered


def _kernel_make_viewless_tensor(inp, requires_grad):
    """Make a viewless tensor.

    View tensors have the undesirable side-affect of retaining a reference
    to the originally-viewed tensor, even after manually setting the '.data'
    field. This method creates a new tensor that links to the old tensor's
    data, without linking the viewed tensor, referenced via the '._base'
    field.
    """
    out = torch.empty(
        (1,),
        dtype=inp.dtype,
        device=inp.device,
        requires_grad=requires_grad,
    )
    out.data = inp.data
    return out


class MakeViewlessTensor(torch.autograd.Function):
    """
    Autograd function to make a viewless tensor.

    This function should be used in cases where the computation graph needs
    to be propagated, but we only want a viewless tensor (e.g.,
    ParallelTransformer's hidden_states). Call this function by passing
    'keep_graph = True' to 'make_viewless_tensor()'.
    """

    @staticmethod
    def forward(ctx, inp, requires_grad):
        return _kernel_make_viewless_tensor(inp, requires_grad)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, None


def make_viewless_tensor(inp, requires_grad, keep_graph):
    """
    Entry-point for creating viewless tensors.

    This method should be used, rather than calling 'MakeViewlessTensor'
    or '_kernel_make_viewless_tensor' directly. This method acts as a
    switch for determining if an autograd function or a regular method
    should be used to create the tensor.
    """

    # return tensor as-is, if not a 'view'
    if inp._base is None:
        return inp

    # create viewless tensor
    if keep_graph:
        return MakeViewlessTensor.apply(inp, requires_grad)
    else:
        return _kernel_make_viewless_tensor(inp, requires_grad)


def assert_viewless_tensor(tensor, extra_msg=None):
    """Assert that a tensor is not a view (i.e., its '._base' field is
    not set)."""
    if isinstance(tensor, list):
        [assert_viewless_tensor(t) for t in tensor]
        return tensor
    if not isinstance(tensor, torch.Tensor):
        return tensor
    assert tensor._base is None, (
        "Ensure tensor._base is None before setting tensor.data or storing "
        "tensor to memory buffer. Otherwise, a memory leak will occur (and "
        "likely accumulate over iterations). %s"
    ) % extra_msg
    return tensor


def safely_set_viewless_tensor_data(tensor, new_data_tensor):
    """Safely set tensor's '.data' field.

    Check first that the tensor is viewless (i.e., '._base' not set). If not,
    raise an exception.
    """
    assert_viewless_tensor(
        tensor,
        extra_msg="FYI, tensor._base has shape %s, and new_data_tensor has shape %s."
        % ("--" if tensor._base is None else tensor._base.shape, new_data_tensor.shape),
    )
    tensor.data = new_data_tensor


class CudaRNGStatesTracker:
    """Tracker for the cuda RNG states.

    Using the `add` method, a cuda rng state is initialized based on
    the input `seed` and is assigned to `name`. Later, by forking the
    rng state, we can perform operations and return to our starting
    cuda state.
    """

    def __init__(self):
        # Map from a string name to the cuda rng state.
        self.states_ = {}
        # Seeds are just for book keeping and ensure no seed is set twice.
        self.seeds_ = set()

    def reset(self):
        """Set to the initial state (no tracker)."""
        self.states_ = {}
        self.seeds_ = set()

    def get_states(self):
        """Get rng states. Copy the dictionary so we have direct
        pointers to the states, not just a pointer to the dictionary."""
        states = {}
        for name in self.states_:
            states[name] = self.states_[name]
        return states

    def set_states(self, states):
        """Set the rng states. For efficiency purposes, we do not check
        the size of seed for compatibility."""
        self.states_ = states

    def add(self, name, seed):
        """Track the rng state."""
        # Check seed is not already used.
        if seed in self.seeds_:
            raise Exception("seed {} already exists".format(seed))
        self.seeds_.add(seed)
        # Check that state is not already defined.
        if name in self.states_:
            raise Exception("cuda rng state {} already exists".format(name))
        # Get the current rng state.
        orig_rng_state = torch.cuda.get_rng_state()
        # Set the new state and store it.
        torch.cuda.manual_seed(seed)
        self.states_[name] = torch.cuda.get_rng_state()
        # Reset rng state to what it was.
        _set_cuda_rng_state(orig_rng_state)

    @contextlib.contextmanager
    def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME):
        """Fork the cuda rng state, perform operations, and exit with
        the original state."""
        # Check if we have added the state
        if name not in self.states_:
            raise Exception("cuda rng state {} is not added".format(name))
        # Store current rng state.
        orig_cuda_rng_state = torch.cuda.get_rng_state()
        # Set rng state to the desired one
        _set_cuda_rng_state(self.states_[name])
        # Do the stuff we wanted to do.
        try:
            yield
        finally:
            # Update the current rng state for later use.
            self.states_[name] = torch.cuda.get_rng_state()
            # And set the state to the original state we started with.
            _set_cuda_rng_state(orig_cuda_rng_state)


# RNG tracker object.
_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker()


def get_cuda_rng_tracker():
    """Get cuda rng tracker."""
    return _CUDA_RNG_STATE_TRACKER


def model_parallel_cuda_manual_seed(seed):
    """Initialize model parallel cuda seed.

    This function should be called after the model parallel is
    initialized. Also, no torch.cuda.manual_seed should be called
    after this function. Basically, this is replacement for that
    function.
    Two set of RNG states are tracked:
        default state: This is for data parallelism and is the same among a
                       set of model parallel GPUs but different across
                       different model paralle groups. This is used for
                       example for dropout in the non-tensor-model-parallel regions.
        tensor-model-parallel state: This state is different among a set of model
                              parallel GPUs, but the same across data parallel
                              groups. This is used for example for dropout in
                              model parallel regions.
    """
    # 2718 is just for fun and any POSITIVE value will work.
    offset = seed + 2718
    tensor_model_parallel_seed = offset + get_tensor_model_parallel_rank()
    # Data parallel gets the original seed.
    data_parallel_seed = seed

    if torch.distributed.get_rank() == 0:
        print(
            "> initializing model parallel cuda seeds on global rank {}, "
            "model parallel rank {}, and data parallel rank {} with "
            "model parallel seed: {} and data parallel seed: {}".format(
                torch.distributed.get_rank(),
                get_tensor_model_parallel_rank(),
                get_data_parallel_rank(),
                tensor_model_parallel_seed,
                data_parallel_seed,
            ),
            flush=True,
        )
    _CUDA_RNG_STATE_TRACKER.reset()
    # Set the default state.
    torch.cuda.manual_seed(data_parallel_seed)
    # and model parallel state.
    _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, tensor_model_parallel_seed)


class CheckpointFunction(torch.autograd.Function):
    """This function is adapted from torch.utils.checkpoint with
    two main changes:
        1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state`
        2) the states in the model parallel tracker are also properly
           tracked/set/reset.
    """

    @staticmethod
    def forward(ctx, run_function, distribute_saved_activations, *args):
        ctx.run_function = run_function
        ctx.distribute_saved_activations = distribute_saved_activations

        # Copy the rng states.
        ctx.fwd_cpu_rng_state = torch.get_rng_state()
        ctx.fwd_cuda_rng_state = torch.cuda.get_rng_state()
        ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()

        with torch.no_grad():
            outputs = run_function(*args)

        # Divide hidden states across model parallel group and only keep
        # the chunk corresponding to the current rank.
        if distribute_saved_activations:
            ctx.input_0_shape = args[0].data.shape
            safely_set_viewless_tensor_data(args[0], split_tensor_into_1d_equal_chunks(args[0].data, new_buffer=True))

        # Store everything.
        ctx.save_for_backward(*args)

        return outputs

    @staticmethod
    def backward(ctx, *args):
        if not torch.autograd._is_checkpoint_valid():
            raise RuntimeError("Checkpointing is not compatible with .grad(), " "please use .backward() if possible")
        inputs = ctx.saved_tensors
        if ctx.distribute_saved_activations:
            safely_set_viewless_tensor_data(inputs[0], gather_split_1d_tensor(inputs[0].data).view(ctx.input_0_shape))

        # Store the current states.
        bwd_cpu_rng_state = torch.get_rng_state()
        bwd_cuda_rng_state = torch.cuda.get_rng_state()
        bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()

        # Set the states to what it used to be before the forward pass.
        torch.set_rng_state(ctx.fwd_cpu_rng_state)
        _set_cuda_rng_state(ctx.fwd_cuda_rng_state)
        get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker)

        # Compute the forward pass.
        detached_inputs = detach_variable(inputs)
        with torch.enable_grad():
            outputs = ctx.run_function(*detached_inputs)

        # Set the states back to what it was at the start of this function.
        torch.set_rng_state(bwd_cpu_rng_state)
        _set_cuda_rng_state(bwd_cuda_rng_state)
        get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker)

        if isinstance(outputs, torch.Tensor):
            outputs = (outputs,)
        torch.autograd.backward(outputs, args)
        grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs)
        return (None, None) + grads


def checkpoint(function, distribute_saved_activations, *args):
    """Checkpoint a model or part of the model.
    This has been directly copied from torch.utils.checkpoint."""
    return CheckpointFunction.apply(function, distribute_saved_activations, *args)

megatron/mpu/random.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import torch


def ensure_divisibility(numerator, denominator):
    """Ensure that numerator is divisible by the denominator."""
    assert numerator % denominator == 0, "{} is not divisible by {}".format(numerator, denominator)


def divide(numerator, denominator):
    """Ensure that numerator is divisible by the denominator and return
    the division value."""
    ensure_divisibility(numerator, denominator)
    return numerator // denominator


def split_tensor_along_last_dim(tensor, num_partitions, contiguous_split_chunks=False):
    """Split a tensor along its last dimension.
    Arguments:
        tensor: input tensor.
        num_partitions: number of partitions to split the tensor
        contiguous_split_chunks: If True, make each chunk contiguous
                                 in memory.
    """
    # Get the size and dimension.
    last_dim = tensor.dim() - 1
    last_dim_size = divide(tensor.size()[last_dim], num_partitions)
    # Split.
    tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
    # Note: torch.split does not create contiguous tensors by default.
    if contiguous_split_chunks:
        return tuple(chunk.contiguous() for chunk in tensor_list)

    return tensor_list


class VocabUtility:
    """Split the vocabulary into `world_size` chunks amd return the
    first and last index of the vocabulary belonging to the `rank`
    partition: Note that indecies in [fist, last)"""

    @staticmethod
    def vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank, world_size):
        index_f = rank * per_partition_vocab_size
        index_l = index_f + per_partition_vocab_size
        return index_f, index_l

    @staticmethod
    def vocab_range_from_global_vocab_size(global_vocab_size, rank, world_size):
        per_partition_vocab_size = divide(global_vocab_size, world_size)
        return VocabUtility.vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank, world_size)

megatron/mpu/utils.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from .api import generate, generate_and_post_process

megatron/text_generation/__init__.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Inference API."""


import base64
from typing import List, Optional, Tuple, Any, Dict

import numpy as np
import torch
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP

from megatron.model.module import MegatronModule
from megatron import mpu
from megatron.model import DistributedDataParallel as LocalDDP
from megatron.model import Float16Module

from megatron.tokenizer.tokenizer import AbstractTokenizer
from megatron.mpu.communication import broadcast_float_list, broadcast_int_list
from megatron.text_generation.generation import (
    generate_tokens_probs_and_return_on_first_stage,
    score_and_return_on_first_stage,
)
from megatron.text_generation.inference_params import InferenceParams
from megatron.text_generation.tokenization import (
    convert_generations_to_human_readable_tokens,
    tokenize_prompts,
)
from megatron.utils import unwrap_model
import numpy.typing as npt

def preprocess_prompts(prompts: Optional[List[List[str]]]) -> Optional[List[List[str]]]:
    """
    Accepts a list of list of subprompts, returns a list of list of processed subprompts.
    """
    if prompts is None:
        return None
    processed_prompts = []
    for prompt in prompts:
        processed_subprompts = []
        for subprompt in prompt:
            processed_subprompts.append(f"human: {subprompt.strip()}\n\nadept:")
        processed_prompts.append(processed_subprompts)
    return processed_prompts


def generate_and_post_process(
    model: MegatronModule,
    params_dtype: torch.dtype,
    max_position_embeddings: int,
    termination_id: int,
    tokenizer: AbstractTokenizer,
    prompts: Optional[List[List[str]]] = None,
    max_tokens_to_generate: int = 0,
    inference_params: Optional[InferenceParams] = None,
    return_output_log_probs: bool = False,
    return_all_log_probs: bool = False,
    log_prob_tokens: Optional[torch.Tensor] = None,
    top_k_sampling: int = 0,
    top_p_sampling: float = 0.0,
    temperature: float = 1.0,
    add_BOS: bool = False,
    random_seed: int = -1,
    process_prompts_for_chat: bool = False
) -> Optional[
    Tuple[
        List[str],
        Optional[Any],
        List[List[int]],
        List[str],
        Optional[str],
        Optional[Any],
        List[List[str]],
        Optional[Any],
    ]
]:
    """Run inference and post-process outputs, i.e., detokenize,
    move to cpu and convert to list.

    prompts: a list of list of strings, where each element in the outer list represents a single sample that
    is an item in the batch. A single sample is represented as a list of strings.
    """
    # Pre-process the prompts.
    if process_prompts_for_chat:
        prompts = preprocess_prompts(prompts)

    # Main inference.
    outputs = generate(
        model,
        max_position_embeddings=max_position_embeddings,
        params_dtype=params_dtype,
        termination_id=termination_id,
        tokenizer=tokenizer,
        prompts=prompts,
        max_tokens_to_generate=max_tokens_to_generate,
        inference_params=inference_params,
        return_output_log_probs=return_output_log_probs,
        return_all_log_probs=return_all_log_probs,
        log_prob_tokens=log_prob_tokens,
        top_k_sampling=top_k_sampling,
        top_p_sampling=top_p_sampling,
        temperature=temperature,
        add_BOS=add_BOS,
        random_seed=random_seed,
    )

    # Only post-process on first stage.
    if mpu.is_pipeline_first_stage():

        all_tokens = outputs["tokens"].cpu().numpy().tolist()

        raw_prompts_plus_generations = [tokenizer.detokenize(ts) for ts in all_tokens]
        processed_generations = [
            tokenizer.detokenize(toks) for toks in outputs.get("generated_tokens", [])
        ]

        processed_generated_tokens = [
            tokenizer.encode(g) for g in processed_generations
        ]

        human_readable_tokens = convert_generations_to_human_readable_tokens(
            processed_generated_tokens, "<s>"
        )

        output_log_probs = None
        if return_output_log_probs:
            output_log_probs = outputs["output_log_probs"].cpu().numpy().tolist()

        all_log_probs = None
        if return_all_log_probs:
            all_log_probs = outputs["all_log_probs"].cpu().numpy().tolist()

        return (
            raw_prompts_plus_generations,
            output_log_probs,
            all_tokens,
            processed_generations,
            all_log_probs,
            human_readable_tokens,
        )

    return None


def generate(
    model: MegatronModule,
    max_position_embeddings: int,
    params_dtype: torch.dtype,
    termination_id: int,
    tokenizer: AbstractTokenizer,
    prompts: Optional[List[List[str]]] = None,
    max_tokens_to_generate: int = 0,
    inference_params: Optional[InferenceParams] = None,
    return_output_log_probs: bool = False,
    return_all_log_probs: bool = False,
    log_prob_tokens: Optional[torch.Tensor] = None,
    top_k_sampling: int = 0,
    top_p_sampling: float = 0.0,
    temperature: float = 1.0,
    add_BOS: bool = False,
    random_seed: int = -1,
) -> Dict[str, Any]:
    """Given prompts and input parameters, run inference and return:
    tokens: prompts plus the generated tokens.
    lengths: length of the prompt + generations. Note that we can
        discard tokens in the tokens tensor that are after the
        corresponding length.
    output_log_probs: log probs of the tokens.
    all_log_probs: log probs of all the vocab (or just log_prob_tokens if provided)
        tokens for each generated token position.
    """
    num_log_prob_tokens = 0 if log_prob_tokens is None else len(log_prob_tokens)

    # Make sure input params are avaialble to all ranks.
    values = [
        max_tokens_to_generate,
        return_output_log_probs,
        top_k_sampling,
        top_p_sampling,
        temperature,
        add_BOS,
        random_seed,
        return_all_log_probs,
        num_log_prob_tokens,
    ]
    values_float_tensor = broadcast_float_list(len(values), float_list=values)
    max_tokens_to_generate = int(values_float_tensor[0].item())
    return_output_log_probs = bool(values_float_tensor[1].item())
    top_k_sampling = int(values_float_tensor[2].item())
    top_p_sampling = values_float_tensor[3].item()
    temperature = values_float_tensor[4].item()
    add_BOS = bool(values_float_tensor[5].item())
    random_seed = int(values_float_tensor[6].item())
    return_all_log_probs = bool(values_float_tensor[7].item())
    num_log_prob_tokens = int(values_float_tensor[8].item())

    if return_all_log_probs and num_log_prob_tokens > 0:
        # Do another broadcast for the log_prob_tokens.
        log_prob_tokens = broadcast_int_list(
            num_log_prob_tokens, int_list=log_prob_tokens
        )
    else:
        log_prob_tokens = None

    if random_seed != -1:
        torch.random.manual_seed(random_seed)

    # Tokenize prompts and get the batch.
    # Note that these tensors are broadcasted to all ranks.
    if torch.distributed.get_rank() == 0:
        assert prompts is not None

    context_tokens_tensor, context_length_tensor = tokenize_prompts(
        prompts=prompts,
        max_tokens_to_generate=max_tokens_to_generate,
        max_position_embeddings=max_position_embeddings,
        add_BOS=add_BOS,
    )

    batch_size = context_tokens_tensor.shape[0]
    num_sub_sequences = context_tokens_tensor.shape[1]

    assert num_sub_sequences == 1
    # Remove subsequence dim
    context_length_tensor = torch.squeeze(context_length_tensor, dim=1)
    context_tokens_tensor = torch.squeeze(context_tokens_tensor, dim=1)

    if max_tokens_to_generate == 0:
        assert inference_params is not None
        return score_and_return_on_first_stage(
            model,
            context_tokens_tensor,
            context_length_tensor,
            inference_params,
            max_position_embeddings=max_position_embeddings,
        )

    # Main inference function.
    # Note that the outputs are available on the first stage.
    assert inference_params is not None
    # Added termination_id to support the case that we want to terminate the
    # generation once that id is generated.

    outputs = generate_tokens_probs_and_return_on_first_stage(
        model,
        context_tokens_tensor,
        context_length_tensor,
        inference_params=inference_params,
        max_position_embeddings=max_position_embeddings,
        termination_id=termination_id,
        vocab_size=tokenizer.vocab_size,
        return_output_log_probs=return_output_log_probs,
        log_prob_tokens=log_prob_tokens,
        return_all_log_probs=return_all_log_probs,
        top_k=top_k_sampling,
        top_p=top_p_sampling,
        temperature=temperature,
    )
    # Now we can figure out what actually got generated and return that too
    generated_tokens = []
    context_lengths = context_length_tensor.detach().cpu().numpy().tolist()
    contexts = context_tokens_tensor.detach().cpu().numpy().tolist()
    lengths_cpu = outputs["lengths"].detach().cpu().numpy().tolist()

    for b in range(batch_size):
        assert lengths_cpu[b] > context_lengths[b]
        gen = contexts[b][context_lengths[b] : lengths_cpu[b]]
        generated_tokens.append(gen)

    outputs["generated_tokens"] = generated_tokens

    return outputs

megatron/text_generation/api.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Forward step utilities."""

from typing import Optional, Dict, Any, List
from collections.abc import Iterable

from torch import dtype
import torch
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP

from megatron import get_args, mpu
from megatron.text_generation.inference_params import InferenceParams
from megatron.model import DistributedDataParallel as LocalDDP
from megatron.model import Float16Module
from megatron.model.module import MegatronModule
from megatron.model.gpt_model import GPTModel
from megatron.utils import unwrap_model

from megatron.mpu.communication import recv_from_prev_pipeline_rank_, send_to_next_pipeline_rank


class ForwardStep:
    """Forward step function with all the communications.
    We use a class here to hide the inference parameters
    from the outside caller."""

    def __init__(
        self,
        model: MegatronModule,
        max_batch_size: int,
        max_sequence_len: int,
        inference_params: InferenceParams,
    ):
        """Set values so we don't need to do it multiple times."""
        # Make sure model is in eval mode.
        assert not isinstance(model, Iterable), "interleaving schedule is not supported for inference"
        model.eval()
        self.model = model
        # Initialize inference parameters.
        self.inference_params = inference_params
        # Pipelining arguments.
        args = get_args()
        self.pipeline_size_larger_than_one = args.pipeline_model_parallel_size > 1
        # Threshold of pipelining.
        self.pipelining_batch_x_seqlen = args.inference_batch_times_seqlen_threshold

    def __call__(
        self,
        tokens: torch.Tensor,
        position_ids: torch.Tensor,
        lm_logits_mask: Optional[torch.Tensor] = None,
    ) -> Optional[Dict[str, Any]]:
        """Invocation of the forward methods. Note that self.inference_params
        is being modified by the forward step."""
        # Pipelining case.
        if self.pipeline_size_larger_than_one:
            current_batch_x_seqlen = tokens.size(0) * tokens.size(1)
            if current_batch_x_seqlen >= self.pipelining_batch_x_seqlen:
                raise ValueError("We deleted _with_pipelining_forward_step")

        return _no_pipelining_forward_step(
            self.model,
            tokens,
            position_ids,
            self.inference_params,
            lm_logits_mask=lm_logits_mask,
        )


def _get_recv_buffer_dtype(args: Any) -> dtype:
    """Receive happens between the layers."""
    if args.fp32_residual_connection:
        return torch.float
    return args.params_dtype


def _allocate_recv_buffer(batch_size: int, sequence_length: int) -> torch.Tensor:
    """Receive happens between the layers with size [s, b, h]."""
    if mpu.is_pipeline_first_stage():
        return None
    args = get_args()
    recv_size = (sequence_length, batch_size, args.hidden_size)
    return torch.empty(recv_size, dtype=_get_recv_buffer_dtype(args), device=torch.cuda.current_device())


def _model_forward_step(
    model: MegatronModule,
    tokens: torch.Tensor,
    position_ids: torch.Tensor,
    inference_params: InferenceParams,
    lm_logits_mask: Optional[torch.Tensor] = None,
) -> Dict[str, Any]:
    # Run a simple forward pass.
    unwrapped_model = unwrap_model(model, (torchDDP, LocalDDP, Float16Module))
    outputs: Dict[str, Any] = {}
    if isinstance(unwrapped_model, GPTModel):
        outputs = model(tokens, position_ids, lm_logits_mask=lm_logits_mask, inference_params=inference_params)
    else:
        assert False, "Unknown model type!" + str(type(unwrapped_model))

    if not isinstance(outputs, dict):
        outputs = {"logits": outputs}

    return outputs


def _forward_step_helper(
    model: MegatronModule,
    tokens: torch.Tensor,
    position_ids: torch.Tensor,
    inference_params: InferenceParams,
    recv_buffer: Optional[torch.Tensor] = None,
    lm_logits_mask: Optional[torch.Tensor] = None,
) -> Dict[str, Any]:
    """Single forward step. Update the allocate memory flag so
    only the first time the memory is allocated."""
    batch_size = tokens.size(0)
    sequence_length = tokens.size(1)
    if recv_buffer is None:
        recv_buffer = _allocate_recv_buffer(batch_size, sequence_length)

    # Receive from previous stage.
    recv_from_prev_pipeline_rank_(recv_buffer)

    # Forward pass through the model.
    model.set_input_tensor(recv_buffer)
    outputs = _model_forward_step(
        model,
        tokens,
        position_ids,
        inference_params=inference_params,
        lm_logits_mask=lm_logits_mask,
    )

    # Send output to the next stage.
    send_to_next_pipeline_rank(outputs)

    return outputs


def _no_pipelining_forward_step(
    model: MegatronModule,
    tokens: torch.Tensor,
    position_ids: torch.Tensor,
    inference_params: InferenceParams,
    recv_buffer: Optional[torch.Tensor] = None,
    lm_logits_mask: Optional[torch.Tensor] = None,
) -> Optional[Dict[str, Any]]:
    """If recv_buffer is none, we will allocate one on the fly."""
    # Run a simple forward pass.
    outputs = _forward_step_helper(
        model,
        tokens,
        position_ids,
        inference_params,
        recv_buffer=recv_buffer,
        lm_logits_mask=lm_logits_mask,
    )
    # Update the sequence length offset.
    if inference_params is not None:
        inference_params.sequence_len_offset += tokens.size(1)

    lm_output = None
    if mpu.is_pipeline_last_stage():
        lm_output = outputs

    return lm_output

megatron/text_generation/forward_step.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Generation utilities."""

from typing import Optional, Dict, Tuple, List

import torch
import torch.nn.functional as F

from megatron import mpu
from megatron.utils import get_ltor_masks_and_position_ids
from megatron.model.module import MegatronModule

from megatron.mpu.communication import (
    broadcast_from_last_pipeline_stage,
    broadcast_from_last_to_first_pipeline_stage,
    copy_from_last_to_first_pipeline_stage,
)
from .inference_params import InferenceParams
from .forward_step import ForwardStep
from .sampling import sample


def score_and_return_on_first_stage(
    model: MegatronModule,
    tokens: torch.Tensor,
    lengths: torch.Tensor,
    inference_params: InferenceParams,
    max_position_embeddings: int,
) -> Dict[str, torch.Tensor]:
    """Function for just scoring.
    Arguments:
        model: no interleaving is supported.
        tokens: prompt tokens extended to be of size [b, max_prompt_length]
        lengths: original prompt length, size: [b]
    Note: Outside of model, other parameters only need to be available on
          rank 0.
    Outputs:
        output_log_probs: log probability of the selected tokens. size: [b, s]
    """

    batch_size = tokens.size(0)
    max_prompt_length = lengths.max().item()
    assert max_prompt_length == tokens.size(1)
    max_sequence_length = min(max_prompt_length, max_position_embeddings)

    # forward step.
    forward_step = ForwardStep(model, batch_size, max_sequence_length, inference_params)

    # ===================
    # Pre-allocate memory
    # ===================

    # Log probability of the sequence (prompt + generated tokens).
    output_log_probs = None
    output_log_probs_size = (batch_size, max_sequence_length - 1)

    if mpu.is_pipeline_last_stage():
        output_log_probs = torch.empty(
            output_log_probs_size,
            dtype=torch.float32,
            device=torch.cuda.current_device(),
        )

    # =============
    # Run inference
    # =============
    with torch.no_grad():
        _, position_ids = _build_attention_mask_and_position_ids(tokens)

        # logits will be meanigful only in the last pipeline stage.
        outputs = forward_step(tokens, position_ids)
        assert outputs is not None
        logits = outputs["logits"]

        if mpu.is_pipeline_last_stage():
            # Always the last stage should have an output.
            assert logits is not None
            assert output_log_probs is not None
            log_probs = F.log_softmax(logits, dim=2)

            # Pick the tokens that we need to get the log
            # probabilities for. Note that next input token is
            # the token which we selected in the current logits,
            # so shift by 1.
            indices = torch.unsqueeze(tokens[:, 1:], 2)
            output_log_probs[:] = torch.gather(log_probs, 2, indices).squeeze(2)

    # ======================================
    # Broadcast to the first pipeline stage.
    # ======================================
    output_log_probs = broadcast_from_last_to_first_pipeline_stage(
        output_log_probs_size, torch.float32, output_log_probs
    )

    return {
        "tokens": tokens,
        "lengths": lengths,
        "output_log_probs": output_log_probs,
    }


def generate_tokens_probs_and_return_on_first_stage(
    model: MegatronModule,
    tokens: torch.Tensor,
    lengths: torch.Tensor,
    inference_params: InferenceParams,
    max_position_embeddings: int,
    termination_id: int,
    vocab_size: int,
    *,
    return_output_log_probs: bool = False,
    return_all_log_probs: bool = False,
    log_prob_tokens: Optional[torch.Tensor] = None,
    top_k: int = 0,
    top_p: float = 0.0,
    temperature: float = 1.0,
) -> Dict[str, torch.Tensor]:
    """Main token generation function.
    Arguments:
        model: no interleaving is supported.
        tokens: prompt tokens extended to be of size [b, max-sequence-length]
        lengths: original prompt length, size: [b]
        return_output_log_probs: flag to calculate the log probability of
            the generated tokens. Note that the log probability is the one
            from the original logit.
        return_all_log_probs: flag to return log probabilities of all the
            vocab tokens for each generated token. Note that unlike output
            log probs, these are only calculated for the generated tokens
            for efficiency.
        log_prob_tokens: if return_all_log_probs is True, then the log probs
            are restricted to the vocab tokens in this tensor.
        top_k, top_p: top-k and top-p sampling parameters.
            Note that top-k = 1 is gready. Also, these paramters are
            exclusive meaning that:
                if top-k > 0 then we expect top-p=0.
                if top-p > 0 then we check for top-k=0.
        temperature: sampling temperature.
    Note: Outside of model, other parameters only need to be available on
          rank 0.
    Outputs: Note that is size is adjusted to a lower value than
             max-sequence-length if generation is terminated early.
        tokens: prompt and generated tokens. size: [b, :]
        generated_sequence_lengths: total length (including prompt) of
            the generated sequence. size: [b]
        output_log_probs: log probability of the selected tokens. size: [b, s]
    """

    batch_size = tokens.size(0)
    min_prompt_length = lengths.min().item()
    max_prompt_length = lengths.max().item()
    max_sequence_length = tokens.size(1)
    tokens_to_generate = max_sequence_length - max_prompt_length

    if max_sequence_length > max_position_embeddings:
        raise ValueError("Length of prompt + tokens_to_generate longer than allowed")

    # forward step.
    forward_step = ForwardStep(model, batch_size, max_sequence_length, inference_params)

    # ===================
    # Pre-allocate memory
    # ===================

    # Log probability of the sequence (prompt + generated tokens).
    output_log_probs = None
    output_log_probs_size = (batch_size, max_sequence_length - 1)

    # Log probability of all vocab tokens for each generated token. We only do generated tokens so that memory
    # usage is not too high.
    all_log_probs = None
    log_prob_vocab_size = (
        vocab_size if log_prob_tokens is None else len(log_prob_tokens)
    )
    all_log_probs_size = (batch_size, tokens_to_generate, log_prob_vocab_size)

    # Lengths of generated seuquence including including prompts.
    generated_sequence_lengths = None
    if mpu.is_pipeline_last_stage():
        if return_output_log_probs:
            output_log_probs = torch.zeros(
                output_log_probs_size,
                dtype=torch.float32,
                device=torch.cuda.current_device(),
            )
        if return_all_log_probs:
            all_log_probs = torch.empty(
                all_log_probs_size,
                dtype=torch.float32,
                device=torch.cuda.current_device(),
            )
        generated_sequence_lengths = (
            torch.ones(
                batch_size, dtype=torch.int64, device=torch.cuda.current_device()
            )
            * max_sequence_length
        )

    # Whether we have reached a termination id.
    is_generation_done = torch.zeros(
        batch_size, dtype=torch.uint8, device=torch.cuda.current_device()
    )

    # Keep track of the indices of the generated token for each prompt in the batch.
    gen_token_indices = torch.full(
        (batch_size,), -1, dtype=torch.int64, device=torch.cuda.current_device()
    )

    # =============
    # Run inference
    # =============

    with torch.inference_mode():
        _, position_ids = _build_attention_mask_and_position_ids(tokens)
        prev_context_length = 0

        if return_output_log_probs:
            lm_logits_mask = None
        else:
            # We do not care about the logits for the prompt tokens, so we can ignore them.
            # This is strictly a performance optimization.
            lm_logits_mask = torch.ones(
                max_sequence_length,
                dtype=torch.bool,
                device=torch.cuda.current_device(),
            )
            lm_logits_mask[: min_prompt_length - 1] = False

        for context_length in range(min_prompt_length, max_sequence_length):
            # Pick the slice that we need to pass through the network.
            tokens2use = tokens[:, prev_context_length:context_length]
            positions2use = position_ids[:, prev_context_length:context_length]
            lm_logits_mask2use = (
                None
                if lm_logits_mask is None
                else lm_logits_mask[prev_context_length:context_length]
            )

            # logits will be meaningful only in the last pipeline stage.
            if (
                not hasattr(model, "_cg_cache")
                or inference_params.sequence_len_offset == 0
            ):
                outputs = forward_step(
                    tokens2use,
                    positions2use,
                    lm_logits_mask=lm_logits_mask2use,
                )
            else:
                outputs = model._cg_cache.run(
                    tokens2use, positions2use, inference_params.sequence_len_offset
                )
                # This is needed since it's on CPU side, and CUDA graph reply doesn't execute this step
                inference_params.sequence_len_offset += 1

            assert outputs is not None
            logits = outputs["logits"]

            if mpu.is_pipeline_last_stage():
                # Always the last stage should have an output.
                assert logits is not None

                # Sample.
                last_token_logits = logits[:, -1, :].float()

                new_sample = sample(
                    last_token_logits,
                    top_k=top_k,
                    top_p=top_p,
                    temperature=temperature,
                    vocab_size=vocab_size,
                )

                # If a prompt length is smaller or equal to the current context
                # length, it means we have started generating tokens
                started = lengths <= context_length
                # Whether we've finished generating the tokens of interest
                finished = context_length >= (lengths + tokens_to_generate)
                # Increment the gen_token_indices if started and not finished
                gen_token_indices[started & ~finished] += 1

                # Update the tokens.
                tokens[started, context_length] = new_sample[started]

                # Calculate the log probabilities.
                if return_output_log_probs:
                    # [b, v]
                    log_probs = F.log_softmax(logits, dim=2)

                    # Pick the tokens that we need to get the log
                    # probabilities for. Note that next input token is
                    # the token which we selected in the current logits,
                    # so shift by 1.
                    indices = torch.unsqueeze(
                        tokens[:, (prev_context_length + 1) : (context_length + 1)], 2
                    )
                    assert output_log_probs is not None
                    output_log_probs[
                        :, prev_context_length:context_length
                    ] = torch.gather(log_probs, 2, indices).squeeze(2)

                if return_all_log_probs:
                    last_token_log_softmax = F.log_softmax(
                        last_token_logits, dim=-1
                    )  # Shape: [b, v]
                    if log_prob_tokens is not None:
                        # Pick only the log probabilities of the tokens in log_prob_tokens
                        last_token_log_softmax = last_token_log_softmax[
                            :, log_prob_tokens
                        ]

                    # Update all_log_probs for the cases where we've started but not finished generating tokens
                    # of interest
                    indices = started & ~finished
                    assert all_log_probs is not None
                    all_log_probs[
                        indices, gen_token_indices[indices], :
                    ] = last_token_log_softmax[indices, :]

            # Update the tokens on the first stage so the next input to
            # the network is correct.
            copy_from_last_to_first_pipeline_stage(
                batch_size, torch.int64, tokens[:, context_length]
            )

            if inference_params is not None:
                # Update the context length for the next token generation if using KV cache.
                prev_context_length = context_length

            # Check if all the sequences have hit the termination_id.
            done = None
            if mpu.is_pipeline_last_stage():
                done_token = (new_sample == termination_id).byte() & started.byte()

                just_finished = (done_token & ~is_generation_done).bool()
                assert generated_sequence_lengths is not None
                generated_sequence_lengths[just_finished.view(-1)] = context_length + 1
                is_generation_done = is_generation_done | done_token
                done = torch.all(is_generation_done)
            done = broadcast_from_last_pipeline_stage(1, torch.uint8, tensor=done)
            if done:  # always use eod token for early termination
                break

    # ===================================================
    # Update the length of based on max generated length.
    # ===================================================

    tokens = tokens[:, : (context_length + 1)]
    if mpu.is_pipeline_last_stage():
        if return_output_log_probs:
            assert output_log_probs is not None
            output_log_probs = output_log_probs[:, :context_length]

    # ======================================
    # Broadcast to the first pipeline stage.
    # ======================================

    generated_sequence_lengths = broadcast_from_last_to_first_pipeline_stage(
        batch_size, torch.int64, generated_sequence_lengths
    )
    if return_output_log_probs:
        output_log_probs_size = (batch_size, context_length)
        output_log_probs = broadcast_from_last_to_first_pipeline_stage(
            output_log_probs_size, torch.float32, output_log_probs
        )
    if return_all_log_probs:
        all_log_probs = broadcast_from_last_to_first_pipeline_stage(
            all_log_probs_size, torch.float32, all_log_probs
        )

    return {
        "tokens": tokens,
        "lengths": generated_sequence_lengths,
        "output_log_probs": output_log_probs,
        "all_log_probs": all_log_probs,
    }


def _build_attention_mask_and_position_ids(
    tokens: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Build the attention mask and postition ids for the input tokens."""

    # Since we are not interested in loss-mask and reset attention/position
    # is also False, eod_token is not used so it is safe to set it to None.
    attention_mask, _, position_ids = get_ltor_masks_and_position_ids(
        data=tokens,
        eod_token=None,
        reset_position_ids=False,
        reset_attention_mask=False,
        eod_mask_loss=False,
    )

    return attention_mask, position_ids

megatron/text_generation/generation.py
persimmon-ai-labs/adept-inference
from typing import Optional, Dict, Tuple, Any

import torch
from torch import Tensor


class InferenceParams:
    """Inference parameters that are passed to the main model in order
    to efficienly calculate and store the context during inference."""

    def __init__(
        self,
        max_batch_size: int,
        max_sequence_len: int,
        lengths_per_sample: Optional[Tensor] = None,
        fused_ft_kernel: bool = False,
    ) -> None:
        # fused_ft_kernel: whether to use FasterTransformer fused single-query attention kernel.
        """Note that offsets are set to zero and we always set the
        flag to allocate memory. After the first call, make sure to
        set this flag to False."""
        self.max_sequence_len = max_sequence_len
        self.max_batch_size = max_batch_size
        self.sequence_len_offset = 0
        self.batch_size_offset = 0
        self.key_value_memory_dict: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {}
        # This is used incase of an encoder-decoder model. The encoder output can be cached
        # and reused for multiple decoder forward passes.
        self.encoder_hidden_state: Dict[Any, Any] = {}
        self.return_encoder_hidden_state = False
        self.fused_ft_kernel = fused_ft_kernel
        self.lengths_per_sample: Tensor = lengths_per_sample
        # Raise import error at initialization time instead of the 1st generation time.
        if fused_ft_kernel:
            try:
                # pylint: disable=import-outside-toplevel,unused-import
                import ft_attention  # type: ignore
            except ImportError as exc:
                raise ImportError("Please install ft_attention from the FlashAttention repo.") from exc

    def reset(self) -> None:
        self.sequence_len_offset = 0
        self.batch_size_offset = 0

    def swap_key_value_dict(self, batch_idx: Any) -> None:
        "swap between batches"
        if len(self.key_value_memory_dict) == 0:
            raise ValueError("should not swap when dict in empty")

        for layer_number in self.key_value_memory_dict.keys():  # pylint: disable=consider-using-dict-items
            inference_key_memory, inference_value_memory = self.key_value_memory_dict[layer_number]
            assert len(batch_idx) == inference_key_memory.shape[1]  ## make sure batch size is the same
            new_inference_key_memory = inference_key_memory[:, batch_idx]
            new_inference_value_memory = inference_value_memory[:, batch_idx]
            self.key_value_memory_dict[layer_number] = (new_inference_key_memory, new_inference_value_memory)

megatron/text_generation/inference_params.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Sampling utilities.
Part of this code is inspired by:
 - https://github.com/ari-holtzman/degen/blob/master/gen.py
 - https://huggingface.co/transformers/_modules/transformers/generation_logits_process.html
"""


from typing import Optional
import torch


def modify_logits_for_top_k_filtering(logits: torch.Tensor, top_k: int) -> None:
    """Set the logits for none top-k values to -inf."""

    filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None]
    logits.masked_fill_(filter_, float("-Inf"))


def modify_logits_for_top_p_filtering(logits: torch.Tensor, top_p: float) -> None:
    """Set the logits for none top-p values to -inf."""

    # First sort and calculate cumulative sum of probabilities.
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

    # Filteration based on the cumulative sum.
    filter_ = cumulative_probs > top_p
    # This shift by 1 is weird and I cannot justify it. This existed
    # in the original implementation:
    #   https://github.com/ari-holtzman/degen/blob/master/gen.py
    # and I guess it is needed so keeping it for now.
    filter_[:, 1:] = filter_[:, :-1].clone()
    # Make sure we at least have one token to select from.
    filter_[..., 0] = 0

    # Fill in the filtered part
    filter_ = filter_.scatter(1, sorted_indices, filter_)
    logits.masked_fill_(filter_, float("-Inf"))


def sample(
    logits: torch.Tensor, top_k: int = 0, top_p: float = 0.0, temperature: float = 1.0, vocab_size: Optional[int] = None
) -> torch.Tensor:
    """Sample and generate a token.
    Note: logits has the dimension [b, v] where b is the batch size
          and v is the vocabulary size.
    If vocab_size is provided, we will make sure the sample that is
    generated is in [0, vocab-size). This will avoid out of vocabulary
    generations due to padding.
    """

    # Check logits for consistency.
    assert logits.ndim == 2, "expected the logits to be of [b, v] shape."
    assert logits.type() == "torch.cuda.FloatTensor", f"input logits should be floats. Got: {logits.type()}"

    # Greedy is just simple argmax.
    if top_k == 1:
        assert top_p == 0.0, "cannot set both greedy and top-p samplings."
        samples = torch.argmax(logits, dim=-1)

    # Top-k or top-p sampling.
    else:
        # Clone so we do not modify the inputs,
        logits = logits.clone()
        # Apply temperature in place.
        if temperature != 1.0:
            logits.div_(temperature)

        if top_k > 1:
            assert top_p == 0.0, "cannot set both top-k and top-p samplings."
            assert top_k <= logits.size(1), "top-k is larger than logit size."
            if vocab_size:
                assert top_k < vocab_size, "top-k is larger than vocab size."
            modify_logits_for_top_k_filtering(logits, top_k)

        elif top_p > 0.0:
            assert top_p <= 1.0, "top-p should be in (0, 1]."
            modify_logits_for_top_p_filtering(logits, top_p)

        # After filtering, we need to recalculate the distribution.
        probs = logits.softmax(dim=-1)
        samples = torch.multinomial(probs, num_samples=1).view(-1)

    # If vocab size is provided, make sure the samples are in
    # in the range [0, vocab-size).
    if vocab_size:
        samples = torch.clamp(samples, min=0, max=(vocab_size - 1))

    return samples

megatron/text_generation/sampling.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tokenization utilities."""

import re
from typing import List, Tuple, Optional, Any, NamedTuple, Dict
import numpy as np

import torch

from megatron import get_args, get_tokenizer
from megatron.tokenizer import AbstractTokenizer
from megatron.mpu.communication import broadcast_int_list, broadcast_tensor

TEXT_REPR_BBOX_OPEN = "<box>"
TEXT_REPR_BBOX_CLOSE = "</box>"
TEXT_REPR_POINT_OPEN = "<point>"
TEXT_REPR_POINT_CLOSE = "</point>"


def convert_generations_to_human_readable_tokens(
    generations: List[List[int]], bos_token: str
) -> List[List[str]]:
    """Convert the list of integers that a model outputs into a human-readable list of tokens.
    Args:
        generations: One list per batch, each of which contains a list of integers to detokenize.
        bos_token: The BOS token that we are using.
    Return:
        A list of lists.
    """
    new_generations = []
    tokenizer = get_tokenizer()

    for generation in generations:
        tokens: List[str] = []
        for i, int_token in enumerate(generation):
            token = tokenizer.inv_vocab[int_token]
            # convert underscore into an empty string when it is first.
            if token[0] == "▁" and (i == 0 or tokens[i - 1] == bos_token):
                token = token[1:]
            # continue processing normally.
            token = re.sub("▁", " ", token)
            tokens.append(token)
        new_generations.append(tokens)
    return new_generations


# ====================================================================== TOKENIZATION  #


def tokenize_prompts(
    prompts: Optional[List[List[str]]],
    max_tokens_to_generate: int,
    max_position_embeddings: int,
    add_BOS: bool,
    rank: int = 0,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Tokenize prompts and make them avaiable on all ranks."""
    assert add_BOS is not None

    # On all ranks set to None so we can pass them to functions
    sizes_list = None
    prompts_tokens_cuda_long_tensor = None
    prompts_length_cuda_long_tensor = None

    # On the specified rank, build the above.
    if torch.distributed.get_rank() == rank:
        assert prompts is not None
        assert max_tokens_to_generate is not None
        # Tensor of tokens padded and their unpadded length.
        (
            prompts_tokens_cuda_long_tensor,
            prompts_length_cuda_long_tensor,
        ) = _tokenize_prompts_and_batch(
            prompts,
            max_tokens_to_generate,
            max_position_embeddings,
            add_BOS,
        )
        # We need the sizes of these tensors for the broadcast
        sizes_list = [
            prompts_tokens_cuda_long_tensor.size(0),  # Batch size
            prompts_tokens_cuda_long_tensor.size(1),  # Num subsequences
            prompts_tokens_cuda_long_tensor.size(2),  # Sequence length
        ]
    # First, broadcast the sizes.
    sizes_tensor = broadcast_int_list(3, int_list=sizes_list, rank=rank)

    # Now that we have the sizes, we can broadcast the tokens
    # and length tensors.
    sizes = sizes_tensor.tolist()
    prompts_tokens_cuda_long_tensor = broadcast_tensor(
        sizes, torch.int64, tensor=prompts_tokens_cuda_long_tensor, rank=rank
    )
    prompts_length_cuda_long_tensor = broadcast_tensor(
        sizes[:2], torch.int64, tensor=prompts_length_cuda_long_tensor, rank=rank
    )

    return prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor


def _tokenize_prompts_and_batch(
    prompts: List[List[str]],
    max_tokens_to_generate: int,
    max_position_embeddings: int,
    add_BOS: bool,  # Same issue with types as above
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Given a set of prompts and number of tokens to generate:
    - tokenize prompts
    - set the sequence length to be the max of length of prompts
      plus the number of tokens we would like to generate
    - pad all the sequences to this length so we can convert them
      into a 3D tensor.
    """
    args = get_args()
    # Tokenize all the prompts.
    tokenizer = get_tokenizer()

    transformed_prompt_tokens = [
        [tokenizer.tokenize(prompt) for prompt in prompt_seq] for prompt_seq in prompts
    ]

    if add_BOS:
        if args.add_bos_prompt_token is not None:
            bos_token = tokenizer.vocab[args.add_bos_prompt_token]
        else:
            bos_token = tokenizer.eod
        prompts_tokens = [
            [[bos_token] + x for x in prompt_seq]
            for prompt_seq in transformed_prompt_tokens
        ]
    else:
        prompts_tokens = transformed_prompt_tokens

    # Now we have a list of list of tokens which each list has a different
    # size. We want to extend this list to:
    #   - incorporate the tokens that need to be generated
    #   - make all the sequences equal length.
    # Get the prompts length.

    prompts_length = [
        [len(x) for x in prompts_tokens_seq] for prompts_tokens_seq in prompts_tokens
    ]
    # Get the max prompts length.
    max_prompt_len: int = np.max(prompts_length)
    # Number of tokens in the each sample of the batch.
    samples_length = min(
        max_prompt_len + max_tokens_to_generate, max_position_embeddings
    )
    if (
        max_prompt_len + max_tokens_to_generate > max_position_embeddings
        and torch.distributed.get_rank() == 0
    ):
        print(
            f"Max subsequence prompt length of {max_prompt_len} + max tokens to generate {max_tokens_to_generate}",
            f"exceeds context length of {max_position_embeddings}. Will generate as many tokens as possible.",
        )
    # Now update the list of list to be of the same size: samples_length.
    for prompt_tokens_seq, prompts_length_seq in zip(prompts_tokens, prompts_length):
        for prompt_tokens, prompt_length in zip(prompt_tokens_seq, prompts_length_seq):
            if len(prompt_tokens) > samples_length:
                raise ValueError(
                    "Length of subsequence prompt exceeds sequence length."
                )
            padding_size = samples_length - prompt_length
            prompt_tokens.extend([tokenizer.eod] * padding_size)

    # Now we are in a structured format, we can convert to tensors.
    prompts_tokens_tensor = torch.cuda.LongTensor(prompts_tokens)
    prompts_length_tensor = torch.cuda.LongTensor(prompts_length)

    return prompts_tokens_tensor, prompts_length_tensor

megatron/text_generation/tokenization.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import collections
import datetime
import json
import os
import threading
import gc
import numbers
from pathlib import Path
import urllib.request
import copy
import traceback
import argparse

from typing import Optional, Callable, Tuple, Dict, Union, Any, List

from dataclasses import dataclass, field

import torch
from torch import Tensor

from flask import Flask, jsonify, request
from flask_cors import CORS, cross_origin
from flask_restful import Api, Resource  # type: ignore
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP

from megatron.tokenizer.tokenizer import AbstractTokenizer
from megatron.model.module import MegatronModule
from megatron.model import DistributedDataParallel as LocalDDP
from megatron.model import Float16Module
from megatron.text_generation.api import generate_and_post_process
from megatron.text_generation.inference_params import InferenceParams
from megatron.text_generation.forward_step import ForwardStep
from megatron.utils import unwrap_model


GENERATE_NUM = 0
BEAM_NUM = 1
lock = threading.Lock()


def add_text_generate_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
    group = parser.add_argument_group(title="text generation")

    group.add_argument(
        "--temperature", type=float, default=1.0, help="Sampling temperature."
    )
    group.add_argument("--top_p", type=float, default=0.0, help="Top p sampling.")
    group.add_argument("--top_k", type=int, default=0, help="Top k sampling.")
    group.add_argument("--port", type=int, default=5000, help="Server port.")
    group.add_argument(
        "--out-seq-length",
        type=int,
        default=1024,
        help="Size of the output generated text.",
    )

    group.add_argument(
        "--use-inference-kv-cache",
        dest="use_inference_kv_cache",
        action="store_true",
        help="Use a KV cache",
    )
    group.add_argument(
        "--no-use-inference-kv-cache",
        dest="use_inference_kv_cache",
        action="store_false",
        help="Don't use a KV cache",
    )
    group.set_defaults(use_inference_kv_cache=True)

    group.add_argument(
        "--fused-ft-kernel",
        dest="fused_ft_kernel",
        action="store_true",
        help="Use fused FT kernel during inference",
    )
    group.add_argument(
        "--no-fused-ft-kernel",
        dest="fused_ft_kernel",
        action="store_false",
        help="Don't use fused FT kernel during inference",
    )
    group.set_defaults(fused_ft_kernel=True)
    group.add_argument(
        "--use-cuda-graph",
        dest="use_cuda_graph",
        action="store_true",
        help="Use CUDA graph",
    )
    group.add_argument(
        "--no-use-cuda-graph",
        dest="use_cuda_graph",
        action="store_false",
        help="Don't use CUDA graph",
    )
    group.set_defaults(use_cuda_graph=True)

    group.add_argument(
        "--add-bos-prompt-token",
        type=str,
        default=None,
        help="Custom BOS token to use instead of EOD",
    )
    group.add_argument(
        "--model-architecture",
        default="GPTModel",
        type=str,
        choices=["GPTModel"],
        help="Model type for Inference.",
    )

    group.add_argument(
        "--inference-max-batch-size",
        type=int,
        help="Maximum batch size for inference.",
    )
    group.add_argument(
        "--inference-max-seqlen",
        type=int,
        help="Maximum sequence length for inference.",
    )

    return parser


def seqlen_to_seqlen_type(seqlen: int) -> int:
    """Convert sequence length to a seqlen_type.
    This is used to determine which cuda graph to use.
    The case work is due to the fact that the attention kernel from FasterTransformer has 3 cases,
    for seqlen < 32, seqlen < 2048, and seqlen >= 2048.
    Arguments:
        seqlen: int
    """
    if seqlen < 32:
        return 0
    elif seqlen < 2048:
        return 1
    else:
        return 2


def seqlen_type_to_max_seqlen(seqlen_type: int) -> int:
    if seqlen_type == 0:
        return 32
    elif seqlen_type == 1:
        return 2048
    elif seqlen_type == 2:
        return 2**32  # Will take the min of this and the actual max_seqlen
    else:
        raise ValueError(f"seqlen_type {seqlen_type} not supported")


BatchSize = int
SeqLenType = int
SeqLenOffset = int
ForwardStepOutput = Optional[Dict[str, Any]]


@dataclass
class DecodingCGCache:
    max_batch_size: int = 0
    max_seqlen: int = 0
    device: Optional[torch.device] = None
    dtype: Optional[torch.dtype] = None
    callables: Dict[
        Tuple[BatchSize, SeqLenType],
        Callable[[Tensor, Tensor, SeqLenOffset], ForwardStepOutput],
    ] = field(default_factory=dict)
    mempool: Optional[torch.Tensor] = None
    inference_params: Optional[InferenceParams] = None
    run: Optional[
        Callable[
            [torch.Tensor, torch.Tensor, int],
            Callable[[torch.Tensor, torch.Tensor, int], Optional[Dict[str, Any]]],
        ]
    ] = None


@torch.inference_mode()  # type: ignore
def update_graph_cache(
    model: MegatronModule,
    cg_cache: Optional[DecodingCGCache],
    inference_params: InferenceParams,
    batch_size: int,
    seqlen_og: int,
    max_seqlen: int,
    dtype: Optional[torch.dtype] = None,
    n_warmups: int = 2,
):
    # PyTorch docs says at least 2 warmups are needed.
    if cg_cache is None:
        cg_cache = DecodingCGCache()
    param_example = next(iter(model.parameters()))
    device = param_example.device
    if dtype is None:
        dtype = param_example.dtype
    if (
        (device, dtype) != (cg_cache.device, cg_cache.dtype)
        or batch_size > cg_cache.max_batch_size
        or max_seqlen > cg_cache.max_seqlen
    ):  # Invalidate the cg_cache
        cg_cache.callables = {}
        cg_cache.mempool = None
        gc.collect()
        cg_cache.device, cg_cache.dtype = device, dtype
        cg_cache.max_batch_size, cg_cache.max_seqlen = batch_size, max_seqlen
        lengths_per_sample = torch.full(
            (batch_size,), seqlen_og, dtype=torch.int32, device=device
        )
        cg_cache.inference_params = InferenceParams(
            max_sequence_len=max_seqlen,
            max_batch_size=batch_size,
            fused_ft_kernel=True,
            lengths_per_sample=lengths_per_sample,
        )
        cg_cache.inference_params.sequence_len_offset = seqlen_og
        cg_cache.inference_params.key_value_memory_dict = (
            inference_params.key_value_memory_dict
        )
        cg_cache.mempool = torch.cuda.graphs.graph_pool_handle()
    for s_type in range(
        seqlen_to_seqlen_type(seqlen_og), seqlen_to_seqlen_type(max_seqlen) + 1
    ):
        if (batch_size, s_type) not in cg_cache.callables:
            max_seqlen_ = min(
                max(seqlen_og, seqlen_type_to_max_seqlen(s_type)), max_seqlen
            )
            assert cg_cache.inference_params is not None
            cg_cache.callables[batch_size, s_type] = capture_graph(
                model,
                cg_cache.inference_params,
                batch_size,
                max_seqlen_,
                mempool=cg_cache.mempool,
                n_warmups=n_warmups,
            )

    def dispatch(
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        seqlen: int,
    ) -> Any:
        batch_size = input_ids.shape[0]
        assert cg_cache is not None
        return cg_cache.callables[batch_size, seqlen_to_seqlen_type(seqlen)](
            input_ids, position_ids, seqlen
        )

    assert cg_cache.inference_params is not None
    cg_cache.run = dispatch
    cg_cache.inference_params.reset()  # Reset so it's not confusing
    return cg_cache


def capture_graph(
    model: MegatronModule,
    inference_params: InferenceParams,
    batch_size: int,
    max_seqlen: int,
    mempool: Optional[torch.Tensor] = None,
    n_warmups: int = 2,
) -> Callable[[torch.Tensor, torch.Tensor, int], Optional[Dict[str, Any]]]:
    device = next(iter(model.parameters())).device
    input_ids = torch.full((batch_size, 1), 0, dtype=torch.long, device=device)
    position_ids = torch.full((batch_size, 1), 0, dtype=torch.long, device=device)
    sequence_len_offset_og = inference_params.sequence_len_offset
    # TD [2023-04-14]: important for correctness of the FT's attention kernel, as seqlen_cpu is
    # used to determine the size of smem. Hence seqlen_cpu must be >= lengths_per_sample.
    inference_params.sequence_len_offset = max_seqlen - 1
    inference_params.lengths_per_sample[:] = max_seqlen - 1

    forward_step = ForwardStep(model, batch_size, max_seqlen, inference_params)
    # Warmup before capture
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(n_warmups):
            outputs = forward_step(input_ids, position_ids)
        s.synchronize()
        # This might be needed for correctness if we run with NCCL_GRAPH_MIXING_SUPPORT=0,
        # which requires that graph launch and non-captured launch to not overlap (I think,
        # that's how I interpret the documentation). I'm not sure if this is required.
        if torch.distributed.is_initialized():
            torch.distributed.barrier()
    torch.cuda.current_stream().wait_stream(s)
    # Captures the graph
    # To allow capture, automatically sets a side stream as the current stream in the context
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph, pool=mempool):
        outputs = forward_step(input_ids, position_ids)

    def run(
        new_input_ids: torch.Tensor, new_position_ids: torch.Tensor, seqlen: int
    ) -> Optional[Dict[str, Any]]:
        inference_params.lengths_per_sample[:] = seqlen
        input_ids.copy_(new_input_ids)
        position_ids.copy_(new_position_ids)
        graph.replay()
        return outputs

    inference_params.sequence_len_offset = sequence_len_offset_og
    return run


def allocate_kv_cache(
    batch_size: int,
    seqlen: int,
    num_layers: int,
    num_heads_per_partition: int,
    head_dim: int,
    fused_ft_kernel: bool = False,
    dtype: torch.dtype = None,
    device: Optional[torch.device] = None,
) -> torch.Tensor:
    kv_cache = {}
    for i in range(1, num_layers + 1):
        if not fused_ft_kernel:
            k_cache = torch.empty(
                seqlen,
                batch_size,
                num_heads_per_partition,
                head_dim,
                dtype=dtype,
                device=device,
            )
            v_cache = torch.empty(
                seqlen,
                batch_size,
                num_heads_per_partition,
                head_dim,
                dtype=dtype,
                device=device,
            )
        else:
            dtype = torch.bfloat16 if dtype is None else dtype
            assert dtype in [torch.float16, torch.bfloat16, torch.float32]
            packsize = 4 if dtype == torch.float32 else 8
            assert head_dim % packsize == 0
            k_cache = torch.empty(
                batch_size,
                num_heads_per_partition,
                head_dim // packsize,
                seqlen,
                packsize,
                dtype=dtype,
                device=device,
            )
            v_cache = torch.empty(
                batch_size,
                num_heads_per_partition,
                seqlen,
                head_dim,
                dtype=dtype,
                device=device,
            )
        kv_cache[i] = (k_cache, v_cache)
    return kv_cache


def setup_model(
    model: MegatronModule,
    model_architecture: str,
    use_inference_kv_cache: bool,
    use_fused_ft_kernel: bool,
    use_cuda_graph: bool,
    max_batch_size: int,
    max_seq_len: int,
) -> Optional[InferenceParams]:
    if use_inference_kv_cache:
        print(f"Initializing KV cache with {max_batch_size=}, {max_seq_len=}")
        inference_params = InferenceParams(
            max_batch_size=max_batch_size,
            max_sequence_len=max_seq_len,
            fused_ft_kernel=use_fused_ft_kernel,
        )
    else:
        print("Disabling KV cache")
        inference_params = None

    example_param = next(iter(model.parameters()))
    layers = (
        model.module.language_model.encoder.layers
        if model_architecture in ["GPTModel"]
        else model.module.decoder.decoder.layers
    )
    attn_layer = layers[0].self_attention

    if inference_params is not None:
        inference_params.key_value_memory_dict = allocate_kv_cache(
            max_batch_size,
            max_seq_len,
            num_layers=len(layers),
            num_heads_per_partition=attn_layer.num_attention_heads_per_partition,
            head_dim=attn_layer.hidden_size_per_attention_head,
            fused_ft_kernel=inference_params.fused_ft_kernel,
            dtype=example_param.dtype,
            device=example_param.device,
        )
        if use_cuda_graph:
            # Capture the graph for the generation step
            model._cg_cache = None
            # Go from large to small batch sizes so we don't need to invalidate the graph cache
            for batch_size in reversed(range(1, max_batch_size + 1)):
                model._cg_cache = update_graph_cache(
                    model, model._cg_cache, inference_params, batch_size, 1, max_seq_len
                )

    return inference_params


def print_request_json() -> None:
    json_to_dump = copy.deepcopy(request.get_json())
    print(json.dumps(json_to_dump), flush=True)


class MegatronGenerate(Resource):  # type: ignore
    def __init__(
        self,
        model: MegatronModule,
        inference_params: InferenceParams,
        params_dtype: torch.dtype,
        max_position_embeddings: int,
        termination_id: int,
        tokenizer: AbstractTokenizer,
    ):
        self.model = model
        if inference_params is not None:
            inference_params.reset()
        self.inference_params = inference_params
        self.params_dtype = params_dtype
        self.max_position_embeddings = max_position_embeddings
        self.termination_id = termination_id
        self.tokenizer = tokenizer

    @staticmethod
    def send_do_generate() -> None:
        choice = torch.cuda.LongTensor([GENERATE_NUM])
        torch.distributed.broadcast(choice, 0)

    @cross_origin()
    def put(self) -> Union[Tuple[str, int], str]:
        print("request IP: " + str(request.remote_addr))
        print_request_json()
        print("current time: ", datetime.datetime.now())

        if not "prompts" in request.get_json():
            return "prompts argument required", 400

        if "max_len" in request.get_json():
            return "max_len is no longer used.  Replace with tokens_to_generate", 400

        if "sentences" in request.get_json():
            return "sentences is no longer used.  Replace with prompts", 400

        prompts = request.get_json()["prompts"]

        if type(prompts) != list:
            prompts = [prompts]

        if len(prompts) > 128:
            return "Maximum number of prompts is 128", 400

        max_tokens_to_generate = (
            64  # Choosing hopefully sane default.  Full sequence is slow
        )
        if "tokens_to_generate" in request.get_json():
            max_tokens_to_generate = request.get_json()["tokens_to_generate"]
            if not isinstance(max_tokens_to_generate, int):
                return "tokens_to_generate must be an integer greater than 0"
            if max_tokens_to_generate < 0:
                return (
                    "tokens_to_generate must be an integer greater than or equal to 0"
                )

        logprobs = False
        if "logprobs" in request.get_json():
            logprobs = request.get_json()["logprobs"]
            if not isinstance(logprobs, bool):
                return "logprobs must be a boolean value"

        if max_tokens_to_generate == 0 and not logprobs:
            return "tokens_to_generate=0 implies logprobs should be True"

        temperature = 1.0
        if "temperature" in request.get_json():
            temperature = request.get_json()["temperature"]
            if not isinstance(temperature, numbers.Number):
                return (
                    "temperature must be a positive number less than or equal to 100.0"
                )
            temperature = float(temperature)
            if not (0.0 < temperature <= 100.0):
                return (
                    "temperature must be a positive number less than or equal to 100.0"
                )

        top_k = 0
        if "top_k" in request.get_json():
            top_k = request.get_json()["top_k"]
            if not isinstance(top_k, int):
                return "top_k must be an integer equal to or greater than 0 and less than or equal to 1000"
            if not (0 <= top_k <= 1000):
                return "top_k must be equal to or greater than 0 and less than or equal to 1000"

        top_p = 0.0
        if "top_p" in request.get_json():
            top_p = request.get_json()["top_p"]
            if not isinstance(top_p, numbers.Number):
                return "top_p must be a positive float less than or equal to 1.0"
            top_p = float(top_p)
            if top_p > 0.0 and top_k > 0.0:
                return "cannot set both top-k and top-p samplings."
            if not (0 <= top_p <= 1.0):
                return "top_p must be less than or equal to 1.0"

        add_BOS = False
        if "add_BOS" in request.get_json():
            add_BOS = request.get_json()["add_BOS"]
            if not isinstance(add_BOS, bool):
                return "add_BOS must be a boolean value"

        random_seed = -1
        if "random_seed" in request.get_json():
            random_seed = request.get_json()["random_seed"]
            if not isinstance(random_seed, int):
                return "random_seed must be integer"
            if random_seed < 0:
                return "random_seed must be a positive integer"

        no_log = False
        if "no_log" in request.get_json():
            no_log = request.get_json()["no_log"]
            if not isinstance(no_log, bool):
                return "no_log must be a boolean value"

        if any(len(prompt) == 0 for prompt in prompts) and not add_BOS:
            return "Empty prompts require add_BOS=true"

        with lock:  # Need to get lock to keep multiple threads from hitting code
            if not no_log:
                print("request IP: " + str(request.remote_addr))
                print_request_json()
                start_time = datetime.datetime.now()
                print("start time: ", start_time, flush=True)

            try:
                MegatronGenerate.send_do_generate()  # Tell other ranks we're doing generate
                retval = generate_and_post_process(
                    self.model,
                    params_dtype=self.params_dtype,
                    max_position_embeddings=self.max_position_embeddings,
                    termination_id=self.termination_id,
                    tokenizer=self.tokenizer,
                    prompts=[[prompt] for prompt in prompts],
                    max_tokens_to_generate=max_tokens_to_generate,
                    inference_params=self.inference_params,
                    return_output_log_probs=logprobs,
                    top_k_sampling=top_k,
                    top_p_sampling=top_p,
                    temperature=temperature,
                    add_BOS=add_BOS,
                    random_seed=random_seed,
                )
                assert retval is not None
                (
                    response,
                    response_logprobs,
                    _,
                    generations,
                    _,
                    human_readable_tokens,
                ) = retval

                end_time = datetime.datetime.now()
                print("Query latency: ", end_time - start_time, flush=True)
                output = {
                    "text": response,
                    "logprobs": response_logprobs,
                    "generations": generations,
                    "human_readable_tokens": human_readable_tokens,
                }
                return jsonify(output)  # type: ignore

            except ValueError as ve:
                traceback.print_exc()
                return str(ve.args[0])


class Healthcheck(Resource):  # type: ignore
    def get(self) -> Tuple[Dict[str, bool], int, Dict[str, str]]:
        return {"success": True}, 200, {"ContentType": "application/json"}


class MegatronServer(object):
    def __init__(
        self,
        model: MegatronModule,
        inference_params: InferenceParams,
        params_dtype: torch.dtype,
        max_position_embeddings: int,
        termination_id: int,
        tokenizer: AbstractTokenizer,
        port: int = 5000,
    ):
        self.port = port
        self.app = Flask(__name__, static_url_path="")
        _ = CORS(self.app)
        self.app.config["CORS_HEADERS"] = "Content-Type"
        api = Api(self.app)
        api.add_resource(
            MegatronGenerate,
            "/api",
            resource_class_args=[
                model,
                inference_params,
                params_dtype,
                max_position_embeddings,
                termination_id,
                tokenizer,
            ],
        )
        api.add_resource(Healthcheck, "/healthz")

    def run(self, url: str) -> None:
        self.app.run(url, port=self.port, threaded=True, debug=False)

megatron/text_generation_server.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from .tokenizer import build_tokenizer, AbstractTokenizer

megatron/tokenizer/__init__.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Megatron tokenizers."""

from abc import ABCMeta
from abc import abstractmethod

from typing import Dict, Any, List, Optional
import sentencepiece as spm  # type: ignore


class AbstractTokenizer(metaclass=ABCMeta):
    """Abstract class for tokenizer."""

    def __init__(self, name: str) -> None:
        self.name = name
        super().__init__()

    @property
    @abstractmethod
    def vocab_size(self) -> int:
        """Number of distinct tokens in the vocabulary."""
        pass

    @property
    @abstractmethod
    def vocab(self) -> Dict[str, int]:
        """Dictionary from vocab text token to id token."""
        pass

    @property
    @abstractmethod
    def inv_vocab(self) -> Dict[int, str]:
        """Dictionary from vocab id token to text token."""
        pass

    @abstractmethod
    def tokenize(self, text: str) -> List[int]:
        """Tokenize the text."""
        pass

    def encode(self, text: str) -> List[int]:
        """Encode the text."""
        return self.tokenize(text)

    def __call__(self, text: str) -> List[int]:
        """Syntactic sugar for tokenize."""
        return self.tokenize(text)

    @abstractmethod
    def detokenize(self, token_ids: List[int]) -> str:
        """Transform tokens back to a string."""
        pass

    def decode(self, ids: List[int]) -> str:
        """Decode the ids."""
        return self.detokenize(ids)

    @property
    def cls(self) -> int:
        raise NotImplementedError(f"CLS is not provided for {self.name} tokenizer")

    @property
    def sep(self) -> int:
        raise NotImplementedError(f"SEP is not provided for {self.name} tokenizer")

    @property
    def pad(self) -> int:
        raise NotImplementedError(f"PAD is not provided for {self.name} tokenizer")

    @property
    def eod(self) -> int:
        raise NotImplementedError(f"EOD is not provided for {self.name} tokenizer")

    @property
    def eod_token_text(self) -> str:
        """The EOD token string."""
        return self.decode([self.eod])

    @property
    def mask(self) -> int:
        """Get the mask token."""
        raise NotImplementedError(f"MASK is not provided for {self.name} tokenizer")

    @property
    def pad_token(self) -> str:
        """Get the pad token."""
        raise NotImplementedError

    @property
    def max_len_single_sentence(self) -> int:
        """Get the max length of a single sentence."""
        raise NotImplementedError

    def __len__(self) -> int:
        """Get the length of the tokenizer."""
        raise NotImplementedError

    @property
    def eos_token_id(self) -> int:
        """Get the id of the EOS token."""
        raise NotImplementedError

    @property
    def eos_token(self) -> str:
        """Get the EOS token."""
        raise NotImplementedError


class _SentencePieceTokenizer(AbstractTokenizer):
    """Sentece piece tokenizer."""

    def __init__(self, model_file: str) -> None:
        name = "Sentence Piece"
        super().__init__(name)

        # no-member is firing here but the code works fine!
        # pylint: disable=no-member
        self._tokenizer = spm.SentencePieceProcessor()
        self._tokenizer.load(model_file)
        self._vocab_size = self._tokenizer.vocab_size()
        self._tokens = [self._tokenizer.id_to_piece(t) for t in range(self.vocab_size)]
        self._vocab = {t: i for i, t in enumerate(self._tokens)}
        self._eod_id = None
        # look for end of document id
        for idx, token in enumerate(self._tokens):
            if token == "|ENDOFTEXT|":
                self._eod_id = idx
                break
        if self._eod_id is None:
            self._eod_id = self._tokenizer.eos_id()
        assert self._eod_id is not None

    @property
    def tokenizer(self) -> spm.SentencePieceProcessor:
        return self._tokenizer

    @property
    def vocab_size(self) -> int:
        return int(self._vocab_size)

    @property
    def vocab(self) -> Dict[str, int]:
        return self._vocab

    @property
    def inv_vocab(self):  # type: ignore
        return self._tokens

    def tokenize(self, text: str):  # type: ignore
        # pylint: disable=bare-except, no-member
        try:
            tokenized = self._tokenizer.encode_as_ids(text)
        except:
            tokenized = None
        return tokenized

    def pieces(self, text: str) -> List[str]:
        # pylint: disable=no-member
        pieces: List[str] = self._tokenizer.encode_as_pieces(text)
        return pieces

    def detokenize(self, token_ids: List[int]) -> str:
        # pylint: disable=no-member
        return self._tokenizer.decode_ids(token_ids)  # type: ignore

    @property
    def eod(self) -> int:
        return self._eod_id  # type: ignore

    @property
    def eos_token_id(self) -> int:
        """Id of the end of sentence token in the vocabulary."""
        eos_id: int = self._eod_id  # type: ignore
        return eos_id


def megatron_initialize_tokenizer(
    tokenizer_type: str,
    sp_model_file: Optional[str] = None,
) -> AbstractTokenizer:
    """Initialize the tokenizer."""
    if tokenizer_type == "SentencePiece":
        assert sp_model_file is not None
        tokenizer = _SentencePieceTokenizer(sp_model_file)
    else:
        raise NotImplementedError(f"{tokenizer_type} tokenizer is not implemented.")
    return tokenizer


def _vocab_size_with_padding(orig_vocab_size: int, args: Any) -> int:
    """Pad vocab size so it is divisible by model parallel size and
    still having GPU friendly size."""

    after = orig_vocab_size
    multiple = args.make_vocab_size_divisible_by * args.tensor_model_parallel_size
    while (after % multiple) != 0:
        after += 1
    if args.rank == 0:
        print(
            f" > padded vocab (sz: {orig_vocab_size}) w/ {after - orig_vocab_size} dummy toks (new sz: {after})",
            flush=True,
        )
    return after


def build_tokenizer(args: Any) -> AbstractTokenizer:
    """Initialize tokenizer."""
    if args.rank == 0:
        print(
            f"> building {args.tokenizer_type} tokenizer ...",
            flush=True,
        )

    # Select and instantiate the tokenizer.
    if args.tokenizer_type == "SentencePiece":
        assert args.sp_model_file is not None
        tokenizer = _SentencePieceTokenizer(f"{args.sp_model_file}")
    else:
        raise NotImplementedError(
            f"{args.tokenizer_type} tokenizer is not implemented."
        )

    # Add vocab size.
    args.padded_vocab_size = _vocab_size_with_padding(tokenizer.vocab_size, args)

    return tokenizer

megatron/tokenizer/tokenizer.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Pretrain utilities."""

import math
import sys
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP

from megatron import (
    get_args,
    mpu,
)
from megatron.checkpointing import (
    load_state_dicts_and_update_args,
    update_model_and_optim_from_loaded_data,
)
from megatron.initialize import initialize_megatron
from megatron.model import DistributedDataParallel as LocalDDP
from megatron.model import Float16Module, ModelType

from megatron.utils import (
    unwrap_model,
)


def get_model(
    model_provider_func,
    model_type=ModelType.encoder_or_decoder,
    wrap_with_ddp=True,
):
    """Build the model."""
    args = get_args()
    args.model_type = model_type

    # Build model.
    if (
        mpu.get_pipeline_model_parallel_world_size() > 1
        and args.virtual_pipeline_model_parallel_size is not None
    ):
        assert (
            model_type != ModelType.encoder_and_decoder
        ), "Interleaved schedule not supported for model with both encoder and decoder"
        model = []
        for i in range(args.virtual_pipeline_model_parallel_size):
            mpu.set_virtual_pipeline_model_parallel_rank(i)
            # Set pre_process and post_process only after virtual rank is set.
            pre_process = mpu.is_pipeline_first_stage()
            post_process = mpu.is_pipeline_last_stage()
            this_model = model_provider_func(
                pre_process=pre_process, post_process=post_process
            )
            this_model.model_type = model_type
            model.append(this_model)
    else:
        pre_process = mpu.is_pipeline_first_stage()
        post_process = mpu.is_pipeline_last_stage()
        add_encoder = True
        add_decoder = True
        if model_type == ModelType.encoder_and_decoder:
            if mpu.get_pipeline_model_parallel_world_size() > 1:
                assert (
                    args.pipeline_model_parallel_split_rank is not None
                ), "Split rank needs to be specified for model with both encoder and decoder"
                rank = mpu.get_pipeline_model_parallel_rank()
                split_rank = args.pipeline_model_parallel_split_rank
                world_size = mpu.get_pipeline_model_parallel_world_size()
                pre_process = rank == 0 or rank == split_rank
                post_process = (rank == (split_rank - 1)) or (rank == (world_size - 1))
                add_encoder = mpu.is_pipeline_stage_before_split()
                add_decoder = mpu.is_pipeline_stage_after_split()
            model = model_provider_func(
                pre_process=pre_process,
                post_process=post_process,
                add_encoder=add_encoder,
                add_decoder=add_decoder,
            )
        else:
            model = model_provider_func(
                pre_process=pre_process, post_process=post_process
            )
        model.model_type = model_type

    if not isinstance(model, list):
        model = [model]

    # Set tensor model parallel attributes if not set.
    # Only parameters that are already tensor model parallel have these
    # attributes set for them. We should make sure the default attributes
    # are set for all params so the optimizer can use them.
    for model_module in model:
        for param in model_module.parameters():
            mpu.set_defaults_if_not_set_tensor_model_parallel_attributes(param)

    # Print number of parameters.
    if mpu.get_data_parallel_rank() == 0:
        print(
            " > number of parameters on (tensor, pipeline) "
            "model parallel rank ({}, {}): {}".format(
                mpu.get_tensor_model_parallel_rank(),
                mpu.get_pipeline_model_parallel_rank(),
                sum(
                    [
                        sum([p.nelement() for p in model_module.parameters()])
                        for model_module in model
                    ]
                ),
            ),
            flush=True,
        )

    # GPU allocation.
    for model_module in model:
        model_module.cuda(torch.cuda.current_device())

    # Fp16 conversion.
    if args.fp16 or args.bf16:
        model = [Float16Module(model_module, args) for model_module in model]

    if wrap_with_ddp:
        if args.DDP_impl == "torch":
            i = torch.cuda.current_device()
            model = [
                torchDDP(
                    model_module,
                    device_ids=[i],
                    output_device=i,
                    process_group=mpu.get_data_parallel_group(),
                )
                for model_module in model
            ]

        elif args.DDP_impl == "local":
            model = [
                LocalDDP(
                    model_module,
                    args.accumulate_allreduce_grads_in_fp32,
                    args.use_contiguous_buffers_in_local_ddp,
                )
                for model_module in model
            ]
            # broad cast params from data parallel src rank to other data parallel ranks
            if args.data_parallel_random_init:
                for model_module in model:
                    model_module.broadcast_params()
        else:
            raise NotImplementedError(
                "Unknown DDP implementation specified: "
                "{}. Exiting.".format(args.DDP_impl)
            )

    return model

megatron/training.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""General utilities."""

import re
import os
import sys
import warnings

import torch
from torch.nn.parallel import DistributedDataParallel as torchDDP
from typing import Optional

from megatron import get_args
from megatron import get_adlr_autoresume
from megatron import mpu
from megatron.model.module import param_is_not_shared
from megatron.mpu.layers import param_is_tensor_parallel_unique


def unwrap_model(model, module_instances=(torchDDP)):
    return_list = True
    if not isinstance(model, list):
        model = [model]
        return_list = False
    unwrapped_model = []
    for model_module in model:
        while isinstance(model_module, module_instances):
            model_module = model_module.module
        unwrapped_model.append(model_module)
    if not return_list:
        return unwrapped_model[0]
    return unwrapped_model


def calculate_num_params_and_params_l2_norm(model):
    """Calculate l2 norm of parameters"""
    import amp_C
    from apex.multi_tensor_apply import multi_tensor_applier

    args = get_args()
    if not isinstance(model, list):
        model = [model]
    # Remove duplicate params.
    params_data = []
    num_params = 0
    for model_ in model:
        for param in model_.parameters():
            is_not_shared = param_is_not_shared(param)
            is_tp_unique = param_is_tensor_parallel_unique(param)
            if is_not_shared and is_tp_unique:
                num_params += param.numel()
                if args.bf16:
                    params_data.append(param.data.float())
                else:
                    params_data.append(param.data)
    # Calculate norm
    dummy_overflow_buf = torch.cuda.IntTensor([0])
    norm, _ = multi_tensor_applier(
        amp_C.multi_tensor_l2norm,
        dummy_overflow_buf,
        [params_data],
        False,  # no per-parameter norm
    )
    norm_2 = norm * norm
    # Sum across all model-parallel GPUs.
    torch.distributed.all_reduce(norm_2, op=torch.distributed.ReduceOp.SUM, group=mpu.get_model_parallel_group())

    # Sum num params across all model-parallel GPUs.
    num_params = torch.tensor(num_params, dtype=torch.int64, device=torch.cuda.current_device())
    torch.distributed.all_reduce(num_params, op=torch.distributed.ReduceOp.SUM, group=mpu.get_model_parallel_group())

    return num_params.item(), norm_2.item() ** 0.5


def calculate_per_layer_grad_norms(model, normalizing_factor):
    """Calculate per-layer grad norms. Assumes gradients have already been all-reduced across data group."""
    import amp_C
    from apex.multi_tensor_apply import multi_tensor_applier

    if mpu.get_pipeline_model_parallel_world_size() > 1:
        warnings.warn(
            "calculate_per_layer_grad_norms is not yet implemented for pipeline parallel models, skipping calculation!"
        )
        return {}
    if mpu.get_tensor_model_parallel_world_size() > 1:
        warnings.warn(
            "calculate_per_layer_grad_norms is not yet implemented for tensor parallel models, skipping calculation!"
        )
        return {}

    if not isinstance(model, list):
        model = [model]
    # Remove duplicate params.
    names = []
    grads = []
    for model_ in model:
        for name, param in model_.named_parameters():
            if hasattr(param, "main_grad"):
                grad = param.main_grad
            elif param.grad is not None:
                grad = param.grad
            else:
                continue
            names.append(name)
            grads.append(grad)
    # Calculate norms.
    dummy_overflow_buf = torch.cuda.IntTensor([0])
    _, norms = multi_tensor_applier(
        amp_C.multi_tensor_l2norm,
        dummy_overflow_buf,
        [grads],
        True,  # include per-parameter norm
    )

    norms_2 = norms * norms
    # Sum across all model-parallel GPUs.
    torch.distributed.all_reduce(norms_2, op=torch.distributed.ReduceOp.SUM, group=mpu.get_model_parallel_group())
    norms = norms_2**0.5

    grad_norms = {}
    for i, name in enumerate(names):
        # Consolidate stats for individual layers into a single combined stat.
        name = re.sub(r"\.layers\.[0-9]+", ".all_layers_combined", name)
        # Consolidate stats for individual qk_layernorms into a single combined stat.
        name = re.sub(r"\.q_layernorms\.[0-9]+", ".all_q_layernorms_combined", name)
        name = re.sub(r"\.k_layernorms\.[0-9]+", ".all_k_layernorms_combined", name)
        if name in grad_norms:
            grad_norms[name] += norms[i]
        else:
            grad_norms[name] = norms[i]

    for k in grad_norms.keys():
        grad_norms[k] = grad_norms[k] / normalizing_factor

    return grad_norms


def average_losses_across_data_parallel_group(losses):
    """Reduce a tensor of losses across all GPUs."""
    averaged_losses = torch.cat([loss.clone().detach().view(1) for loss in losses])
    torch.distributed.all_reduce(averaged_losses, group=mpu.get_data_parallel_group())
    averaged_losses = averaged_losses / torch.distributed.get_world_size(group=mpu.get_data_parallel_group())

    return averaged_losses


def report_memory(name):
    """Simple GPU memory report."""
    mega_bytes = 1024.0 * 1024.0
    string = name + " memory (MB)"
    string += " | allocated: {}".format(torch.cuda.memory_allocated() / mega_bytes)
    string += " | max allocated: {}".format(torch.cuda.max_memory_allocated() / mega_bytes)
    string += " | reserved: {}".format(torch.cuda.memory_reserved() / mega_bytes)
    string += " | max reserved: {}".format(torch.cuda.max_memory_reserved() / mega_bytes)
    if mpu.get_data_parallel_rank() == 0:
        print("[Rank {}] {}".format(torch.distributed.get_rank(), string), flush=True)


def check_adlr_autoresume_termination(iteration, model, optimizer, opt_param_scheduler):
    """Check for autoresume signal and exit if it is received."""
    from megatron.checkpointing import save_checkpoint

    args = get_args()
    autoresume = get_adlr_autoresume()
    # Add barrier to ensure consistnecy.
    torch.distributed.barrier()
    if autoresume.termination_requested():
        if args.save:
            save_checkpoint(iteration, model, optimizer, opt_param_scheduler)
        print_rank_0(">>> autoresume termination request found!")
        if torch.distributed.get_rank() == 0:
            autoresume.request_resume()
        print_rank_0(">>> training terminated. Returning")
        sys.exit(0)


def get_ltor_masks_and_position_ids(data, eod_token, reset_position_ids, reset_attention_mask, eod_mask_loss):
    """Build masks and position id for left to right model."""

    # Extract batch size and sequence length.
    micro_batch_size, seq_length = data.size()

    # Attention mask (lower triangular).
    if reset_attention_mask:
        att_mask_batch = micro_batch_size
    else:
        att_mask_batch = 1
    attention_mask = torch.tril(torch.ones((att_mask_batch, seq_length, seq_length), device=data.device)).view(
        att_mask_batch, 1, seq_length, seq_length
    )

    # Loss mask.
    loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device)
    if eod_mask_loss:
        loss_mask[data == eod_token] = 0.0

    # Position ids.
    position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)
    position_ids = position_ids.unsqueeze(0).expand_as(data)
    # We need to clone as the ids will be modifed based on batch index.
    if reset_position_ids:
        position_ids = position_ids.clone()

    if reset_position_ids or reset_attention_mask:
        # Loop through the batches:
        for b in range(micro_batch_size):
            # Find indecies where EOD token is.
            eod_index = position_ids[b, data[b] == eod_token]
            # Detach indecies from positions if going to modify positions.
            if reset_position_ids:
                eod_index = eod_index.clone()

            # Loop through EOD indecies:
            prev_index = 0
            for j in range(eod_index.size()[0]):
                i = eod_index[j]
                # Mask attention loss.
                if reset_attention_mask:
                    attention_mask[b, 0, (i + 1) :, : (i + 1)] = 0
                # Reset positions.
                if reset_position_ids:
                    position_ids[b, (i + 1) :] -= i + 1 - prev_index
                    prev_index = i + 1

    # Convert attention mask to binary:
    attention_mask = attention_mask < 0.5

    return attention_mask, loss_mask, position_ids


def print_rank_0(*args):
    """If distributed is initialized, print only on rank 0."""
    if torch.distributed.is_initialized():
        if torch.distributed.get_rank() == 0:
            print(*args, flush=True)
    elif os.environ.get("RANK", "0") == "0":
        print(*args, flush=True)


def print_once_per_data_rank(*args):
    if not torch.distributed.is_initialized() or (
        mpu.get_tensor_model_parallel_rank() == 0 and mpu.get_pipeline_model_parallel_rank() == 0
    ):
        print(*args, flush=True)


def is_last_rank():
    return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1)


def print_rank_last(*args):
    """If distributed is initialized, print only on last rank."""
    if torch.distributed.is_initialized():
        if is_last_rank():
            print(*args, flush=True)
    else:
        print(*args, flush=True)


def get_total_eval_iters(train_iters: Optional[int], eval_interval: int, iter_per_eval: int) -> Optional[int]:
    """Helper function to calculate evaluation iters for iteration based training, otherwise ignore."""
    if train_iters is not None:
        return (train_iters // eval_interval + 1) * iter_per_eval

    return None

megatron/utils.py
persimmon-ai-labs/adept-inference
# coding=utf-8
# Copyright (c) 2023 ADEPT AI LABS INC.
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Sample Generate GPT"""
import os
import sys
from typing import Any, Optional, Tuple

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)))
import torch

from megatron import get_args, get_tokenizer, mpu
from megatron.checkpointing import load_checkpoint
from megatron.initialize import initialize_megatron
from megatron.model import GPTModel
from megatron.model.module import MegatronModule
from megatron.model.utils import print_named_parameters
from megatron.text_generation.api import generate_and_post_process
from megatron.text_generation.inference_params import InferenceParams
from megatron.text_generation_server import (MegatronServer,
                                             add_text_generate_args,
                                             setup_model)
from megatron.training import get_model

MAX_BATCH_SIZE = 1  # You can increase this depending on your desired max sequence length and GPU memory
MAX_SEQLEN = 16 * 1024


def model_provider(
        pre_process: bool=True,
        post_process: bool=True
    ) -> MegatronModule:
    """Build the model."""

    args = get_args()
    if args.model_architecture == "GPTModel":
        model : MegatronModule = GPTModel(num_tokentypes=0, parallel_output=False, pre_process=pre_process, post_process=post_process)
    else:
        raise ValueError(f"Unsupported model type: {args.model_architecture}")
    print_named_parameters(model)
    return model


def initialize_model_from_args() -> Tuple[Any, Optional[InferenceParams]]:
    # Needed for tensor parallel inference with CUDA graph
    os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"
    os.environ["NCCL_GRAPH_MIXING_SUPPORT"] = "0"

    initialize_megatron(
        extra_args_provider=add_text_generate_args,
        args_defaults={
            "tokenizer_type": "GPT2BPETokenizer",
            "no_load_rng": True,
            "no_load_optim": True,
            "inference_max_batch_size": MAX_BATCH_SIZE,
            "inference_max_seqlen": MAX_SEQLEN,
        },
    )

    args = get_args()
    if not args.fused_ft_kernel:
        args.use_cuda_graph = False  # CUDA graph requires fused FT kernel
    if hasattr(args, "iteration"):
        args.curr_iteration = args.iteration
        print("curr_iteration", args.curr_iteration)
    if args.num_layers_per_virtual_pipeline_stage is not None:
        print("Interleaved pipeline schedule is not yet supported for text generation.")
        exit()
    # Set up model and load checkpoint
    model = get_model(model_provider, wrap_with_ddp=False)

    if args.load is not None:
        _ = load_checkpoint(model, None, None)

    assert len(model) == 1, "Above condition should have caught this"
    model = model[0]

    args.model_architecture = "GPTModel"
    inference_params = setup_model(
        model,
        args.model_architecture,
        args.use_inference_kv_cache,
        args.fused_ft_kernel,
        args.use_cuda_graph,
        args.inference_max_batch_size,
        args.inference_max_seqlen,
    )

    return model, inference_params


if __name__ == "__main__":
    model, inference_params = initialize_model_from_args()

    args = get_args()
    tokenizer = get_tokenizer()
    if hasattr(args, "eos_id"):
        termination_id = args.eos_id
    else:
        termination_id = tokenizer.eod
    if mpu.is_pipeline_first_stage() and mpu.get_tensor_model_parallel_rank() == 0:
        assert inference_params is not None
        server = MegatronServer(
            model=model,
            inference_params=inference_params,
            params_dtype=args.params_dtype,
            max_position_embeddings=args.max_position_embeddings,
            termination_id=termination_id,
            tokenizer=tokenizer,
            port=args.port,
            )
        server.run("0.0.0.0")

    while True:
        if inference_params is not None:
            inference_params.reset()
        choice = torch.cuda.LongTensor(1)
        torch.distributed.broadcast(choice, 0)
        if choice[0].item() == 0:
            try:
                assert inference_params is not None
                generate_and_post_process(
                    model=model,
                    params_dtype=args.params_dtype,
                    max_position_embeddings=args.max_position_embeddings,
                    termination_id=termination_id,
                    tokenizer=tokenizer,
                    inference_params=inference_params)
            except ValueError as ve:
                pass

run_text_generation_server.py
persimmon-ai-labs/adept-inference
"""Utils for megatron tool use."""

from enum import Enum


class ModelBackend(Enum):
    """Model backend."""

    MEGATRON = "megatron"

tool_use/megatron/megatron_utils.py
persimmon-ai-labs/adept-inference