nemo/collections/llm/__init__.py

# Copyright (c) 2025, 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 is here to import it once, which improves the speed of launch when in debug-mode
from nemo.utils.import_utils import safe_import

safe_import("transformer_engine")

from nemo.collections.llm import peft
from nemo.collections.llm.bert.data import BERTMockDataModule, BERTPreTrainingDataModule, SpecterDataModule
from nemo.collections.llm.bert.model import (
    BertConfig,
    BertEmbeddingLargeConfig,
    BertEmbeddingMiniConfig,
    BertEmbeddingModel,
    BertModel,
    HuggingFaceBertBaseConfig,
    HuggingFaceBertConfig,
    HuggingFaceBertLargeConfig,
    HuggingFaceBertModel,
    MegatronBertBaseConfig,
    MegatronBertConfig,
    MegatronBertLargeConfig,
)
from nemo.collections.llm.gpt.data import (  # noqa: F401
    AlpacaDataModule,
    ChatDataModule,
    CustomReRankerDataModule,
    CustomRetrievalDataModule,
    DollyDataModule,
    FineTuningDataModule,
    HFDatasetDataModule,
    HFDatasetDataModulePacked,
    HFMockDataModule,
    MockDataModule,
    PreTrainingDataModule,
    SpecterReRankerDataModule,
    SquadDataModule,
)
from nemo.collections.llm.gpt.data.api import dolly, hf_dataset, mock, squad
from nemo.collections.llm.gpt.model import (  # noqa: F401
    Baichuan2Config,
    Baichuan2Config7B,
    Baichuan2Model,
    BaseMambaConfig1_3B,
    BaseMambaConfig2_7B,
    BaseMambaConfig130M,
    BaseMambaConfig370M,
    BaseMambaConfig780M,
    ChatGLM2Config6B,
    ChatGLM3Config6B,
    ChatGLMConfig,
    ChatGLMModel,
    CodeGemmaConfig2B,
    CodeGemmaConfig7B,
    CodeLlamaConfig7B,
    CodeLlamaConfig13B,
    CodeLlamaConfig34B,
    CodeLlamaConfig70B,
    DeepSeekModel,
    DeepSeekV2Config,
    DeepSeekV2LiteConfig,
    DeepSeekV3Config,
    Gemma2Config,
    Gemma2Config2B,
    Gemma2Config9B,
    Gemma2Config27B,
    Gemma2Model,
    Gemma3Config1B,
    Gemma3Config4B,
    Gemma3Config12B,
    Gemma3Config27B,
    Gemma3Model,
    GemmaConfig,
    GemmaConfig2B,
    GemmaConfig7B,
    GemmaModel,
    GPTConfig,
    GPTConfig5B,
    GPTConfig7B,
    GPTConfig20B,
    GPTConfig40B,
    GPTConfig126M,
    GPTConfig175B,
    GPTModel,
    GPTOSSConfig,
    GPTOSSConfig20B,
    GPTOSSConfig120B,
    GPTOSSModel,
    Hyena1bConfig,
    Hyena7bARCLongContextConfig,
    Hyena7bConfig,
    Hyena40bARCLongContextConfig,
    Hyena40bConfig,
    HyenaConfig,
    HyenaModel,
    HyenaNV1bConfig,
    HyenaNV7bConfig,
    HyenaNV40bConfig,
    HyenaNVTestConfig,
    HyenaTestConfig,
    Llama2Config7B,
    Llama2Config13B,
    Llama2Config70B,
    Llama3Config8B,
    Llama3Config70B,
    Llama4Config,
    Llama4Experts16Config,
    Llama4Experts128Config,
    Llama31Config8B,
    Llama31Config70B,
    Llama31Config405B,
    Llama31Nemotron70BConfig,
    Llama31NemotronNano8BConfig,
    Llama31NemotronUltra253BConfig,
    Llama32Config1B,
    Llama32Config3B,
    Llama32EmbeddingConfig1B,
    Llama32EmbeddingConfig3B,
    Llama32Reranker1BConfig,
    Llama32Reranker500MConfig,
    Llama33NemotronSuper49BConfig,
    LlamaConfig,
    LlamaEmbeddingModel,
    LlamaModel,
    LlamaNemotronModel,
    MambaModel,
    MaskedTokenLossReduction,
    MistralConfig7B,
    MistralModel,
    MistralNeMoConfig12B,
    MistralSmall3Config24B,
    MixtralConfig,
    MixtralConfig8x3B,
    MixtralConfig8x7B,
    MixtralConfig8x22B,
    MixtralModel,
    Nemotron3Config4B,
    Nemotron3Config8B,
    Nemotron3Config22B,
    Nemotron4Config15B,
    Nemotron4Config340B,
    NemotronConfig,
    NemotronHConfig4B,
    NemotronHConfig8B,
    NemotronHConfig47B,
    NemotronHConfig56B,
    NemotronModel,
    NemotronNano9Bv2,
    NemotronNano12Bv2,
    NVIDIAMambaConfig8B,
    NVIDIAMambaHybridConfig8B,
    Phi3Config,
    Phi3ConfigMini,
    Phi3Model,
    Qwen2Config,
    Qwen2Config1P5B,
    Qwen2Config7B,
    Qwen2Config72B,
    Qwen2Config500M,
    Qwen2Model,
    Qwen3Config,
    Qwen3Config1P7B,
    Qwen3Config4B,
    Qwen3Config8B,
    Qwen3Config14B,
    Qwen3Config30B_A3B,
    Qwen3Config32B,
    Qwen3Config235B_A22B,
    Qwen3Config600M,
    Qwen3Model,
    Qwen25Config1P5B,
    Qwen25Config3B,
    Qwen25Config7B,
    Qwen25Config14B,
    Qwen25Config32B,
    Qwen25Config72B,
    Qwen25Config500M,
    ReRankerModel,
    SSMConfig,
    Starcoder2Config,
    Starcoder2Config3B,
    Starcoder2Config7B,
    Starcoder2Config15B,
    Starcoder2Model,
    StarcoderConfig,
    StarcoderConfig15B,
    StarcoderModel,
    gpt_data_step,
    gpt_forward_step,
)
from nemo.collections.llm.t5.data import FineTuningDataModule as T5FineTuningDataModule
from nemo.collections.llm.t5.data import MockDataModule as T5MockDataModule
from nemo.collections.llm.t5.data import PreTrainingDataModule as T5PreTrainingDataModule
from nemo.collections.llm.t5.data import SquadDataModule as T5SquadDataModule
from nemo.collections.llm.t5.model import (
    T5Config,
    T5Config3B,
    T5Config11B,
    T5Config220M,
    T5Model,
    t5_data_step,
    t5_forward_step,
)

