모델 테스트 작성하기

모델 테스트 작성하기 (Writing model tests)

Transformers 테스트 스위트는 믹스인(mixin) 기반 아키텍처를 사용해서 최소한의 코드로 100개 이상의 테스트를 자동 생성해요. 모델별 코드를 조금만 작성하면, 믹스인이 저장/로드, 생성, 파이프라인, 학습, 텐서 병렬화를 처리해 줘요.

출처: 문서

본문

다음 명령으로 모델의 테스트를 실행해요.

# run your model's tests
pytest tests/models/mymodel/test_modeling_mymodel.py -v

# run a specific test
pytest tests/models/mymodel/test_modeling_mymodel.py::MyModelTest::test_model

# run tests matching a keyword pattern (useful to run all integration tests)
pytest tests/models/mymodel/ -k integration -v

# include slow integration tests
RUN_SLOW=1 pytest tests/models/mymodel/ -v

Hugging Face CI는 모든 풀 리퀘스트에서 @slow가 붙지 않은 모델 테스트를 실행하고, 느린 테스트는 야간 일정에 따라 돌려요 (CI가 무엇을 검증하는지는 Pull request checks 참고).

기본 테스트 클래스 고르기

가장 흔한 모델군을 다루는 세 가지 기본 클래스가 있어요. 모델의 모달리티에 맞는 것을 고르세요.

Base class Use for Mixins
CausalLMModelTest Causal language models ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, TrainingTesterMixin, TensorParallelTesterMixin
VLMModelTest Vision-language models ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin
ALMModelTest Audio-language models ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin

VLMModelTest와 ALMModelTest는 공통의 MultiModalModelTest 부모를 공유하는데, 이 부모는 하위 config들을 복합(composite) 최상위 config 안에 중첩시키고, 원시 모달리티 특징(오디오 또는 비전)과 함께 모달리티 플레이스홀더 토큰을 input_ids에 배치해요. CausalLMModelTest는 다중 모달 부모를 사용하지 않아요. 이 클래스는 공유된 세 가지 믹스인에 기반하며, 학습과 텐서 병렬화 커버리지를 위해 TrainingTesterMixin과 TensorParallelTesterMixin을 추가해요.

세 가지 중 어느 것에도 맞지 않는 아키텍처(인코더 전용, 인코더-디코더 등)라면 아래 설명하는 two-class pattern과 test mixins에서 직접 테스트 인프라를 구성해요.

CausalLMModelTest

CausalLMModelTest는 causal language model을 테스트하기 위해 권장되는 기본 클래스예요. 다섯 개의 test mixins에서 상속받아 저장/로드, 생성, 파이프라인, 학습, 텐서 병렬화 테스트를 자동 생성해요.

import unittest

from transformers.testing_utils import require_torch
from transformers import is_torch_available

from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester

if is_torch_available():
    from transformers import MyModel

class MyModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = MyModel

@require_torch
class MyModelTest(CausalLMModelTest, unittest.TestCase):
    model_tester_class = MyModelTester

이 두 클래스는 MyModel과 그 모든 헤드 클래스(MyModelForCausalLM, MyModelForSequenceClassification 등)에 대해 완전한 테스트 커버리지를 제공해요. 실제 예시는 tests/models/llama/test_modeling_llama.py를 참고하세요.

CausalLMModelTester는 base_model_class만 요구해요. tester는 Model 접미사를 떼어 기본 이름을 얻고(LlamaModel은 Llama가 됨), 그다음 Config나 ForCausalLM 같은 접미사를 붙여서 관련 클래스를 찾아요. 모듈에 클래스가 없으면 해당 속성은 None으로 남고 tester는 해당 테스트를 건너뛰어요.

CausalLMTester에서 기본값 재정의하기

모델이 표준 명명 규칙을 따르지 않거나 동작을 커스터마이즈해야 한다면, tester 또는 테스트 클래스의 속성을 재정의해요.

class MyModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = MyModel
        # override if the class name doesn't follow the convention
        causal_lm_class = MyCustomCausalLM

