tests/mlops/test_base_quantization.py

"""Le point de depart de la recette : un checkpoint DEJA quantifie.

La recette partait de `Qwen/Qwen3.8-27B` (poids bf16, ~54 Go telecharges) et
quantifiait a la volee via un `BitsAndBytesConfig` passe dans
`model_init_kwargs`. Elle part maintenant de
`unsloth/Qwen3.8-27B-unsloth-bnb-4bit` (22,4 Go).

Le piege de ce changement n'est pas la taille, c'est l'INTERACTION. Quand le
checkpoint porte deja un `quantization_config` dans son `config.json` ET qu'on
passe un `BitsAndBytesConfig` explicite, `transformers` n'echoue pas et ne
fusionne pas : il IGNORE celui qu'on passe et garde celui du checkpoint, avec
pour seule trace un `warnings.warn` que 25 h de log noient. Une recette qui
continue de le passer decrit donc une quantification qui ne sera jamais
appliquee.

Ce fichier mesure ce fait sur la pile INSTALLEE plutot que de le supposer, puis
verifie la decision qu'on en tire.
"""

import pathlib
import textwrap
import types
import warnings

import torch

from scripts.training.base_quantization import (
    build_model_init_kwargs,
    is_already_quantized,
)

# --- Le fait, MESURE sur la pile installee --------------------------------


def test_transformers_ignores_a_passed_config_when_the_checkpoint_carries_one():
    """La reponse a « erreur ? ecrasement silencieux ? avertissement ignore ? ».

    On appelle la fonction que `from_pretrained` appelle
    (`transformers/quantizers/auto.py`, `get_hf_quantizer` ->
    `merge_quantization_configs`) avec un config de checkpoint et un config
    d'appel VOLONTAIREMENT differents. Si un jour transformers se mettait a
    honorer celui qu'on passe, ce test rougit et la decision ci-dessous doit
    etre relue -- pas conservee telle quelle."""
    from transformers import BitsAndBytesConfig
    from transformers.quantizers.auto import AutoHfQuantizer

    from_checkpoint = {
        "quant_method": "bitsandbytes",
        "load_in_4bit": True,
        "load_in_8bit": False,
        "bnb_4bit_quant_type": "nf4",
        "bnb_4bit_use_double_quant": True,
        "bnb_4bit_compute_dtype": "bfloat16",
    }
    ours = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="fp4",
        bnb_4bit_use_double_quant=False,
    )

    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        merged = AutoHfQuantizer.merge_quantization_configs(from_checkpoint, ours)

    # Ni erreur, ni fusion : les DEUX attributs qu'on a poses differents sont
    # revenus a la valeur du checkpoint.
    assert merged.bnb_4bit_quant_type == "nf4", (
        "transformers a honore le quant_type passe a from_pretrained : la "
        "raison d'etre de ce module tombe, relire la decision."
    )
    assert merged.bnb_4bit_use_double_quant is True
    # Et la seule trace est un avertissement -- pas une exception.
    assert any(
        "already has a `quantization_config`" in str(w.message) for w in caught
    ), f"aucun avertissement emis : {[str(w.message) for w in caught]}"


