Repository: oumi-ai/oumi
File: src/oumi/core/configs/params/model_params.py (Python)

Code before:
    """

    compile: bool = False
    """Whether to JIT compile the model.

    For training, do not set this param, and instead set `TrainingParams.compile`.
    """

    chat_template: Optional[str] = None
    """The chat template to use for formatting inputs.

    Options:
    - None: Uses fallback hierarchy (internal config → built-in template → default)
    - "auto": Directly uses model's built-in chat template (recommended for clarity)
    - Custom string: Uses specified Oumi template name (e.g., "llama3-instruct")

    Recommendation: Use explicit "auto" instead of None for less ambiguous behavior.

    Note:
        Different models may require specific chat templates. Consult the model's
        documentation for the appropriate template to use.
    """

    chat_template_kwargs: dict[str, Any] = field(default_factory=dict)
    """Additional keyword args to pass to the chat template renderer.
    Currently, it is being access by vLLM inference engine to pass additional kwargs
    for generation of assistant messages.
    """

    attn_implementation: Optional[str] = None
    """The attention implementation to use.

    Valid options include:

    - None: Use the default attention implementation (spda for torch>=2.1.1, else eager)
    - "sdpa": Use PyTorch's scaled dot-product attention
    - "flash_attention_2": Use Flash Attention 2 for potentially faster computation.
      Requires "flash-attn" package to be installed
    - "eager": Manual implementation of attention
    - "kernels-community/vllm-flash-attn3": Use vLLM Flash Attention 3 kernel from
      HF Hub
    - Custom kernel paths: Any HuggingFace Hub path to attention kernels
    """

    device_map: Optional[str] = "auto"
    """Specifies how to distribute the model's layers across available devices.

    - "auto": Automatically distribute the model across available devices
    - None: Load the entire model on the default device

    Note:

Review comment (suggestion):
```suggestion
    Currently, it is only used by the vLLM inference engine to pass additional kwargs
```

Code after:
    """

    compile: bool = False
    """Whether to JIT compile the model.

    For training, do not set this param, and instead set `TrainingParams.compile`.
    """

    chat_template: Optional[str] = None
    """The chat template to use for formatting inputs.

    Options:
    - None: Uses fallback hierarchy (internal config → built-in template → default)
    - "auto": Directly uses model's built-in chat template (recommended for clarity)
    - Custom string: Uses specified Oumi template name (e.g., "llama3-instruct")

    Recommendation: Use explicit "auto" instead of None for less ambiguous behavior.

    Note:
        Different models may require specific chat templates. Consult the model's
        documentation for the appropriate template to use.
    """

    chat_template_kwargs: dict[str, Any] = field(default_factory=dict)
    """Additional keyword args to pass to the chat template renderer.
    Currently, it is only used by the vLLM inference engine to pass additional kwargs
    for generation of assistant messages.
    """

    attn_implementation: Optional[str] = None
    """The attention implementation to use.

    Valid options include:

    - None: Use the default attention implementation (spda for torch>=2.1.1, else eager)
    - "sdpa": Use PyTorch's scaled dot-product attention
    - "flash_attention_2": Use Flash Attention 2 for potentially faster computation.
      Requires "flash-attn" package to be installed
    - "eager": Manual implementation of attention
    - "kernels-community/vllm-flash-attn3": Use vLLM Flash Attention 3 kernel from
      HF Hub
    - Custom kernel paths: Any HuggingFace Hub path to attention kernels
    """

    device_map: Optional[str] = "auto"
    """Specifies how to distribute the model's layers across available devices.

    - "auto": Automatically distribute the model across available devices
    - None: Load the entire model on the default device

    Note: