Model Support

LLaMA-Factory allows users to add support for custom models. We will take the LLaMA-4 multimodal model as an example to introduce how to add support for a new model. For multimodal models, we need to complete two main tasks:

1. Register the model’s template

2. Parse multimodal data and construct messages

Register Template

First, we can obtain the template of the LLaMA-4 model using the following method:

from transformers import AutoTokenizer, AutoProcessor

tokenizer = AutoTokenizer.from_pretrained("unsloth/Llama-4-Scout-17B-16E-Instruct")
messages = [
    {"role": "user", "content": r"{{content}}"},
    {"role": "assistant", "content": r"{{content}}"},
    {"role": "system", "content": r"{{content}}"},
    {"role": "tool", "content": r"{{content}}"}
]

text = tokenizer.apply_chat_template(messages, tokenize=False,add_generation_prompt=True)
print("========== Template ==========")
print(text)

The output is as follows. By observing the output, we can obtain the model’s chat_template. Additionally, you can also obtain the model’s template via the huggingface repo.

========== Template ==========
<|begin_of_text|><|header_start|>user<|header_end|>

{{content}}<|eot|><|header_start|>assistant<|header_end|>

{{content}}<|eot|><|header_start|>system<|header_end|>

{{content}}<|eot|><|header_start|>ipython<|header_end|>

"{{content}}"<|eot|><|header_start|>assistant<|header_end|>

By observing the output, we can see that the chat_template of LLaMA-4 mainly consists of the following parts:

Message Type	Template Format
User message	<|header_start|>user<|header_end|>\n\n{{content}}<|eot|>
Assistant message	<|header_start|>assistant<|header_end|>\n\n{{content}}<|eot|>
System message	<|header_start|>system<|header_end|>\n\n{{content}}<|eot|>
Tool message	<|header_start|>ipython<|header_end|>\n\n"{{content}}"<|eot|>

We can use the register_template method in src/llamafactory/data/template.py to register the chat_template for the custom model.

In practice, we often append the assistant response header <|header_start|>assistant<|header_end|> after the user’s input to guide the model’s response.

Therefore, we can see that the user message and tool output templates both include the assistant response header, and thus the assistant message format format_assistant omits the assistant response header, keeping only its content part {{content}}<|eot|>.

We can complete the name, format_user, format_assistant, format_system, and format_observation fields based on the output above.

The format_prefix field specifies the beginning of the model sequence and can usually be found in tokenizer_config.json.

The stop_words field specifies the model’s stop words. You can find the eos_token_id in generation_config.json and fill in the corresponding token.

For multimodal models, we also need to specify the multimodal plugin in the mm_plugin field.

register_template(
    # Template
    name="llama4",
    # User Message Format, with a generation prompt template at the end
    format_user=StringFormatter(
        slots=["<|header_start|>user<|header_end|>\n\n{{content}}<|eot|><|header_start|>assistant<|header_end|>\n\n"]
    ),
    # Assistant Message format
    format_assistant=StringFormatter(slots=["{{content}}<|eot|>"]),
    # System Message Format
    format_system=StringFormatter(slots=["<|header_start|>system<|header_end|>\n\n{{content}}<|eot|>"]),
    # Function Call Format
    format_function=FunctionFormatter(slots=["{{content}}<|eot|>"], tool_format="llama3"),
    # Tool Output Format, with a generation prompt template at the end
    format_observation=StringFormatter(
        slots=[
            "<|header_start|>ipython<|header_end|>\n\n{{content}}<|eot|><|header_start|>assistant<|header_end|>\n\n"
        ]
    ),
    # Tool Call Format
    format_tools=ToolFormatter(tool_format="llama3"),
    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
    stop_words=["<|eot|>", "<|eom|>"],
    mm_plugin=get_mm_plugin(name="llama4", image_token="<|image|>"),
)

Multimodal Data Construction

For multimodal models, we refer to the original model to implement multimodal data parsing in LLaMA-Factory.

We can implement the Llama4Plugin class in src/llamafactory/data/mm_plugin.py to parse multimodal data.

The Llama4Plugin class inherits from the BasePlugin class and implements the get_mm_inputs and process_messages methods to parse multimodal data.

Note

@dataclass
class Llama4Plugin(BasePlugin):
    @override
    def process_messages(
        ...
    @override
    def get_mm_inputs(
        ...

The function of get_mm_inputs is to convert multimodal data such as images and videos into inputs that the model can accept, such as pixel_values. To implement get_mm_inputs, we first need to check if the llama4 processor is compatible with the existing implementation. The processing_llama4.py in the official model repository indicates that the data returned by the llama4 processor contains the field pixel_values, which is compatible with the existing implementation in LLaMA-Factory. Therefore, we only need to implement it by referring to the existing get_mm_inputs method.

Note

# https://github.com/hiyouga/LLaMA-Factory/blob/da971c37640de20f97b4d774e77e6f8d5c00b40a/src/llamafactory/data/mm_plugin.py#L264
def _get_mm_inputs(
    self,
    images: list["ImageInput"],
    videos: list["VideoInput"],
    audios: list["AudioInput"],
    processor: "MMProcessor",
    imglens: Optional[list[int]] = None,
) -> dict[str, "torch.Tensor"]:
    r"""Process visual inputs.

    Returns: (llava and paligemma)
        pixel_values: tensor with shape (B, C, H, W)

    Returns: (qwen2-vl)
        pixel_values: tensor with shape (num_patches, patch_dim)
        image_grid_thw: tensor with shape (num_images, 3), where the three numbers are time, width, height
        where num_patches == torch.prod(image_grid_thw)

    Returns: (mllama)
        pixel_values: tensor with shape
                    (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width)
                    For example, (2, 1, 4, 3, 560, 560).
        aspect_ratio_ids: tensor with shape (batch_size, max_num_images). For example, (2, 1).
        aspect_ratio_mask: tensor with shape (batch_size, max_num_images, max_image_tiles). For example, (2, 1, 4).
        num_tiles: List[List[int]] with shape (batch_size, num_images_in_batch). For example, (2, 1).

The function of process_messages is to insert a corresponding number of placeholders in the messages based on the size, quantity, etc., of the input images/videos so that the model can correctly parse the multimodal data. We need to refer to the original repository implementation and the specifications in LLaMA-Factory to return messages of type list[dict[str, str]].

Provide Model Path

Finally, provide the model download path in src/llamafactory/extras/constants.py. For example:

register_model_group(
models={
    "Llama-4-Scout-17B-16E": {
        DownloadSource.DEFAULT: "meta-llama/Llama-4-Scout-17B-16E",
        DownloadSource.MODELSCOPE: "LLM-Research/Llama-4-Scout-17B-16E",
    },
    "Llama-4-Scout-17B-16E-Instruct": {
        DownloadSource.DEFAULT: "meta-llama/Llama-4-Scout-17B-16E-Instruct",
        DownloadSource.MODELSCOPE: "LLM-Research/Llama-4-Scout-17B-16E-Instruct",
    },
    "Llama-4-Maverick-17B-128E": {
        DownloadSource.DEFAULT: "meta-llama/Llama-4-Maverick-17B-128E",
        DownloadSource.MODELSCOPE: "LLM-Research/Llama-4-Maverick-17B-128E",
    },
    "Llama-4-Maverick-17B-128E-Instruct": {
        DownloadSource.DEFAULT: "meta-llama/Llama-4-Maverick-17B-128E-Instruct",
        DownloadSource.MODELSCOPE: "LLM-Research/Llama-4-Maverick-17B-128E-Instruct",
    },
},
template="llama4",
multimodal=True,
)