@require_torch
class MyModelTest(CausalLMModelTest, unittest.TestCase):
    model_tester_class = MyModelTester
    # disable embedding resize tests
    test_resize_embeddings = False

tester에 커스텀 생성자 파라미터가 필요한 모델의 경우 __init__을 재정의하고, 추가 속성을 설정하기 전에 super().__init__(parent=parent)를 호출해요. 실제 예시는 tests/models/youtu/test_modeling_youtu.py를 참고하세요.

class YoutuModelTester(CausalLMModelTester):
    if is_torch_available():
        base_model_class = YoutuModel

    def __init__(self, parent, kv_lora_rank=16, q_lora_rank=32):
        super().__init__(parent=parent)
        self.kv_lora_rank = kv_lora_rank
        self.q_lora_rank = q_lora_rank

VLMModelTest

VLMModelTest는 vision-language 모델을 위한 기본 클래스예요. 세 가지 믹스인(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin)에서 상속받고, 여러 하위 모델을 처리하기 위해 _is_composite = True로 설정해요.

import unittest

from transformers.testing_utils import require_torch
from transformers import is_torch_available

from ...vlm_tester import VLMModelTest, VLMModelTester

if is_torch_available():
    from transformers import (
        MyVLMConfig,
        MyVLMModel,
        MyVLMTextConfig,
        MyVLMVisionConfig,
        MyVLMForConditionalGeneration,
    )

class MyVLMTester(VLMModelTester):
    if is_torch_available():
        base_model_class = MyVLMModel
        config_class = MyVLMConfig
        text_config_class = MyVLMTextConfig
        vision_config_class = MyVLMVisionConfig
        conditional_generation_class = MyVLMForConditionalGeneration

@require_torch
class MyVLMTest(VLMModelTest, unittest.TestCase):
    model_tester_class = MyVLMTester

VLMModelTester에서 기본값 재정의하기

VLM이 커스텀 비전 파라미터나 기본값이 아닌 config 값이 필요할 때는 __init__을 재정의해요. super().__init__(parent, **kwargs)를 호출하기 전에 setdefault로 기본값을 설정해요. 아래 예시는 tests/models/qianfan_ocr/test_modeling_qianfan_ocr.py의 첫 몇 가지 기본값을 보여줘요.

class QianfanOCRVisionText2TextModelTester(VLMModelTester):
    base_model_class = QianfanOCRModel
    config_class = QianfanOCRConfig
    text_config_class = Qwen3Config
    vision_config_class = QianfanOCRVisionConfig
    conditional_generation_class = QianfanOCRForConditionalGeneration

    def __init__(self, parent, **kwargs):
        kwargs.setdefault("image_token_id", 1)
        kwargs.setdefault("image_size", 32)
        kwargs.setdefault("patch_size", 4)
        kwargs.setdefault("num_channels", 3)
        # ... more defaults
        super().__init__(parent, **kwargs)

VLM 테스트는 CausalLMModelTest와 몇 가지 점에서 달라요.

  • tester에 config_class, text_config_class, vision_config_class, conditional_generation_class를 반드시 설정해야 해요.
  • VLMModelTest는 TrainingTesterMixin 또는 TensorParallelTesterMixin을 포함하지 않아요.
  • tester의 __init__은 **kwargs와 setdefault()에서 비전 파라미터(image_size, patch_size, num_channels, num_image_tokens)를 받아요.
  • ConfigTester는 최상위 config가 텍스트 모델 config가 아닌 복합 config이기 때문에 has_text_modality=False를 사용해요.

ALMModelTest

ALMModelTest는 Qwen2Audio, AudioFlamingo3, GraniteSpeech 같은 오디오-언어 모델(ALM, Audio-Language Model)을 위한 기본 클래스예요. 이것은 같은 MultiModalModelTest 부모와 헤드 클래스 자동 발견으로 VLM 패턴을 따르되, 비전 쪽 메커니즘을 오디오 특징, 오디오 하위 config, 오디오 토큰 배치 전략으로 바꾼 것이에요.

