Repository: unslothai/unsloth
File: unsloth/models/loader_utils.py (Python)

Code before:
    temp_dir = tempfile.gettempdir()
    new_model_name = model_name.split("/")[-1] + "-fp8"
    new_model_name = os.path.join(temp_dir, new_model_name)
    print(
        f"Quantizing '{model_name}' to fp8, using model_name='{new_model_name}' instead"
    )
    if not os.path.isdir(new_model_name):
        qconfig = _get_torchao_fp8_config()
        qconfig = TorchAoConfig(qconfig)
        model = AutoModel.from_pretrained(
            model_name,
            torch_dtype = "auto",
            device_map = "auto",
            quantization_config = qconfig,
        )
        tokenizer = AutoProcessor.from_pretrained(model_name)
        model.save_pretrained(new_model_name, safe_serialization = False)
        tokenizer.save_pretrained(new_model_name)
    return new_model_name


def _tag_model_with_fp8_torchao_config(model: torch.nn.Module):
    """
    Tag a model with a `TorchAOConfig` so downstream callers will know what to do with it.
    """
    base_config = _get_torchao_fp8_config()
    model.torchao_config = TorchAOConfig(
        qat_scheme = None,
        base_config_and_filter_fns = [(base_config, None)],
    )


def _check_load_in_fp8_settings(
    fast_inference: bool,
    full_finetuning: bool,
    load_in_4bit: bool,
    load_in_8bit: bool,
    load_in_16bit: bool,
    use_exact_model_name: bool,
):
    """
    Assuming `load_in_fp8=True`, raise appropriate errors on incompatible settings
    and environment. Currently this feature requires:
    1. H100 GPUs or after
    2. torchao 0.15.0+ (or nightly)
    3. torch 2.9.0+
    4. If fbgemm_gpu_genai is installed, require 1.4.1+
    """
    if not fast_inference:
        raise ValueError(
            "Unsloth: `load_in_fp8` is only supported for `fast_inference` for now"

Review comment (suggestion):
@danielhanchen I had to change this back for this to work. When I tried `AutoModel` it removed the `lm_head` from the state dict, which later caused an out of bounds exception on this line: https://github.com/unslothai/unsloth-zoo/blob/54dce973426ee61670e15b720619c8539bf05104/unsloth_zoo/vllm_utils.py#L1110. Maybe we can generalize this to multimodal models later

Code after:
    Quantizes the model to fp8 using torchao and saving the quantized model to a
    temporary location. Return the path to the quantized model.

    Note: Once on-the-fly quantization is added in vllm in
    https://github.com/vllm-project/vllm/pull/26327, we should
    dynamically quantize the model there instead:

      llm = LLM(
        ...
        hf_overrides={"quantization_config_file": "torchao_config.json"},
      )
    """
    temp_dir = tempfile.gettempdir()
    new_model_name = model_name.split("/")[-1] + "-fp8-" + fp8_mode
    new_model_name = os.path.join(temp_dir, new_model_name)
    print(
        f"Unsloth: Quantizing '{model_name}' to fp8, using model_name='{new_model_name}' instead"
    )

    if not os.path.isdir(new_model_name):
        qconfig = _get_torchao_fp8_config(fp8_mode)
        qconfig = TorchAoConfig(qconfig)
        # TODO: generalize this to beyond text models?
        # Right now using AutoModel removes the `lm_head` layer,
        # which is expected later when loading the vllm state dict
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype = "auto",
            device_map = "auto",
            quantization_config = qconfig,
        )
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model.save_pretrained(new_model_name, safe_serialization = False)
        tokenizer.save_pretrained(new_model_name)
    return new_model_name


def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
    """
    Tag a model with a `TorchAOConfig` so downstream callers will know what to do with it.
    """
    base_config = _get_torchao_fp8_config(fp8_mode)
    model.torchao_config = TorchAOConfig(
        qat_scheme = None,
        base_config_and_filter_fns = [(base_config, None)],
    )


def _get_fp8_mode_and_check_settings(
    load_in_fp8: Union[bool, str],
    fast_inference: bool,