__all__ = [
    "MockDataModule",
    "T5MockDataModule",
    "CustomRetrievalDataModule",
    "CustomReRankerDataModule",
    "SpecterReRankerDataModule",
    "GPTModel",
    "GPTConfig",
    "HyenaTestConfig",
    "Hyena7bConfig",
    "Hyena40bConfig",
    "Hyena7bARCLongContextConfig",
    "Hyena40bARCLongContextConfig",
    "HyenaNVTestConfig",
    "HyenaNV40bConfig",
    "HyenaNV7bConfig",
    "HyenaConfig",
    "HyenaModel",
    "Hyena1bConfig",
    "HyenaNV1bConfig",
    "gpt_data_step",
    "gpt_forward_step",
    "T5Model",
    "T5Config",
    "T5Config220M",
    "T5Config3B",
    "T5Config11B",
    "BertConfig",
    "BertEmbeddingModel",
    "BertModel",
    "BertEmbeddingLargeConfig",
    "BertEmbeddingMiniConfig",
    "t5_data_step",
    "t5_forward_step",
    "MaskedTokenLossReduction",
    "MistralConfig7B",
    "MistralNeMoConfig12B",
    "MistralSmall3Config24B",
    "MistralModel",
    "MixtralConfig",
    "MixtralConfig8x3B",
    "MixtralConfig8x7B",
    "MixtralConfig8x22B",
    "MixtralModel",
    "Starcoder2Config15B",
    "Starcoder2Config",
    "Starcoder2Model",
    "NemotronModel",
    "Nemotron3Config4B",
    "Nemotron3Config8B",
    "Nemotron3Config22B",
    "Nemotron4Config15B",
    "Nemotron4Config340B",
    "NemotronConfig",
    "LlamaEmbeddingModel",
    "Llama32EmbeddingConfig1B",
    "Llama32EmbeddingConfig3B",
    "Phi3Config",
    "Phi3ConfigMini",
    "Phi3Model",
    "SSMConfig",
    "BaseMambaConfig130M",
    "BaseMambaConfig370M",
    "BaseMambaConfig780M",
    "BaseMambaConfig1_3B",
    "BaseMambaConfig2_7B",
    "NVIDIAMambaConfig8B",
    "NVIDIAMambaHybridConfig8B",
    "NemotronHConfig4B",
    "NemotronHConfig8B",
    "NemotronHConfig47B",
    "NemotronHConfig56B",
    "NemotronNano9Bv2",
    "NemotronNano12Bv2",
    "MambaModel",
    "LlamaConfig",
    "Llama2Config7B",
    "Llama2Config13B",
    "Llama2Config70B",
    "Llama3Config8B",
    "Llama3Config70B",
    "Llama31Config8B",
    "Llama31Config70B",
    "Llama31Config405B",
    "Llama32Config1B",
    "Llama32Config3B",
    "Llama4Experts16Config",
    "Llama4Experts128Config",
    "Llama4Config",
    "Llama31NemotronNano8BConfig",
    "Llama31Nemotron70BConfig",
    "Llama33NemotronSuper49BConfig",
    "Llama31NemotronUltra253BConfig",
    "Llama32Reranker500MConfig",
    "Llama32Reranker1BConfig",
    "CodeLlamaConfig7B",
    "CodeLlamaConfig13B",
    "CodeLlamaConfig34B",
    "CodeLlamaConfig70B",
    "LlamaModel",
    "LlamaNemotronModel",
    "GPTOSSConfig",
    "GPTOSSConfig120B",
    "GPTOSSConfig20B",
    "GPTOSSModel",
    "GemmaConfig",
    "GemmaConfig2B",
    "GemmaConfig7B",
    "CodeGemmaConfig2B",
    "CodeGemmaConfig7B",
    "GemmaModel",
    "Gemma2Model",
    "Gemma2Config9B",
    "Gemma2Config",
    "Gemma2Config27B",
    "Gemma2Config2B",
    "Gemma3Model",
    "Gemma3Config1B",
    "Gemma3Config4B",
    "Gemma3Config12B",
    "Gemma3Config27B",
    "Baichuan2Config",
    "Baichuan2Config7B",
    "Baichuan2Model",
    "ChatGLMConfig",
    "ChatGLM2Config6B",
    "ChatGLM3Config6B",
    "ChatGLMModel",
    "Qwen2Model",
    "Qwen2Config7B",
    "Qwen2Config",
    "Qwen2Config500M",
    "Qwen2Config1P5B",
    "Qwen25Config3B",
    "Qwen2Config72B",
    "Qwen25Config500M",
    "Qwen25Config1P5B",
    "Qwen25Config7B",
    "Qwen25Config14B",
    "Qwen25Config32B",
    "Qwen25Config72B",
    "Qwen3Config",
    "Qwen3Config600M",
    "Qwen3Config1P7B",
    "Qwen3Config4B",
    "Qwen3Config8B",
    "Qwen3Config14B",
    "Qwen3Config32B",
    "Qwen3Config30B_A3B",
    "Qwen3Config235B_A22B",
    "Qwen3Model",
    "PreTrainingDataModule",
    "FineTuningDataModule",
    "ChatDataModule",
    "SquadDataModule",
    "T5PreTrainingDataModule",
    "T5FineTuningDataModule",
    "T5SquadDataModule",
    "T5MockDataModule",
    "DeepSeekModel",
    "DeepSeekV2Config",
    "DeepSeekV2LiteConfig",
    "DeepSeekV3Config",
    "HuggingFaceBertBaseConfig",
    "HuggingFaceBertConfig",
    "HuggingFaceBertLargeConfig",
    "HuggingFaceBertModel",
    "MegatronBertBaseConfig",
    "MegatronBertConfig",
    "MegatronBertLargeConfig",
    "BERTMockDataModule",
    "BERTPreTrainingDataModule",
    "SpecterDataModule",
    "DollyDataModule",
    "tokenizer",
    "mock",
    "squad",
    "dolly",
    "peft",
    "hf_dataset",
    "HFMockDataModule",
]


