nemo/collections/llm/api.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.
import json
import warnings
from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Optional, Union

import lightning.pytorch as pl
import nemo_run as run
import torch
from megatron.core import parallel_state
from rich.console import Console
from torch.distributed import all_gather_object
from typing_extensions import Annotated

import nemo.lightning as nl
from nemo.collections.llm import GPTModel
from nemo.collections.llm.gpt.data.fine_tuning import FineTuningDataModule
from nemo.collections.llm.modelopt import (
    DistillationGPTModel,
    ExportConfig,
    PruningConfig,
    QuantizationConfig,
    Quantizer,
    prune_language_model,
    save_pruned_model,
    set_modelopt_spec_if_exists_in_ckpt,
    setup_trainer_and_restore_model_with_modelopt_spec,
)
from nemo.lightning import (
    AutoResume,
    NeMoLogger,
    OptimizerModule,
    Trainer,
    configure_no_restart_validation_training_loop,
    io,
)
from nemo.lightning.base import NEMO_MODELS_CACHE
from nemo.lightning.callback_group import CallbackGroup
from nemo.lightning.ckpt_utils import ckpt_to_context_subdir
from nemo.lightning.pytorch.callbacks import PEFT, JitTransform, ModelTransform
from nemo.utils import logging
from nemo.utils.get_rank import is_global_rank_zero

if TYPE_CHECKING:
    from megatron.core.inference.common_inference_params import CommonInferenceParams
    from megatron.core.inference.inference_request import InferenceRequest


TokenizerType = Any
AnyPath = Union[Path, str]


@run.cli.entrypoint(namespace="llm")
def train(
    model: Union[pl.LightningModule, AnyPath],
    data: pl.LightningDataModule,
    trainer: Trainer,
    log: Annotated[Optional[NeMoLogger], run.Config[NeMoLogger]] = None,
    resume: Annotated[Optional[AutoResume], run.Config[AutoResume]] = None,
    optim: Optional[OptimizerModule] = None,
    tokenizer: Optional[TokenizerType] = None,
    model_transform: Optional[Union[PEFT, ModelTransform, Callable]] = None,
    # TODO: Fix export export: Optional[str] = None,
) -> Path:
    """
    Trains a model using the specified data and trainer, with optional tokenizer, source, and export.

    Args:
        model (Union[pl.LightningModule, AnyPath]): The model to be trained or a path to the NeMo 2 checkpoint.
        data (pl.LightningDataModule): The data module containing training data.
        trainer (Trainer): The trainer instance configured with a MegatronStrategy.
        log (NeMoLogger): A nemologger instance.
        resume (Optional[Union[AutoResume, Resume]]): Resume training from a checkpoint.
        optim (Optional[OptimizerModule]): The optimizer module to be used. If not provided, the default optimizer
            from the model will be used.
        tokenizer (Optional[TokenizerType]): Tokenizer setting to be applied. Can be 'data' or 'model'
            or an instance of TokenizerSpec.
        export (Optional[str]): Filename to save the exported checkpoint after training.
        model_transform (Optional[Union[Callable[[nn.Module], nn.Module], PEFT]]): A model transform to be applied.

    Returns
    -------
        Path: The directory path where training artifacts are saved.

    Examples
    --------
        >>> from nemo.collections import llm
        >>> from nemo import lightning as nl
        >>> model = llm.MistralModel()
        >>> data = llm.SquadDataModule(seq_length=4096, global_batch_size=16, micro_batch_size=2)
        >>> precision = nl.MegatronMixedPrecision(precision="bf16-mixed")
        >>> trainer = nl.Trainer(strategy=nl.MegatronStrategy(tensor_model_parallel_size=2), plugins=precision)
        >>> llm.train(model, data, trainer, tokenizer="data")
        PosixPath('/path/to/log_dir')
    """
    model = _load_model_from_path(model)

    # [ModelOpt]: If modelopt_state exists, overwrite transformer_layer_spec to modelopt spec
    if resume:
        if resume.restore_config and resume.restore_config.path:
            set_modelopt_spec_if_exists_in_ckpt(model, resume.restore_config.path)
        elif resume.resume_from_path:
            set_modelopt_spec_if_exists_in_ckpt(model, resume.resume_from_path)

    app_state = _setup(
        model=model,
        data=data,
        trainer=trainer,
        log=log,
        resume=resume,
        optim=optim,
        tokenizer=tokenizer,
        model_transform=model_transform,
    )

    trainer.fit(model, data)

    return app_state.exp_dir


@run.cli.entrypoint(namespace="llm")
def pretrain(
    model: Union[pl.LightningModule, AnyPath],
    data: pl.LightningDataModule,
    trainer: Trainer,
    log: Annotated[Optional[NeMoLogger], run.Config[NeMoLogger]] = None,
    resume: Annotated[Optional[AutoResume], run.Config[AutoResume]] = None,
    optim: Optional[OptimizerModule] = None,
) -> Path:
    """
    Pretrains a model using the specified data and trainer, with optional logging, resuming, and optimization.

    This function is a wrapper around the `train` function, specifically configured for pretraining tasks.
    Note, by default it will use the tokenizer from the model.

    Args:
        model (Union[pl.LightningModule, AnyPath]): The model to be pretrained or a path to the NeMo 2 checkpoint.
        data (pl.LightningDataModule): The data module containing pretraining data.
        trainer (Trainer): The trainer instance configured with a MegatronStrategy.
        log (NeMoLogger): A nemologger instance.
        resume (Optional[AutoResume]): Resume training from a checkpoint.
        optim (Optional[OptimizerModule]): The optimizer module to be used. If not provided, the default
            optimizer from the model will be used.

    Returns:
        Path: The directory path where pretraining artifacts are saved.

    Examples:
        >>> from nemo.collections import llm
        >>> from nemo import lightning as nl
        >>> model = llm.MistralModel()
        >>> data = llm.PretrainingDataModule(paths=[...], seq_length=4096, global_batch_size=16, micro_batch_size=2)
        >>> precision = nl.MegatronMixedPrecision(precision="bf16-mixed")
        >>> trainer = nl.Trainer(strategy=nl.MegatronStrategy(tensor_model_parallel_size=2), plugins=precision)
        >>> llm.pretrain(model, data, trainer)
        PosixPath('/path/to/log_dir')
    """
    model = _load_model_from_path(model)
    _validate_config(model, data, trainer, log=log, resume=resume, optim=optim)

    return train(
        model=model,
        data=data,
        trainer=trainer,
        log=log,
        resume=resume,
        optim=optim,
        tokenizer="data",
    )