def test_the_prequantized_fact_is_read_where_transformers_reads_it():
    """`is_already_quantized` doit lire EXACTEMENT le meme fait que
    `get_hf_quantizer` (`pre_quantized = hasattr(config, "quantization_config")`),
    pas une heuristique sur le nom du depot -- un nom peut mentir, la config
    non."""
    import inspect
    import pathlib
    import sys
    import sysconfig

    candidates = [
        pathlib.Path(sysconfig.get_paths()[key])
        / "transformers"
        / "quantizers"
        / "auto.py"
        for key in ("purelib", "platlib")
    ]
    candidates += list(
        pathlib.Path(sys.prefix).glob(
            "**/site-packages/transformers/quantizers/auto.py"
        )
    )
    source = None
    for candidate in candidates:
        if candidate.is_file():
            source = candidate.read_text(encoding="utf-8", errors="replace")
            break
    assert source is not None, (
        "transformers est installe mais sa source est introuvable : le garde "
        "se sauterait lui-meme alors qu'il a de quoi mesurer."
    )
    assert 'pre_quantized = hasattr(config, "quantization_config")' in source
    # Et notre lecture est la meme. Verifie sur l'AST du CORPS, docstring
    # exclue : la seule chaine que la decision a le droit de connaitre est
    # "quantization_config". Une mutation vers une heuristique de nom
    # (`return "4bit" in repo`) introduit une autre chaine, et rougit ici.
    import ast

    tree = ast.parse(textwrap.dedent(inspect.getsource(is_already_quantized)))
    body = tree.body[0].body
    if isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant):
        body = body[1:]  # docstring
    strings = {
        node.value
        for statement in body
        for node in ast.walk(statement)
        if isinstance(node, ast.Constant) and isinstance(node.value, str)
    }
    assert strings == {"quantization_config"}, (
        "la decision connait d'autres chaines que le nom de l'attribut de "
        f"config : {sorted(strings)}. Une heuristique sur le nom du depot "
        "(« -4bit », « bnb ») ferait mentir un depot mal nomme."
    )


# --- La decision qu'on en tire --------------------------------------------


def _prequantized_config():
    """Un config de la forme que rend `AutoConfig.from_pretrained` sur un
    checkpoint bnb-4bit : l'attribut existe, c'est tout ce qui compte."""
    return types.SimpleNamespace(
        quantization_config={
            "quant_method": "bitsandbytes",
            "load_in_4bit": True,
            "bnb_4bit_quant_type": "nf4",
        }
    )


def test_no_bitsandbytes_config_is_handed_to_an_already_quantized_base():
    kwargs = build_model_init_kwargs(_prequantized_config(), torch.bfloat16)

    assert "quantization_config" not in kwargs, (
        "la recette repasse un BitsAndBytesConfig a un checkpoint qui en porte "
        "deja un : transformers l'ignore (voir le test mesure ci-dessus) et la "
        "recette decrit alors une quantification qui n'est jamais appliquee."
    )


def test_the_full_weights_base_is_still_quantized_on_the_fly():
    """Le chemin ne doit PAS disparaitre : `TRAIN_BASE_MODEL` reste
    surchargeable, et pointer un depot de poids pleins sans quantifier ferait
    tenir 54 Go de bf16 en VRAM -- OOM immediat sur toute carte visee."""
    kwargs = build_model_init_kwargs(types.SimpleNamespace(), torch.bfloat16)

    quantization = kwargs.get("quantization_config")
    assert quantization is not None, (
        "un depot de poids pleins ne serait plus quantifie du tout : 54 Go de "
        "bf16 en VRAM."
    )
    assert quantization.load_in_4bit is True
    assert quantization.bnb_4bit_quant_type == "nf4"
    assert quantization.bnb_4bit_use_double_quant is True
    assert quantization.bnb_4bit_compute_dtype == torch.bfloat16


def test_the_compute_dtype_travels_with_the_rest_in_both_cases():
    """`dtype` et `trust_remote_code` ne dependent pas du point de depart : les
    perdre sur la branche pre-quantifiee chargerait en fp32 par defaut."""
    for base_config in (types.SimpleNamespace(), _prequantized_config()):
        kwargs = build_model_init_kwargs(base_config, torch.bfloat16)
        assert kwargs["dtype"] == torch.bfloat16
        assert kwargs["trust_remote_code"] is True


# --- Un LoRA entraine sur du 4-bit s'applique-t-il a la meme architecture ? --
#
# La question n'est pas rhetorique : c'est elle qui decide si l'adaptateur
# publie pourra etre converti en GGUF contre la config de poids PLEINS que
# `export_otaku_adapter.sh` donne a `convert_lora_to_gguf.py`. Si les noms de
# modules ciblables changeaient entre les deux points de depart, le
# convertisseur ne saurait plus a quoi rattacher les tenseurs -- et rien en
# aval ne le dirait (ce depot a deja vu convert + ollama create + ollama run
# sortir tous en code 0 sur un adaptateur inutilisable).
#
# On la tranche ICI, sans GPU : `replace_with_bnb_linear` s'execute sur CPU, et
# `get_peft_model` rend les cles reelles de l'etat sauvegarde.