from nemo.utils import logging

try:
    import nemo_run as run  # noqa: F401

    from nemo.collections.llm.api import (  # noqa: F401
        distill,
        export_ckpt,
        finetune,
        generate,
        import_ckpt,
        pretrain,
        prune,
        ptq,
        train,
        validate,
    )
    from nemo.collections.llm.recipes import *  # noqa

    __all__.extend(
        [
            "train",
            "import_ckpt",
            "export_ckpt",
            "pretrain",
            "validate",
            "finetune",
            "generate",
            "prune",
            "ptq",
            "distill",
        ]
    )
except ImportError as error:
    logging.warning(f"Failed to import nemo.collections.llm.[api,recipes]: {error}")

try:
    from nemo.collections.llm.api import deploy  # noqa: F401

    __all__.append("deploy")
except ImportError as error:
    logging.warning(f"The deploy module could not be imported: {error}")

try:
    from nemo.collections.llm.api import evaluate  # noqa: F401

    __all__.append("evaluate")
except ImportError as error:
    logging.warning(f"The evaluate module could not be imported: {error}")


import warnings


warnings.warn(
    "nemo.collections.llm is deprecated and will be removed in a future major NeMo FW container release. "
    "Please refer to the new Megatron-Bridge repository: https://github.com/NVIDIA-NeMo/Megatron-Bridge",
    DeprecationWarning,
    stacklevel=2,
)