@run.cli.entrypoint(namespace="llm")
def finetune(
    model: Union[pl.LightningModule, AnyPath],
    data: pl.LightningDataModule,
    trainer: Trainer,
    log: Annotated[Optional[NeMoLogger], run.Config[NeMoLogger]] = None,
    resume: Annotated[Optional[AutoResume], run.Config[AutoResume]] = None,
    optim: Optional[OptimizerModule] = None,
    peft: Optional[Union[PEFT, ModelTransform, Callable]] = None,
    tokenizer: Optional[TokenizerType] = "model",
) -> Path:
    """
    Finetunes a model using the specified data and trainer, with optional logging, resuming, and PEFT.

    Note, by default it will use the tokenizer from the model.

    Args:
        model (Union[pl.LightningModule, AnyPath]): The model to be finetuned.
        data (pl.LightningDataModule): The data module containing finetuning data.
        trainer (Trainer): The trainer instance configured with a MegatronStrategy.
        log (NeMoLogger): A nemologger instance.
        resume (Optional[AutoResume]): Resume training from a checkpoint.
        optim (Optional[OptimizerModule]): The optimizer module to be used. If not provided, the default
            optimizer from the model will be used.
        peft (Optional[PEFT]): A PEFT (Parameter-Efficient Fine-Tuning) configuration to be applied.
        tokenizer (Optional[TokenizerType]): Tokenizer setting to be applied. Can be 'data' or 'model'
            or an instance of TokenizerSpec. If 'data' uses the data loader's tokenizer instead of the tokenizer
            from the model checkpoint, which is useful for expanding vocabulary or adding special tokens
            (such as chat template tokens).

    Returns:
        Path: The directory path where finetuning artifacts are saved.

    Examples:
        >>> from nemo.collections import llm
        >>> from nemo import lightning as nl
        >>> model = llm.MistralModel()
        >>> data = llm.SquadDataModule(seq_length=4096, global_batch_size=16, micro_batch_size=2)
        >>> precision = nl.MegatronMixedPrecision(precision="bf16-mixed")
        >>> trainer = nl.Trainer(strategy=nl.MegatronStrategy(tensor_model_parallel_size=2), plugins=precision)
        >>> llm.finetune(model, data, trainer, peft=llm.peft.LoRA()])
        PosixPath('/path/to/log_dir')
    """
    model = _load_model_from_path(model)
    _validate_config(model, data, trainer, log=log, resume=resume, optim=optim, model_transform=peft)
    return train(
        model=model,
        data=data,
        trainer=trainer,
        log=log,
        resume=resume,
        optim=optim,
        tokenizer=tokenizer,
        model_transform=peft,
    )


@run.cli.entrypoint(namespace="llm")
def validate(
    model: pl.LightningModule,
    data: pl.LightningDataModule,
    trainer: Trainer,
    log: Annotated[Optional[NeMoLogger], run.Config[NeMoLogger]] = None,
    resume: Annotated[Optional[AutoResume], run.Config[AutoResume]] = None,
    optim: Optional[OptimizerModule] = None,
    tokenizer: Optional[TokenizerType] = None,
    model_transform: Optional[Union[PEFT, ModelTransform, Callable]] = None,
) -> Path:
    """
    Validates a model using the specified data and trainer, with optional logging, resuming, and model transformations.

    Args:
        model (pl.LightningModule): The model to be validated.
        data (pl.LightningDataModule): The data module containing validation data.
        trainer (Trainer): The trainer instance configured with a MegatronStrategy.
        log (NeMoLogger): A nemologger instance.
        resume (Optional[AutoResume]): Resume from a checkpoint for validation.
        optim (Optional[OptimizerModule]): The optimizer module to be used. If not provided, the default optimizer
            from the model will be used.
        tokenizer (Optional[TokenizerType]): Tokenizer setting to be applied. Can be 'data' or 'model'
            or an instance of TokenizerSpec.
        model_transform (Optional[Union[Callable[[nn.Module], nn.Module], PEFT]]): A model transform to be applied.

    Returns:
        Path: The directory path where validation artifacts are saved.

    Examples:
        >>> from nemo.collections import llm
        >>> from nemo import lightning as nl
        >>> model = llm.MistralModel()
        >>> data = llm.SquadDataModule(seq_length=4096, global_batch_size=16, micro_batch_size=2)
        >>> precision = nl.MegatronMixedPrecision(precision="bf16-mixed")
        >>> trainer = nl.Trainer(strategy=nl.MegatronStrategy(tensor_model_parallel_size=2), plugins=precision)
        >>> llm.validate(model, data, trainer, tokenizer="data")
        PosixPath('/path/to/log_dir')
    """
    app_state = _setup(
        model=model,
        data=data,
        trainer=trainer,
        log=log,
        resume=resume,
        optim=optim,
        tokenizer=tokenizer,
        model_transform=model_transform,
    )

    trainer.validate(model, data)

    return app_state.exp_dir