def _script_target_modules() -> list:
    """Les `target_modules` REELS de la recette, lus dans son AST -- pas une
    liste recopiee qui divergerait au premier changement."""
    import ast

    script = (
        pathlib.Path(__file__).resolve().parents[2]
        / "scripts"
        / "training"
        / "train_otaku_qwen38.py"
    ).read_text(encoding="utf-8")
    for node in ast.walk(ast.parse(script)):
        if (
            isinstance(node, ast.Call)
            and getattr(node.func, "id", None) == "LoraConfig"
        ):
            for keyword in node.keywords:
                if keyword.arg == "target_modules":
                    return [element.value for element in keyword.value.elts]
    raise AssertionError("LoraConfig(target_modules=...) introuvable dans la recette")


def _adapter_keys(model, target_modules):
    from peft import LoraConfig, get_peft_model

    peft_model = get_peft_model(
        model, LoraConfig(r=4, target_modules=target_modules, task_type=None)
    )
    return sorted(key for key in peft_model.state_dict() if "lora_" in key)


def _toy_model(target_modules):
    import torch.nn as nn

    class Block(nn.Module):
        def __init__(self):
            super().__init__()
            for name in target_modules:
                setattr(self, name, nn.Linear(8, 8))
            # Les noms fusionnes de la tour vision, qui ne doivent JAMAIS etre
            # matches par la recette.
            self.qkv = nn.Linear(8, 24)
            self.linear_fc1 = nn.Linear(8, 8)

    class Model(nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = nn.ModuleList([Block()])
            self.lm_head = nn.Linear(8, 8)

    return Model()


def test_the_adapter_keys_are_the_same_on_4bit_and_on_full_weights():
    """Reponse mesuree : `bitsandbytes` remplace `nn.Linear` par
    `bnb.nn.Linear4bit` SOUS LE MEME NOM
    (`transformers/integrations/bitsandbytes.py` :
    `model._modules[name] = bnb.nn.Linear4bit(...)`), donc `target_modules`
    matche a l'identique et l'etat sauvegarde porte les memes cles. C'est ce
    qui rend l'adaptateur convertible contre la config de poids pleins que
    `export_otaku_adapter.sh` passe a convert_lora_to_gguf.py."""
    from transformers import BitsAndBytesConfig
    from transformers.integrations.bitsandbytes import replace_with_bnb_linear

    targets = _script_target_modules()

    on_full_weights = _adapter_keys(_toy_model(targets), targets)

    quantized = _toy_model(targets)
    quantized = replace_with_bnb_linear(
        quantized,
        modules_to_not_convert=["lm_head"],
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            llm_int8_skip_modules=["lm_head"],
        ),
    )
    # Anti-vacuite : sans ca, un `replace_with_bnb_linear` devenu sans effet
    # ferait passer le test en comparant deux fois le meme modele.
    replaced = {
        type(module).__name__
        for name, module in quantized.named_modules()
        if name.endswith(tuple(targets))
    }
    assert replaced == {"Linear4bit"}, (
        "les modules cibles n'ont pas ete quantifies : la comparaison "
        f"ci-dessous ne prouverait rien ({replaced})."
    )

    on_4bit = _adapter_keys(quantized, targets)

    assert on_full_weights == on_4bit, (
        "les cles de l'adaptateur different selon le point de depart : "
        "convert_lora_to_gguf.py ne saurait plus a quoi rattacher les "
        "tenseurs, et rien en aval ne le dirait.\n"
        f"poids pleins : {on_full_weights}\n4-bit        : {on_4bit}"
    )
    assert on_4bit, "aucune cle LoRA : le harnais ne mesure rien"
    # Et la tour vision reste hors du LoRA, des deux cotes.
    assert not any(
        fragment in key for key in on_4bit for fragment in (".qkv.", ".linear_fc1.")
    ), "la recette matche des modules fusionnes de la tour vision"