class MyALMTester(ALMModelTester):
    config_class = MyALMConfig
    text_config_class = MyALMTextConfig
    audio_config_class = MyALMAudioConfig
    conditional_generation_class = MyALMForConditionalGeneration
    audio_mask_key = "feature_attention_mask"

class MyALMTest(ALMModelTest, unittest.TestCase):
    model_tester_class = MyALMTester

ALMModelTester에서 기본값 재정의하기

tester의 __init__은 ALM 특유의 기본값(feat_seq_length=128, num_mel_bins=80, audio_token_id=0)을 설정해요. super().__init__(parent, **kwargs)를 호출하기 전에 setdefault로 재정의해요.

모델이 무엇을 어떻게 부르는지 알려주는 클래스 속성이 두 개 있어요.

  • audio_mask_key: 모델이 오디오 마스크에 기대하는 kwarg 이름("feature_attention_mask", "input_features_mask" 등). 모델이 별도의 오디오 마스크를 사용하지 않으면 None으로 남겨둬요.
  • audio_config_key: 최상위 config가 오디오 하위 config를 중첩할 때 사용하는 속성 이름. 기본값은 "audio_config"이지만 GraniteSpeech 같은 모델은 "encoder_config"를 사용해요.
class Qwen2AudioModelTester(ALMModelTester):
    def __init__(self, parent, **kwargs):
        kwargs.setdefault("feat_seq_length", 60)
        kwargs.setdefault("max_source_positions", kwargs["feat_seq_length"] // 2)
        super().__init__(parent, **kwargs)

ALMModelTester는 하나의 훅 get_audio_embeds_mask(audio_mask)를 반드시 재정의해야 하며, 커스터마이즈를 위해 선택적인 훅 몇 개를 더 제공해요.

  • get_audio_embeds_mask(audio_mask): 인코더의 다운샘플링 후 오디오 임베딩 위치의 배치별 마스크를 반환해요. tester는 행별 합계를 사용해 input_ids에 삽입할 audio_token_id 플레이스홀더 수를 정하므로, 그 개수는 인코더가 내보내는 것과 일치해야 해요.
  • create_audio_features(): 오디오 특징 텐서를 반환해요. 기본 형상은 [batch_size, num_mel_bins, feat_seq_length]예요. GraniteSpeech처럼 시간 우선 특징([batch_size, feat_seq_length, num_mel_bins])을 기대하는 모델에서는 재정의해요.
  • create_audio_mask(): 오디오 레벨 어텐션 마스크를 반환해요. 기본값은 배치의 각 행에 무작위의 연속된 유효 영역을 만듭니다. 테스트가 두 prepare_config_and_inputs_for_common() 호출을 서로 비교하거나, 오디오 인코더가 null이 아닌 마스크를 거부하는 백엔드로 디스패치한다면, 결정적인 전체 길이 마스크로 재정의해요.
  • place_audio_tokens(input_ids, config, num_audio_tokens): BOS 뒤에 오디오 플레이스홀더 토큰을 연속적으로 배치해요. 모델이 다른 레이아웃을 필요로 할 때만 재정의해요.
  • get_audio_feature_key(): 오디오 특징용 inputs-dict 키("input_features" 기본값)를 반환해요.

상속된 다중 모달 테스트 외에도 ALMModelTest는 test_mismatching_num_audio_tokens을 추가해요. 이 테스트는 오디오 특징 개수가 input_ids의 오디오 플레이스홀더 토큰 개수와 일치하지 않을 때 모델이 명확한 ValueError를 발생시키는지 확인하고, 여러 오디오 세그먼트가 있는 프롬프트가 여전히 성공적으로 forward 되는지 검증해요.

다른 아키텍처용 테스트 작성하기

인코더 전용, 인코더-디코더, 오디오 또는 기타 비표준 아키텍처의 경우 아래 설명하는 two-class 패턴과 테스트 믹스인에서 직접 테스트 인프라를 구성해요.

ModelTester와 ModelTest

모든 모델 테스트 파일은 같은 구조를 따라요.

  1. ModelTester(일반 클래스)는 테스트용으로 아주 작은 config와 더미 입력을 만들고, 모델 특유의 작은 회귀 테스트를 보유할 수도 있어요.
  2. ModelTest(unittest.TestCase + 믹스인)는 자동 생성된 테스트를 상속받아 모든 모델 변형에 대해 실행해요.

ModelTest는 tester에서 prepare_config_and_inputs_for_common()을 호출해서 (config, inputs_dict) 튜플을 얻어요. 모든 믹스인은 테스트 데이터를 위해 prepare_config_and_inputs_for_common()을 사용해요.

테스트 믹스인

모델에 필요한 믹스인을 고르세요.

Mixin Source file What it tests
ModelTesterMixin tests/test_modeling_common.py Save/load, gradient checkpointing, forward signature, common attributes
GenerationTesterMixin tests/generation/test_utils.py Greedy, sampling, beam search, assisted decoding
PipelineTesterMixin tests/test_pipeline_mixin.py One test per pipeline task
TrainingTesterMixin tests/test_training_mixin.py Overfitting on a small batch
TensorParallelTesterMixin tests/test_tensor_parallel_mixin.py Distributed tensor parallelism

모델 테스트 작성하기

완전한 작동 예시는 tests/models/modernbert/test_modeling_modernbert.py를 참고하세요. 핵심 단계는 아래에 정리돼 있어요.

  1. ModelTester 클래스가 아주 작은 config와 더미 입력을 만들어요. 테스트가 CPU에서 몇 초 안에 끝나도록 차원을 작게 유지하세요. 입력을 만들려면 아래의 세 가지 텐서 헬퍼를 사용해요.

    • ids_tensor(shape, vocab_size): [0, vocab_size) 범위의 무작위 정수 텐서. input_ids, token_type_ids, 라벨 텐서에 사용해요.
    • random_attention_mask(shape): 첫 토큰이 항상 1인 이진(0과 1) 텐서. attention_mask에 사용해요.
    • floats_tensor(shape, scale=1.0): 무작위 실수 텐서. pixel_values나 inputs_embeds 같은 연속 입력에 사용해요.

    tester는 get_config(), prepare_config_and_inputs(), prepare_config_and_inputs_for_common()을 구현해야 해요. 각 작업 헤드(기본 모델, 시퀀스 분류, 토큰 분류 등)마다 create_and_check_* 메서드를 추가해요.

  2. 모델이 필요한 믹스인에서 상속받고, all_model_classes와 pipeline_model_mapping을 설정하며, setUp()을 정의해요. tester의 create_and_check_* 메서드에 위임하는 test_* 메서드를 작성해요.

  3. 각 작업 헤드에 대해 tester에 모델을 인스턴스화하고 forward pass를 실행하며 출력 형상을 검증하는 create_and_check_* 메서드를 추가하고, 테스트 클래스에 대응하는 test_* 메서드를 추가해요.

파일 구성

테스트 파일은 아래 구조처럼 tests/models/mymodel/에 있어요.

tests/models/mymodel/
├── __init__.py
├── test_modeling_mymodel.py            # model tests (required)
├── test_tokenization_mymodel.py        # tokenizer tests (if custom tokenizer)
├── test_image_processing_mymodel.py    # image processor tests (if vision model)
├── test_feature_extraction_mymodel.py  # feature extractor tests (if audio/speech model)
└── test_processing_mymodel.py          # processor tests (if multimodal)

토크나이저 테스트도 같은 패턴을 따라요. tests/test_tokenization_common.py에서 TokenizerTesterMixin을 상속받고 몇 가지 속성을 설정하면 자동 생성된 테스트를 얻을 수 있어요. 예시는 tests/models/llama/test_tokenization_llama.py를 참고하세요.

Config 테스트

ConfigTester는 config 클래스가 직렬화, 저장/로드, 표준 속성을 올바르게 처리하는지 검증해요. CausalLMModelTest와 VLMModelTest는 config 테스트를 자동으로 포함해요. ModelTester와 ModelTest를 쓰는 일반 경로에서는 setUp()에서 config tester를 수동으로 정의해요.

from tests.test_configuration_common import ConfigTester

def setUp(self):
    self.config_tester = ConfigTester(self, config_class=MyModelConfig, hidden_size=32)

def test_config(self):
    self.config_tester.run_common_tests()

run_common_tests()는 여러 가지 검사를 실행해요.

  • hidden_size, num_attention_heads, num_hidden_layers 같은 공통 속성(그리고 has_text_modality=True면 vocab_size)이 존재하는지 확인해요.
  • to_json_string()과 to_json_file()로 JSON 직렬화를 테스트해요.
  • save_pretrained()와 from_pretrained()를 왕복(round-trip) 검사해요.
  • id2label과 label2id의 일관성을 확인해요.
  • 인자 없이 config를 만들어 기본 초기화를 검증해요.
  • output_hidden_states 같은 공통 kwarg를 설정하고 올바르게 저장되는지 확인해요.

vocab_size가 없는 비전 전용 모델에는 has_text_modality=False를 전달하고, config 기본값을 재정의하려면 추가 **kwargs를 전달해요.

self.config_tester = ConfigTester(
    self, config_class=MyVisionConfig, has_text_modality=False, hidden_size=64
)

통합 테스트와 작은 모델 (tiny models)

믹스인 테스트는 무작위 가중치를 가진 작은 config를 사용해 모델 동작을 빠르게 검증해요. 통합 테스트는 실제 사전 훈련된 가중치로 추론을 실행해 출력 정확성을 검증해요. Hub의 작은 모델은 빠른 CI에 충분히 작지만, 실제 체크포인트처럼 구조화돼 있어요.

통합 테스트 작성하기

통합 테스트는 별도의 테스트 클래스에 두고 @slow로 표시해요. 각 테스트는 실제 가중치를 다운로드하고 추론을 실행한 다음 출력을 기대값과 대조해요. 메모리 잔여물을 피하려면 setUp과 tearDown에서 cleanup(torch_device, gc_collect=False)를 호출해요.

import torch
from transformers import AutoTokenizer
from transformers.testing_utils import cleanup, require_torch, slow, torch_device

class MyModelIntegrationTest(unittest.TestCase):
    def setUp(self):
        cleanup(torch_device, gc_collect=False)

    def tearDown(self):
        cleanup(torch_device, gc_collect=False)

    @slow
    @require_torch
    def test_inference(self):
        model = MyModelForCausalLM.from_pretrained("myorg/mymodel-base").to(torch_device)
        tokenizer = AutoTokenizer.from_pretrained("myorg/mymodel-base")
        inputs = tokenizer("Hello, world", return_tensors="pt").to(torch_device)

        with torch.no_grad():
            outputs = model(**inputs)

        # check against expected values
        expected_slice = torch.tensor([[-0.1234, 0.5678, -0.9012]])
        torch.testing.assert_close(outputs.logits[0, :1, :3], expected_slice, atol=1e-4, rtol=1e-4)

가중치를 다운로드하거나, 큰 데이터셋을 로드하거나, 몇 초 이상 걸리는 테스트는 @slow로 표시해요. pull request CI는 slow 테스트를 건너뛰지만, 야간 일정에서 실행돼요.

생성 통합 테스트

생성 테스트에서는 do_sample=False를 사용해서 출력이 실행 환경과 하드웨어에 걸쳐 결정적이도록 해요. MoE(Mixture-of-Experts) 모델의 경우, 생성 전에 model.set_experts_implementation("eager")를 호출해서 안정적인 전문가 디스패치 경로를 강제해요. 이렇게 하지 않으면 라우터의 아주 작은 수치 차이로 어떤 전문가가 토큰을 처리할지 뒤집혀 출력이 바뀔 수 있어요.

@slow
@require_torch
def test_generate(self):
    model = MyModelForCausalLM.from_pretrained("myorg/mymodel-base").to(torch_device)
    tokenizer = AutoTokenizer.from_pretrained("myorg/mymodel-base")
    inputs = tokenizer("Hello, world", return_tensors="pt").to(torch_device)

    # model.set_experts_implementation("eager")  # uncomment for MoE models
    generated_ids = model.generate(**inputs, max_new_tokens=20, do_sample=False)
    output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
    self.assertEqual(output, ["Hello, world! This is the expected continuation..."])

하드웨어별 기대값

Transformers CI는 NVIDIA A10에서 slow 테스트를 실행해요. 수치 결과는 GPU 세대에 따라 약간 달라질 수 있으므로, 통합 테스트는 Expectations 클래스를 사용해 장치별 기대값을 등록해요. Expectations는 (device_type, (major, minor)) SM 키를 기준으로 현재 하드웨어에 가장 잘 맞는 값을 고르고, 맞는 것이 없으면 기본값으로 폴백해요.

torch.cuda.get_device_capability()를 실행해 로컬 SM 버전을 출력해요 (예: A10은 (8, 6), H100은 (9, 0)).

from transformers.testing_utils import Expectations

expected_texts = Expectations(
    {
        ("cuda", (8, 6)): ["Hello, world! This is the A10 continuation..."],
        ("cuda", (9, 0)): ["Hello, world! This is the H100 continuation..."],
    }
).get_expectation()

self.assertEqual(output, expected_texts)

작은 모델 만들기

무작위 가중치를 가진 작은 모델은 hf-internal-testing 조직 아래 Hub에 있어요. 파이프라인 테스트는 Hub에 호스팅된 체크포인트가 필요하지만 출력 품질은 신경 쓰지 않을 때 작은 모델에 의존해요. 빠른 스모크 테스트(smoke test)도 큰 체크포인트를 다운로드하지 않고 forward pass 형상을 검증하기 위해 작은 모델을 로드해요.

작은 모델은 통합 테스트의 최후의 수단이에요. 사용 가능한 가장 작은 체크포인트가 약 24GB VRAM을 초과할 때만 사용하세요. 가능하면 원본 사전 훈련된 가중치를 사용해서 실제 수치 회귀(numerical regression)를 잡아내세요.

utils/create_dummy_models.py 스크립트는 ModelTester.get_config()에서 작은 모델을 생성해요. 이 스크립트는 tester에서 작은 하이퍼파라미터를 추출하고, 무작위 가중치로 모델을 만들어 결과를 Hub에 업로드해요.

로컬에서 작은 모델을 생성해요.

python utils/create_dummy_models.py output_dir -m your_model_type

Hub에 업로드해요.

python utils/create_dummy_models.py output_dir -m your_model_type --upload --organization hf-internal-testing

각 모델은 hf-internal-testing/tiny-random-{ModelClassName} 이름을 사용하며 tests/utils/tiny_model_summary.json에 기록돼요. CI 워크플로(.github/workflows/check_tiny_models.yml)는 매일 작은 모델을 다시 생성해요.

무엇이 테스트되는지 제어하기

ModelTesterMixin의 불리언 플래그가 자동 생성된 테스트를 켜고 꺼요. 테스트 클래스에서 아무 플래그나 재정의해서 특정 검사를 활성화하거나 비활성화해요.

class MyModelTest(CausalLMModelTest, unittest.TestCase):
    model_tester_class = MyModelTester
    test_resize_embeddings = False
    test_all_params_have_gradient = False  # when not all parameters are activated in every forward pass
Flag Default What it controls
test_resize_embeddings True Embedding layer resizing
test_resize_position_embeddings False Position embedding resizing
test_mismatched_shapes True Mismatched input/output shape handling
test_missing_keys True Missing key warnings on load
test_torch_exportable True torch.export compatibility
test_all_params_have_gradient True All parameters receive gradients (set False when not all parameters are activated in every forward pass, such as MoE experts)
is_encoder_decoder False Encoder-decoder specific tests
has_attentions True Attention output tests
_is_composite False Composite/multimodal model handling
model_split_percents [0.5, 0.7, 0.9] Split percentages for model parallelism tests

다음 단계

  • 테스트 선택, 픽스처, 로깅 등에 대해 더 배우려면 pytest 문서를 살펴보세요.