@run.cli.entrypoint(name="prune", namespace="llm")
def prune(
    nemo_checkpoint: str,
    save_path: str,
    pruning_config: PruningConfig,
    devices: int = 1,
    num_nodes: int = 1,
    tp_size: int = 1,
    pp_size: int = 1,
    num_layers_in_first_pipeline_stage: int | None = None,
    num_layers_in_last_pipeline_stage: int | None = None,
    num_train_samples: int = 1024,
    data: pl.LightningDataModule | None = None,
    tokenizer_path: str | None = None,
    legacy_ckpt: bool = False,
) -> str:
    """
    Prunes a model using the specified data and trainer. Currently only supports GPT models.

    Args:
        nemo_checkpoint (str): The path to the NeMo checkpoint to be pruned.
        save_path (str): The path to save the pruned NeMo checkpoint.
        pruning_config (PruningConfig): The pruning configuration.
        devices (int): The number of devices to use for pruning.
        num_nodes (int): The number of nodes to use for pruning.
        tp_size (int): The tensor parallel size.
        pp_size (int): The pipeline parallel size.
        num_train_samples (int): Number of training samples for importance estimation using forward pass.
        num_layers_in_first_pipeline_stage (int): The number of layers in the first pipeline stage.
        num_layers_in_last_pipeline_stage (int): The number of layers in the last pipeline stage.
        data (pl.LightningDataModule): The data module for forward pass.
            Required if not dropping layers.
        tokenizer_path (str): Path to the tokenizer if not using model's tokenizer.
        legacy_ckpt (bool): If True, allow loading ckpt saved with older version of TE.
            Use for cases like missing state dict keys ending with `_extra_state`.

    Returns:
        str: The path to the pruned NeMo checkpoint.

    Examples:
        >>> from nemo.collections import llm
        >>> from nemo.collections.llm.modelopt.prune import PruningConfig
        >>> data = llm.PretrainingDataModule(
                paths=["1.0", "path/to/tokenized/data"],
                seq_length=256,
                global_batch_size=1,
                micro_batch_size=1,
            )
        >>> llm.prune(
                nemo_checkpoint="path/to/llama3.1-8b",
                save_path="path/to/pruned_llama_model",
                pruning_config=PruningConfig(target_ffn_hidden_size=9216, target_hidden_size=3072),
                data=data
            )
    """
    if data is not None:
        assert data.global_batch_size == data.micro_batch_size, "Global batch size must be equal to micro batch size"
        steps = num_train_samples // data.global_batch_size
    else:
        steps = num_train_samples

    model, trainer = setup_trainer_and_restore_model_with_modelopt_spec(
        model_path=nemo_checkpoint,
        tensor_model_parallel_size=tp_size,
        pipeline_model_parallel_size=pp_size,
        num_layers_in_first_pipeline_stage=num_layers_in_first_pipeline_stage,
        num_layers_in_last_pipeline_stage=num_layers_in_last_pipeline_stage,
        devices=devices,
        num_nodes=num_nodes,
        inference_only=True,
        tokenizer_path=tokenizer_path,
        legacy_ckpt=legacy_ckpt,
        strategy_kwargs={"sequence_parallel": False, "replace_progress_bar": False},
        trainer_kwargs={"max_steps": steps, "limit_val_batches": steps, "val_check_interval": steps},
        model_config_overrides={"sequence_parallel": False},
    )
    prune_language_model(model, pruning_config, data, trainer)
    save_pruned_model(trainer, save_path)

    console = Console()
    console.print(f"[green]✓ Pruning succeded, pruned checkpoint saved to {save_path}[/green]")

    return save_path


@run.cli.entrypoint(name="distill", namespace="llm")
def distill(
    student_model_path: AnyPath,
    teacher_model_path: AnyPath,
    data: pl.LightningDataModule,
    trainer: Trainer,
    distillation_config_path: Optional[AnyPath] = None,
    log: Annotated[Optional[NeMoLogger], run.Config[NeMoLogger]] = None,
    resume: Annotated[Optional[AutoResume], run.Config[AutoResume]] = None,
    optim: Optional[OptimizerModule] = None,
    tokenizer: Optional[TokenizerType] = None,
    model_transform: Optional[Union[PEFT, ModelTransform, Callable]] = None,
) -> Path:
    """
    Distills a teacher model into a student model using special Knowledge-Distillation losses.

    Note that this requires an existing NeMo 2.0 checkpoint of the student model as well, as
    the model class is not known beforehand.
    This script currently supports instances of ``nemo.collections.llm.GPTModel`` for now.

    Args:
        student_model_path (Path): Path to student model NeMo checkpoint to be trained.
        teacher_model_path (Path): Path to teacher model NeMo checkpoint to distill from.
        data (pl.LightningDataModule): The data module containing training data.
        trainer (Trainer): The trainer instance configured with a MegatronStrategy.
        distillation_config_path (Optional[Path]): Path to distillation config YAML file.
            If not provided, by default will perform logits-only distillation.
        log (NeMoLogger): A nemologger instance.
        resume (Optional[Union[AutoResume, Resume]]): Resume training from a checkpoint.
        optim (Optional[OptimizerModule]): The optimizer module to be used. If not provided, the default optimizer
            from the model will be used.
        tokenizer (Optional[TokenizerType]): Tokenizer setting to be applied. Can be 'data' or 'model'
            or an instance of TokenizerSpec.
        export (Optional[str]): Filename to save the exported checkpoint after training.
        model_transform (Optional[Union[Callable[[nn.Module], nn.Module], PEFT]]): A model transform to be applied.

    Returns
    -------
        Path: The directory path where training artifacts are saved.

    Examples
    --------
        >>> from nemo.collections import llm
        >>> from nemo import lightning as nl
        >>> student = "/path/to/student/nemo/ckpt"  # <-- change me
        >>> teacher = "/path/to/teacher/nemo/ckpt"  # <-- change me
        >>> data = llm.SquadDataModule(seq_length=4096, global_batch_size=16, micro_batch_size=2)
        >>> precision = nl.MegatronMixedPrecision(precision="bf16-mixed")
        >>> trainer = nl.Trainer(strategy=nl.MegatronStrategy(tensor_model_parallel_size=2), plugins=precision)
        >>> llm.distill(student, teacher, data, trainer, tokenizer="model")
        PosixPath('/path/to/log_dir')
    """
    _student_model = io.load_context(ckpt_to_context_subdir(student_model_path), subpath="model")
    _teacher_model = io.load_context(ckpt_to_context_subdir(teacher_model_path), subpath="model")
    assert isinstance(_student_model, GPTModel), "Only models based on `llm.GPTModel` are supported currently."
    assert isinstance(_teacher_model, GPTModel), "Only models based on `llm.GPTModel` are supported currently."

    if tokenizer is None:
        tokenizer = getattr(_student_model, "tokenizer", None) or getattr(_teacher_model, "tokenizer", None)
        assert tokenizer is not None, "Tokenizer neither provided nor found in models."

    model = DistillationGPTModel(
        _student_model.config,
        _teacher_model.config,
        teacher_ckpt_path=teacher_model_path,
        distillation_config_path=distillation_config_path,
    )
    model.__io__ = _student_model.__io__

    if resume is None:
        resume = AutoResume()
    if resume.restore_config is None:
        resume.restore_config = nl.RestoreConfig(path=student_model_path)

    return train(
        model=model,
        data=data,
        optim=optim,
        tokenizer=tokenizer,
        trainer=trainer,
        log=log,
        resume=resume,
        model_transform=model_transform,
    )


