Repository: oumi-ai/oumi
File: src/oumi/builders/models.py (Python)

Code before:
        )

    if model_params.adapter_model is not None:
        raise NotImplementedError(
            "Loading PEFT adapters is not yet supported for custom Oumi models. "
            "If you need to use PEFT adapters, please use a HuggingFace model instead "
            "or open a feature request at https://github.com/oumi-ai/oumi."
        )

    dtype = model_params.torch_dtype
    model = model.to(dtype=dtype)
    # Needed for MFUTrainerCallback
    model.dtype = dtype
    return model


def _get_quantization_config_for_training(model_params: ModelParams):
    """Get the appropriate quantization config for training.

    For MXFP4 quantized models that need dequantization during training,
    this creates the appropriate Mxfp4Config(dequantize=True) configuration.

    Args:
        model_params: Model parameters that may contain quantization config.

    Returns:
        Quantization config object or None if no quantization needed.
    """
    # Check if model_kwargs contains quantization_config
    if not model_params.model_kwargs:
        return None

    quant_config = model_params.model_kwargs.get("quantization_config")
    if not quant_config:
        return None

    # Handle MXFP4 quantization for training (requires dequantization)
    if isinstance(quant_config, dict) and quant_config.get("quant_method") == "mxfp4":
        if not MXFP4_AVAILABLE:
            raise ImportError(
                "MXFP4 quantization requires transformers>=4.55.0 with "
                "Mxfp4Config support. Please upgrade transformers: "
                "pip install 'transformers>=4.55.0'"
            )

        logger.info(
            "Detected MXFP4 quantized model. Creating Mxfp4Config(dequantize=True) "
            "for training."
        )
        return Mxfp4Config(dequantize=True)  # pyright: ignore[reportOptionalCall]


Review comment (suggestion):
Having it as a model kwarg is a bit hacky as opposed to a built-in field, but we can proceed with this for now to unblock, and add an issue to make it more formal (could merge with some of the qlora fields in PeftParams perhaps)

Code after:
        )

    dtype = model_params.torch_dtype
    model = model.to(dtype=dtype)
    # Needed for MFUTrainerCallback
    model.dtype = dtype
    return model


def _get_quantization_config_for_training(model_params: ModelParams):
    """Get the appropriate quantization config for training.

    For MXFP4 quantized models that need dequantization during training,
    this creates the appropriate Mxfp4Config(dequantize=True) configuration.

    Args:
        model_params: Model parameters that may contain quantization config.

    Returns:
        Quantization config object or None if no quantization needed.
    """
    # Check if model_kwargs contains quantization_config
    if not model_params.model_kwargs:
        return None

    quant_config = model_params.model_kwargs.get("quantization_config")
    if not quant_config:
        return None

    # Handle MXFP4 quantization for training (requires dequantization)
    if isinstance(quant_config, dict) and quant_config.get("quant_method") == "mxfp4":
        logger.info(
            "Detected MXFP4 quantized model. Creating Mxfp4Config(dequantize=True) "
            "for training."
        )
        return Mxfp4Config(dequantize=True)  # pyright: ignore[reportOptionalCall]

    return None


def _disable_cache_in_model_config(model: transformers.PreTrainedModel) -> None:
    # Required for FSDP.
    # Context: https://github.com/huggingface/transformers/issues/28499
    model.config.use_cache = False
    if hasattr(model.config, "text_config"):
        # This may be needed for VLM-s.
        model.config.text_config.use_cache = False


def build_huggingface_model(
    model_params: ModelParams,