api_ref_modules.rst

torchtune.modules

.. currentmodule:: torchtune.modules

Modeling Components and Building Blocks

.. autosummary::
    :toctree: generated/
    :nosignatures:

    MultiHeadAttention
    FeedForward
    KVCache
    RotaryPositionalEmbeddings
    RMSNorm
    Fp32LayerNorm
    TanhGate
    TiedLinear
    TransformerSelfAttentionLayer
    TransformerCrossAttentionLayer
    TransformerDecoder
    VisionTransformer
    LayerDropout
    prepare_layer_dropout
    classifier_model

Losses

.. autosummary::
    :toctree: generated/
    :nosignatures:

    loss.LinearCrossEntropyLoss
    loss.ForwardKLLoss
    loss.ForwardKLWithChunkedOutputLoss
    loss.ReverseKLWithChunkedOutputLoss
    loss.SymmetricKLWithChunkedOutputLoss

Optimizers

.. autosummary::
    :toctree: generated/
    :nosignatures:

    optim.OptimizerInBackward

Base Tokenizers

Base tokenizers are tokenizer models that perform the direct encoding of text into token IDs and decoding of token IDs into text. These are typically byte pair encodings that underlie the model specific tokenizers.

.. autosummary::
    :toctree: generated/
    :nosignatures:

    transforms.tokenizers.SentencePieceBaseTokenizer
    transforms.tokenizers.TikTokenBaseTokenizer
    transforms.tokenizers.HuggingFaceBaseTokenizer
    transforms.tokenizers.ModelTokenizer
    transforms.tokenizers.BaseTokenizer

Tokenizer Utilities

These are helper methods that can be used by any tokenizer.

.. autosummary::
    :toctree: generated/
    :nosignatures:

    transforms.tokenizers.tokenize_messages_no_special_tokens
    transforms.tokenizers.parse_hf_tokenizer_json

PEFT Components

.. autosummary::
    :toctree: generated/
    :nosignatures:

    peft.LoRALinear
    peft.DoRALinear
    peft.AdapterModule
    peft.get_adapter_params
    peft.set_trainable_params
    peft.get_adapter_state_dict
    peft.validate_missing_and_unexpected_for_lora
    peft.disable_adapter

Fusion Components

Components for building models that are a fusion of two+ pre-trained models.

.. autosummary::
    :toctree: generated/
    :nosignatures:

    model_fusion.DeepFusionModel
    model_fusion.FusionLayer
    model_fusion.FusionEmbedding
    model_fusion.register_fusion_module
    model_fusion.get_fusion_params

Module Utilities

These are utilities that are common to and can be used by all modules.

.. autosummary::
   :toctree: generated/
   :nosignatures:

   common_utils.reparametrize_as_dtype_state_dict_post_hook
   common_utils.local_kv_cache
   common_utils.disable_kv_cache
   common_utils.delete_kv_caches
   embedding_utils.resize_token_embeddings

Vision Transforms

Functions used for preprocessing images.

.. autosummary::
   :toctree: generated/
   :nosignatures:

    transforms.Transform
    transforms.VisionCrossAttentionMask