@run.cli.entrypoint(name="ptq", namespace="llm")
def ptq(
    model_path: str,
    export_config: ExportConfig,
    calibration_tp: int = 1,
    calibration_pp: int = 1,
    calibration_ep: int = 1,
    num_layers_in_first_pipeline_stage: int | None = None,
    num_layers_in_last_pipeline_stage: int | None = None,
    devices: int | None = None,
    num_nodes: int | None = None,
    quantization_config: Annotated[Optional[QuantizationConfig], run.Config[QuantizationConfig]] = None,
    forward_loop: Callable | None = None,
    tokenizer_path: str | None = None,
    legacy_ckpt: bool = False,
    trust_remote_code: bool = False,
) -> Path:
    """
    Applies Post-Training Quantization (PTQ) for a model using the specified quantization and export configs. It runs
    calibration for a small dataset to collect scaling factors low-precision GEMMs used by desired quantization method.
    By default, this function produces TensorRT-LLM checkpoint ready for deployment using the Export-Deploy repository
    (https://github.com/NVIDIA-NeMo/Export-Deploy) or directly using TensorRT-LLM library.

    The function can be used through the NeMo CLI in the following way:
    ```bash
    # Run calibration using tensor parallel set to 8 and export quantized checkpoint with tensor parallel equal 2
    nemo llm ptq run.executor=torchrun run.executor.ntasks_per_node=8 \
        model_path=/models/Llama-3-70B \
        export_config.path=/models/Llama-3-70B-FP8 \
        calibration_tp=8 \
        export_config.inference_tp=2

    # Choose different quantization method, for example, INT8 SmoothQuant
    nemo llm ptq run.executor=torchrun run.executor.ntasks_per_node=1 \
        model_path=/models/Llama-3-8B \
        export_config.path=/models/Llama-3-8B-INT8_SQ \
        quantization_config.algorithm=int8_sq

    # Export as NeMo checkpoint instead
    nemo llm ptq run.executor=torchrun \
        model_path=/models/Llama-3-8B \
        export_config.path=/models/Llama-3-8B-INT8_SQ \
        quantization_config.algorithm=int8_sq \
        export_config.export_format=nemo

    # Quantize HF AutoModel checkpoint.
    nemo llm ptq run.executor=torchrun run.executor.ntasks_per_node=1 \
        model_path=/models/Llama-3-70B-HF \
        export_config.path=/models/Llama-3-70B-HF-FP8 \
        export_config.export_format=hf
    ```

    Args:
        model_path (str): The path to model to be quantized.
        calibration_tp (int): Calibration tensor parallelism.
        calibration_pp (int): Calibration pipeline parallelism.
        num_layers_in_first_pipeline_stage (int): Number of layers in the first pipeline stage.
        num_layers_in_last_pipeline_stage (int): Number of layers in the last pipeline stage.
        export_config (ExportConfig): Export configuration for output checkpoint.
        devices (int): Number of devices to use for calibration. Default: calibration_tp.
        num_nodes (int): Number of nodes to use for calibration. Default: calibration_pp.
        quantization_config (QuantizationConfig): Configuration for quantization algorithm.
        forward_loop (Callable): Forward loop to use for calibration.
            If not provided, a forward loop will be created using the calibration dataset.
        tokenizer_path (str): Path to the tokenizer if not using model's tokenizer.
        legacy_ckpt (bool): If True, allow loading ckpt saved with older version of TE.
        trust_remote_code (bool): Trust remote code when loading HuggingFace models.

    Returns:
        Path: The path where the quantized checkpoint has been saved after calibration.
    """
    if not quantization_config:
        quantization_config = QuantizationConfig()
    if devices is None:
        devices = calibration_tp
    if num_nodes is None:
        num_nodes = calibration_pp

    quantizer = Quantizer(quantization_config, export_config)
    assert Path(model_path).exists(), f"Path {model_path} does not exist"

    trainer = None
    model, trainer = setup_trainer_and_restore_model_with_modelopt_spec(
        model_path=model_path,
        tensor_model_parallel_size=calibration_tp,
        pipeline_model_parallel_size=calibration_pp,
        num_layers_in_first_pipeline_stage=num_layers_in_first_pipeline_stage,
        num_layers_in_last_pipeline_stage=num_layers_in_last_pipeline_stage,
        expert_model_parallel_size=calibration_ep,
        devices=devices,
        num_nodes=num_nodes,
        inference_only=True,
        tokenizer_path=tokenizer_path,
        legacy_ckpt=legacy_ckpt,
        strategy_kwargs={"sequence_parallel": False, "lazy_init": True},
        trainer_kwargs={},
        model_config_overrides={"sequence_parallel": False},
    )

    model = quantizer.quantize(model, forward_loop)
    quantizer.export(model, model_path, trainer)

    if is_global_rank_zero():
        console = Console()
        console.print(f"[green]✓ PTQ succeded, quantized checkpoint exported to {export_config.path}[/green]")
    return export_config.path


