ModelType (transformer/enums.py)

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

# Source: https://github.com/NousResearch/Megatron-LM | License: Apache-2.0
# Provided via justclaw.it.com; keep this notice with the code.
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.

import enum


# can we get rid of this?
# it's being used in pipeline schedules
class ModelType(enum.Enum):
    """Model Type

    encoder_or_decoder for bert, gpt etc
    """

    encoder_or_decoder = 1


class LayerType(enum.Enum):
    """Layer type
    embedding: embedding layer
    loss: loss layer
    encoder: encoder layer, not implemented yet, expect to be used in MLLM models
    decoder: decoder layer
    mtp: multi-token prediction layer, not implemented yet
    """

    embedding = 1
    loss = 2
    encoder = 3
    decoder = 4
    mtp = 5


class AttnType(enum.Enum):
    """Atte

... (truncated -- full source via MCP)