@run.cli.entrypoint(name="import", namespace="llm")
def import_ckpt(
    model: pl.LightningModule,
    source: str,
    output_path: Optional[AnyPath] = None,
    overwrite: bool = False,
    **kwargs,
) -> Path:
    """
    Imports a checkpoint into a model using the model's associated importer, typically for
    the purpose of fine-tuning a community model trained in an external framework, such as
    Hugging Face.

    This function can be used both programmatically and through the NeMo CLI:

    CLI Usage:
    ```bash
    # Import Llama 3 8B from HuggingFace (saves to $NEMO_MODELS_CACHE)
    nemo llm import model=llama3_8b source="hf://meta-llama/Llama-3.1-8B"

    # Import with custom output path
    nemo llm import model=llama3_8b source="hf://meta-llama/Llama-3.1-8B" output_path="/path/to/save"

    # Force overwrite existing checkpoint
    nemo llm import model=llama3_8b source="hf://meta-llama/Llama-3.1-8B" overwrite=true
    ```

    Python Usage:
    ```python
    model = Mistral7BModel()
    imported_path = import_ckpt(model, "hf://mistralai/Mistral-7B-v0.1")
    ```

    The importer component of the model reads the checkpoint data from the specified source
    and transforms it into the right format. This is particularly useful for adapting
    models that have been pre-trained in different environments or frameworks to be fine-tuned
    or further developed within the current system.

    For instance, using `import_ckpt(Mistral7BModel(), "hf")` initiates the import process
    by searching for a registered model importer tagged with "hf". In NeMo, `HFMistral7BImporter`
    is registered under this tag via:
    `@io.model_importer(Mistral7BModel, "hf", default_path="mistralai/Mistral-7B-v0.1")`.
    This links `Mistral7BModel` to `HFMistral7BImporter`, designed for HuggingFace checkpoints.

    Args:
        model (pl.LightningModule): The model into which the checkpoint will be imported.
            This model must implement the ConnectorMixin.
        source (str): The source from which the checkpoint will be imported. This can be
            a file path, URL, or any other string identifier that the model's importer
            can recognize.
        output_path (Optional[Path]): The path where the imported checkpoint will be stored.
            If not specified, the checkpoint will be saved to $NEMO_MODELS_CACHE
            (defaults to ~/.cache/nemo/models/ if the environment variable is not set).
        overwrite (bool): If set to True, existing files at the output path will be overwritten.
            This is useful for model updates where retaining old checkpoint files is not required.

    Returns:
        Path: The path where the checkpoint has been saved after import.

    Raises:
        ValueError: If the model does not implement ConnectorMixin, indicating a lack of
            necessary importer functionality.
        FileExistsError: If the output path is provided (that is, when not using models cache)
            and it exists and overwrite is not set to True.
    """
    if output_path:
        output_path = Path(output_path)
        if output_path.exists() and not overwrite:
            raise FileExistsError(f"Output path {output_path} exists. Use overwrite=True to force overwrite.")

    output = io.import_ckpt(model=model, source=source, output_path=output_path, overwrite=overwrite, **kwargs)

    console = Console()
    if output_path:
        console.print(f"[green]✓ Checkpoint imported to {output}[/green]")
    else:
        console.print(f"[green] $NEMO_MODELS_CACHE={NEMO_MODELS_CACHE} [/green]")

    # Display directory structure as a tree
    dir_tree = _build_directory_tree(output, root_name="Imported Checkpoint")
    console.print(dir_tree)

    return output


def load_connector_from_trainer_ckpt(path: AnyPath, target: str) -> io.ModelConnector:
    # pylint: disable=C0116
    if not isinstance(path, Path):
        path = Path(path)
    return io.load_context(path, subpath="model").exporter(target, path)


@run.cli.entrypoint(name="export", namespace="llm")
def export_ckpt(
    path: AnyPath,
    target: str,
    output_path: Optional[AnyPath] = None,
    overwrite: bool = False,
    load_connector: Callable[[Path, str], io.ModelConnector] = load_connector_from_trainer_ckpt,
    modelopt_export_kwargs: dict[str, Any] = None,
    **kwargs,
) -> Path:
    """
    Exports a checkpoint from a model using the model's associated exporter, typically for
    the purpose of sharing a model that has been fine-tuned or customized within NeMo.

    This function can be used both programmatically and through the NeMo CLI:

    CLI Usage:
    ```bash
    # Export model to HuggingFace format (saves to {checkpoint_path}/hf/)
    nemo llm export path=/path/to/model.nemo target="hf"

    # Export with custom output path
    nemo llm export path=/path/to/model.nemo target="hf" output_path="/path/to/save"

    # Force overwrite existing export
    nemo llm export path=/path/to/model.nemo target="hf" overwrite=true
    ```

    Python Usage:
    ```python
    nemo_ckpt_path = Path("/path/to/model.nemo")
    export_path = export_ckpt(nemo_ckpt_path, "hf")
    ```

    The exporter component of the model reads the model's state from the specified path and
    exports it into the format specified by the 'target' identifier. This is particularly
    useful for adapting models that have been developed or fine-tuned within NeMo to be
    compatible with other environments or frameworks.

    Args:
        path (Path): The path to the model's checkpoint file from which data will be exported.
        target (str): The identifier for the exporter that defines the format of the export
            (e.g., "hf" for HuggingFace format).
        output_path (Optional[Path]): The path where the exported checkpoint will be saved.
            If not specified, defaults to {checkpoint_path}/{target}/.
        overwrite (bool): If set to True, existing files at the output path will be overwritten.
            This is useful for model updates where retaining old checkpoint files is not required.
        load_connector (Callable[[Path, str], ModelConnector]): A function to load the appropriate
            exporter based on the model and target format. Defaults to `load_connector_from_trainer_ckpt`.
        modelopt_export_kwargs (Dict[str, Any]): Additional keyword arguments for ModelOpt export to HuggingFace.

    Returns:
        Path: The path where the checkpoint has been saved after export.

    Raises:
        ValueError: If the model does not implement ConnectorMixin, indicating a lack of
            necessary exporter functionality.
        FileExistsError: If the output path is provided (that is, when not using models cache)
            and it exists and overwrite is not set to True.
    """
    if not isinstance(path, Path):
        path = Path(path)
    if output_path and not isinstance(output_path, Path):
        output_path = Path(output_path)
        if output_path.exists() and not overwrite:
            raise FileExistsError(f"Output path {output_path} exists. Use overwrite=True to force overwrite.")

    output = io.export_ckpt(path, target, output_path, overwrite, load_connector, modelopt_export_kwargs, **kwargs)

    console = Console()
    console.print(f"[green]✓ Checkpoint exported to {output}[/green]")

    return output


@run.cli.entrypoint(name="generate", namespace="llm")
def generate(
    path: AnyPath,
    trainer: nl.Trainer,
    prompts: Optional[list[str]] = None,
    encoder_prompts: Optional[list[str]] = None,
    input_dataset: Optional[Union[pl.LightningDataModule, str]] = None,
    params_dtype: torch.dtype = torch.bfloat16,
    add_BOS: bool = False,
    max_batch_size: int = 4,
    random_seed: Optional[int] = None,
    inference_batch_times_seqlen_threshold: int = 1000,
    inference_params: Optional["CommonInferenceParams"] = None,
    text_only: bool = False,
    output_path: Optional[AnyPath] = None,
    enable_flash_decode: bool = True,
    **kwargs,
) -> list[Union["InferenceRequest", str]]:
    """
    Generates text using a NeMo LLM model.

    This function takes a checkpoint path and a list of prompts,
    and generates text based on the loaded model and parameters.
    It returns a list of generated text, either as a string or as an InferenceRequest object.

    Python Usage:
    ```python
    strategy = nl.MegatronStrategy(
        tensor_model_parallel_size=2,
        pipeline_model_parallel_size=1,
        context_parallel_size=1,
        sequence_parallel=False,
        setup_optimizers=False,
        store_optimizer_states=False,
    )

    trainer = nl.Trainer(
        accelerator="gpu",
        devices=2,
        num_nodes=1,
        strategy=strategy,
        plugins=nl.MegatronMixedPrecision(
            precision="bf16-mixed",
            params_dtype=torch.bfloat16,
            pipeline_dtype=torch.bfloat16,
            autocast_enabled=False,
            grad_reduce_in_fp32=False,
        ),
    )
    prompts = [
        "Hello, how are you?",
        "How many r's are in the word 'strawberry'?",
        "Which number is bigger? 10.119 or 10.19?",
    ]

    if __name__ == "__main__":
        results = api.generate(
            path=os.path.join(os.environ["NEMO_HOME"], "models", "meta-llama/Meta-Llama-3-8B"),
            prompts=prompts,
            trainer=trainer,
            inference_params=CommonInferenceParams(temperature=0.1, top_k=10, num_tokens_to_generate=512),
            text_only=True,
        )
    ```

    Args:
        path (Union[Path, str]): The path to the model checkpoint.
        prompts (list[str]): The list of prompts to generate text for.
        trainer (nl.Trainer): The trainer object.
        encoder_prompts (Optional[list[str]], optional): The list of encoder prompts. Defaults to None.
        input_dataset (Optional[Union[pl.LightningDataModule, str]], optional): The input data module or jsonl file.
            Test set will be used for generation for data modules. Defaults to None.
        params_dtype (torch.dtype, optional): The data type of the model parameters. Defaults to torch.bfloat16.
        add_BOS (bool, optional): Whether to add the beginning of sequence token. Defaults to False.
        max_batch_size (int, optional): The maximum batch size. Defaults to 4.
        random_seed (Optional[int], optional): The random seed. Defaults to None.
        inference_batch_times_seqlen_threshold (int, optional): If batch-size times sequence-length is smaller than
            this threshold then we will not use pipelining, otherwise we will. Defaults to 1000.
        inference_params (Optional["CommonInferenceParams"], optional): The inference parameters defined in
            Mcore's CommonInferenceParams. Defaults to None.
        text_only (bool, optional): Whether to return only the generated text as a string. Defaults to False.
        output_path (Optional[Union[Path, str]], optional): The path to save the generated text or test dataset
            predictions. Defaults to None.
        enable_flash_decode (bool, optional): Whether to enable flash decode. Defaults to True.
        **kwargs: Additional keyword arguments passed to setup_model_and_tokenizer.

    Returns:
        list[Union["InferenceRequest", str]]: A list of generated text,
            either as a string or as an InferenceRequest object.
    """
    from nemo.collections.llm import inference

    if input_dataset is not None:
        input_path = input_dataset if isinstance(input_dataset, str) else input_dataset.test_path
        with open(input_path) as f:
            dataset = [json.loads(sample) for sample in f.readlines()]
            inputs = [sample["input"] for sample in dataset]
    elif prompts is not None:
        inputs = prompts
    else:
        raise ValueError("Either prompts or input_dataset must be provided.")

    inference_wrapped_model, mcore_tokenizer = inference.setup_model_and_tokenizer(
        path=path,
        trainer=trainer,
        params_dtype=params_dtype,
        inference_batch_times_seqlen_threshold=inference_batch_times_seqlen_threshold,
        enable_flash_decode=enable_flash_decode,
        **kwargs,
    )

    max_seq_length = inference_params.num_tokens_to_generate + max(len(mcore_tokenizer.tokenize(p)) for p in inputs)
    # set kv cache allocation to only num tokens in prompt + max tokens to generate
    inference_wrapped_model.inference_wrapper_config.inference_max_seq_length = max_seq_length
    inference_wrapped_model.inference_context.max_sequence_length = max_seq_length

    if trainer.strategy.expert_model_parallel_size > 1:
        inputs_on_this_dp_rank = inputs
    else:
        dp_size = trainer.strategy.distributed_sampler_kwargs['num_replicas']
        dp_rank = trainer.strategy.distributed_sampler_kwargs['rank']
        chunk_size = (len(inputs) + dp_size - 1) // dp_size
        start_idx = dp_rank * chunk_size
        end_idx = min(start_idx + chunk_size, len(inputs))
        inputs_on_this_dp_rank = inputs[start_idx:end_idx]

    results_on_this_dp_rank = inference.generate(
        model=inference_wrapped_model,
        tokenizer=mcore_tokenizer,
        prompts=inputs_on_this_dp_rank,
        encoder_prompts=encoder_prompts,
        add_BOS=add_BOS,
        max_batch_size=max_batch_size,
        random_seed=random_seed,
        inference_params=inference_params,
    )

    if trainer.strategy.expert_model_parallel_size > 1:
        gathered_results = [r.generated_text if text_only else r for r in results_on_this_dp_rank]
    else:
        gathered_results = [None] * dp_size

        all_gather_object(
            gathered_results,
            [r.generated_text if text_only else r for r in results_on_this_dp_rank],
            group=parallel_state.get_data_parallel_group(),
        )
        gathered_results = [result for sublist in gathered_results for result in sublist]

        assert len(gathered_results) == len(inputs)

    if output_path is not None and is_global_rank_zero():
        with open(output_path, "w") as f:
            for sample, pred in zip(dataset if input_dataset else inputs, gathered_results):
                if type(sample) == dict:
                    sample["label"] = sample.pop("output", None)
                    sample["prediction"] = pred if text_only else pred.generated_text
                elif type(sample) == str:
                    sample = {"input": sample, "prediction": pred if text_only else pred.generated_text}
                f.write(json.dumps(sample) + "\n")
        logging.info(f"Predictions written to {output_path}")

    return gathered_results


def _use_tokenizer(model: pl.LightningModule, data: pl.LightningDataModule, tokenizer: TokenizerType) -> None:
    if tokenizer == "data":
        _set_with_io(model, "tokenizer", data.tokenizer)
    elif tokenizer == "model":
        _set_with_io(data, "tokenizer", model.tokenizer)
    else:
        try:
            from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec

            if isinstance(tokenizer, TokenizerSpec):
                _set_with_io(model, "tokenizer", tokenizer)
                _set_with_io(data, "tokenizer", tokenizer)
            else:
                raise ValueError(f"Expected TokenizerSpec or 'data' or 'model', got: {tokenizer}")
        except ImportError:
            raise ValueError("TokenizerSpec is not available")


def _setup(
    model: pl.LightningModule,
    data: pl.LightningDataModule,
    trainer: Trainer,
    log: Optional[NeMoLogger],
    resume: Optional[AutoResume],
    optim: Optional[OptimizerModule],
    tokenizer: Optional[TokenizerType],
    model_transform: Optional[Union[PEFT, ModelTransform, Callable]],
) -> Any:  # Return type is Any because app_state's type is not specified
    configure_no_restart_validation_training_loop(trainer)
    _log = log or NeMoLogger()
    if resume and isinstance(model_transform, PEFT) and _log.ckpt:
        logging.info("Disabling try_restore_best_ckpt restoration for adapters")
        _log.ckpt.try_restore_best_ckpt = False

    app_state = _log.setup(
        trainer,
        resume_if_exists=getattr(resume, "resume_if_exists", False),
        task_config=getattr(train, "__io__", None),
    )

    # Configure telemetry via CallbackGroup
    CallbackGroup.get_instance().update_config(nemo_version='v2', trainer=trainer, data=data)

    if resume is not None:
        CallbackGroup.get_instance().on_load_checkpoint_start()
        resume.setup(trainer, model)
        CallbackGroup.get_instance().on_load_checkpoint_end()

    if optim:
        CallbackGroup.get_instance().on_optimizer_init_start()
        optim.connect(model)
        CallbackGroup.get_instance().on_optimizer_init_end()
    if tokenizer:  # TODO: Improve this
        _use_tokenizer(model, data, tokenizer)

    if model_transform:
        _set_with_io(model, "model_transform", model_transform)

    # Add ModelTransform callback to Trainer if needed
    if getattr(model, "model_transform", None):
        if not any(isinstance(cb, ModelTransform) for cb in trainer.callbacks):
            if isinstance(model_transform, ModelTransform):
                trainer.callbacks.append(model_transform)
            else:
                trainer.callbacks.append(ModelTransform())
    # Move jit callback at the end ensure it's applied on top of any model transformations (peft)
    jit_cb = None
    for i, cb in enumerate(trainer.callbacks):
        if isinstance(cb, JitTransform):
            assert jit_cb is None
            jit_cb = trainer.callbacks.pop(i)
    if jit_cb is not None:
        trainer.callbacks.append(jit_cb)
    return app_state


def _set_with_io(obj, attr, value):
    setattr(obj, attr, value)
    if hasattr(obj, "__io__") and hasattr(value, "__io__"):
        setattr(obj.__io__, attr, deepcopy(value.__io__))


def _validate_config(
    model: pl.LightningModule,
    data: pl.LightningDataModule,
    trainer: Trainer,
    log: Optional[NeMoLogger] = None,
    resume: Optional[AutoResume] = None,
    optim: Optional[OptimizerModule] = None,
    tokenizer: Optional[TokenizerType] = None,
    model_transform: Optional[Union[PEFT, ModelTransform, Callable]] = None,
) -> None:

    # Model validation
    if hasattr(model, "config"):
        assert getattr(model.config, "seq_length", 1) > 0
        assert getattr(model.config, "max_position_embeddings", 1) > 0
        assert model.config.num_layers > 0
        assert model.config.hidden_size > 0
        assert model.config.num_attention_heads > 0
        assert model.config.ffn_hidden_size > 0
    else:
        assert not isinstance(trainer.strategy, nl.MegatronStrategy), "Expected model.config to exist"

    # Data validation
    assert data.micro_batch_size > 0
    if isinstance(trainer.strategy, nl.MegatronStrategy):
        assert data.global_batch_size > 0
        assert data.seq_length > 0

        assert (
            data.global_batch_size % data.micro_batch_size == 0
        ), "Global batch size must be divisible by micro batch size in data module."

    # Trainer validation

    # MegatronStrategy validation
    if isinstance(trainer.strategy, nl.MegatronStrategy):
        # Basic validation
        assert trainer.strategy.tensor_model_parallel_size > 0
        assert trainer.strategy.pipeline_model_parallel_size > 0
        assert trainer.strategy.context_parallel_size > 0

        # DP validation
        assert (trainer.num_devices * trainer.num_nodes) % (
            trainer.strategy.tensor_model_parallel_size
            * trainer.strategy.pipeline_model_parallel_size
            * trainer.strategy.context_parallel_size
        ) == 0, "Number of GPUs must be divisible by the product of all parallelism sizes for data parallel."

        assert (
            data.global_batch_size
            % (
                data.micro_batch_size
                * (
                    (trainer.num_devices * trainer.num_nodes)
                    / (
                        trainer.strategy.tensor_model_parallel_size
                        * trainer.strategy.pipeline_model_parallel_size
                        * trainer.strategy.context_parallel_size
                    )
                )
            )
            == 0
        ), "Global batch size must be divisible by the product of micro batch size and data parallel size"

        # TP/SP validation
        if trainer.strategy.tensor_model_parallel_size == 1:
            if trainer.strategy.sequence_parallel == True:
                warnings.warn("Disabling sequence parallelism because tensor model parallelism is disabled")
                trainer.strategy.sequence_parallel = False

        # PP/VP validation
        if trainer.strategy.pipeline_model_parallel_size > 1:
            assert (
                trainer.strategy.pipeline_dtype is not None
            ), "pipeline_dtype must be set if pipeline model parallelism is enabled"
        else:
            if trainer.strategy.virtual_pipeline_model_parallel_size is not None:
                warnings.warn("Disabling virtual pipeline parallelism because pipeline model parallelism is disabled")
                trainer.strategy.virtual_pipeline_model_parallel_size = None
            if trainer.strategy.pipeline_dtype is not None:
                warnings.warn("Setting pipeline dtype to None because pipeline model parallelism is disabled")
                trainer.strategy.pipeline_dtype = None

        # CP validation
        if trainer.strategy.context_parallel_size > 1:
            if hasattr(model, "config"):
                if model.config.seq_length is not None:
                    assert (
                        model.config.seq_length % (trainer.strategy.context_parallel_size * 2) == 0
                    ), 'Sequence length must be divisible by 2 * context parallel size if context parallel is used.'
                if isinstance(data, FineTuningDataModule):
                    # check calculate_per_token_loss to be True
                    # check average_in_collective to be False
                    # for context parallel to solve the issue of nan loss on ranks with all tokens masked
                    # (only happens in SFT)
                    assert (
                        model.config.calculate_per_token_loss
                    ), "When finetuning with CP>1, model.config.calculate_per_token_loss must be True"
                    assert (
                        not trainer.strategy.ddp_config.average_in_collective
                    ), "When finetuning with CP>1, average_in_collective must be False"

        # EP validation
        if trainer.strategy.expert_model_parallel_size > 1:
            if hasattr(model, "config"):
                assert (
                    model.config.num_moe_experts is not None
                ), "num_experts must be non None to use expert model parallelism"
                assert (
                    model.config.num_moe_experts % trainer.strategy.expert_model_parallel_size == 0
                ), "Number of experts should be a multiple of expert model parallel_size."


def _build_directory_tree(path, tree=None, root_name=None):
    """Build a Rich Tree representation of a directory structure."""
    from rich.tree import Tree

    path = Path(path)
    if tree is None:
        tree = Tree(f"[bold blue]{root_name or path.name}[/bold blue]")

    # Sort to have directories first, then files
    items = sorted(path.iterdir(), key=lambda x: (not x.is_dir(), x.name))

    for item in items:
        if item.is_dir():
            branch = tree.add(f"[bold cyan]{item.name}/[/bold cyan]")
            _build_directory_tree(item, branch)
        else:
            # Color differently based on file extension
            if item.suffix in ('.json', '.jsonl'):
                tree.add(f"[yellow]{item.name}[/yellow]")
            elif item.suffix in ('.pt', '.bin', '.ckpt', '.nemo'):
                tree.add(f"[magenta]{item.name}[/magenta]")
            elif item.suffix in ('.py', '.sh'):
                tree.add(f"[green]{item.name}[/green]")
            else:
                tree.add(f"[white]{item.name}[/white]")

    return tree


def _load_model_from_path(model: Union[pl.LightningModule, AnyPath]):
    if isinstance(model, AnyPath):
        model = io.load_context(ckpt_to_context_subdir(model), subpath="model")
    return model