# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Testing suite for the PyTorch Step3p7 model.""" import re import unittest from unittest.mock import patch from transformers import is_torch_available from transformers.conversion_mapping import get_model_conversion_mapping from transformers.models.step3p7.configuration_step3p7 import ( Step3p7Config, Step3p7TextConfig, Step3p7VisionConfig, ) from transformers.testing_utils import ( cleanup, require_torch, slow, torch_device, ) from ... import test_modeling_common from ...vlm_tester import VLMModelTest, VLMModelTester if is_torch_available(): import torch from transformers import Step3p7ForConditionalGeneration, Step3p7Model _REAL_CHECKPOINT = "stepfun-ai/Step-3.7-Flash" # FP8-quantized release: ~200GB on disk (vs. ~400GB for the bf16 checkpoint above), and its # `config.json` carries a native `quant_method: fp8` block so `from_pretrained` dequantizes/ # dispatches it automatically, with no hand-built `quantization_config` needed. _REAL_CHECKPOINT_FP8 = "stepfun-ai/Step-3.7-Flash-FP8" # Vision: image_size=16, patch_size=4 → 4×4=16 patches → after 2×stride-2 downsampler → 1×1=1 token per image. # The projector maps vision_hidden_size*4 (=32) → text_hidden_size (=16). _NUM_IMAGE_TOKENS = 1 # tokens per image after the vision downsampler class Step3p7VisionText2TextModelTester(VLMModelTester): base_model_class = Step3p7Model if is_torch_available() else None config_class = Step3p7Config conditional_generation_class = Step3p7ForConditionalGeneration if is_torch_available() else None text_config_class = Step3p7TextConfig vision_config_class = Step3p7VisionConfig def __init__(self, parent, **kwargs): # Vision downsampler reduces (image_size/patch_size)^2 → (image_size/patch_size/4)^2 # For image_size=16, patch_size=4: 16 patches → 1 token after 2×stride-2 conv kwargs.setdefault("num_image_tokens", _NUM_IMAGE_TOKENS) kwargs.setdefault("image_token_id", 4) kwargs.setdefault("image_size", 16) kwargs.setdefault("patch_size", 4) kwargs.setdefault("num_hidden_layers", 2) kwargs.setdefault("hidden_size", 16) kwargs.setdefault("intermediate_size", 37) kwargs.setdefault("num_attention_heads", 2) kwargs.setdefault("num_key_value_heads", 1) kwargs.setdefault("head_dim", 8) kwargs.setdefault("max_position_embeddings", 64) kwargs.setdefault("pad_token_id", 1) kwargs.setdefault("bos_token_id", 0) kwargs.setdefault("eos_token_id", 2) kwargs.setdefault("moe_intermediate_size", 8) kwargs.setdefault("n_routed_experts", 4) kwargs.setdefault("num_experts_per_tok", 2) kwargs.setdefault("share_expert_dim", 8) # layer_types is required (Step3p7Attention accesses it by index) kwargs.setdefault("layer_types", ["full_attention", "full_attention"]) # mlp_layer_types default heuristic (MoE from layer index 3 onward) never fires for a # 2-layer model; set explicitly so at least one layer builds a real `experts` submodule # (needed for `base_model_tp_plan`'s `mlp.experts.*` entries to match a real parameter). kwargs.setdefault("mlp_layer_types", ["dense", "sparse"]) # sliding_window required by create_sliding_window_causal_mask even when no sliding layers are used kwargs.setdefault("sliding_window", 64) super().__init__(parent, **kwargs) def get_vision_config(self): return self.vision_config_class( num_hidden_layers=1, hidden_size=8, num_attention_heads=2, num_channels=self.num_channels, image_size=self.image_size, patch_size=self.patch_size, mlp_ratio=1.0, ) def get_text_config(self): return self.text_config_class( vocab_size=self.vocab_size, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size, num_attention_heads=self.num_attention_heads, num_key_value_heads=self.num_key_value_heads, head_dim=self.head_dim, num_hidden_layers=self.num_hidden_layers, max_position_embeddings=self.max_position_embeddings, pad_token_id=self.pad_token_id, bos_token_id=self.bos_token_id, eos_token_id=self.eos_token_id, moe_intermediate_size=self.moe_intermediate_size, n_routed_experts=self.n_routed_experts, num_experts_per_tok=self.num_experts_per_tok, share_expert_dim=self.share_expert_dim, layer_types=self.layer_types, mlp_layer_types=self.mlp_layer_types, sliding_window=self.sliding_window, ) @require_torch class Step3p7ModelTest(VLMModelTest, unittest.TestCase): model_tester_class = Step3p7VisionText2TextModelTester # Vision encoder outputs hidden_size*4 channels after the stride-2 conv downsampler, # so last_hidden_state.shape[-1] != vision_config.hidden_size skip_test_image_features_output_shape = True # Training tests: only Step3p7ForConditionalGeneration has a loss head. # Step3p7Model is a backbone (returns BaseModelOutputWithPast, no .loss). def _for_cond_gen_only(self, fn): orig = self.all_model_classes self.all_model_classes = (Step3p7ForConditionalGeneration,) if is_torch_available() else () try: fn() finally: self.all_model_classes = orig def test_reverse_loading_mapping(self, check_keys_were_modified=True, skip_base_model=True): # The `vit_large_projector` -> `model.multi_modal_projector` mapping carries the base-model # prefix, so it's only visible on the model-with-head; skip the base-model check. # `.mlp.experts.down_proj_scale_inv` only matches a real FP8 checkpoint's `weight_scale_inv` # keys, never this bf16-style test config, so drop it before running the shared assertion. def get_model_conversion_mapping_without_fp8_scale(*args, **kwargs): dropped_targets = {".mlp.experts.down_proj_scale_inv"} return [ conversion for conversion in get_model_conversion_mapping(*args, **kwargs) if not dropped_targets.intersection(conversion._original_target_patterns) ] with patch.object( test_modeling_common, "get_model_conversion_mapping", get_model_conversion_mapping_without_fp8_scale, ): super().test_reverse_loading_mapping( check_keys_were_modified=check_keys_were_modified, skip_base_model=skip_base_model ) def test_training(self): self._for_cond_gen_only(super().test_training) def test_training_gradient_checkpointing(self): self._for_cond_gen_only(super().test_training_gradient_checkpointing) def test_training_gradient_checkpointing_use_reentrant_false(self): self._for_cond_gen_only(super().test_training_gradient_checkpointing_use_reentrant_false) def test_training_gradient_checkpointing_use_reentrant_true(self): self._for_cond_gen_only(super().test_training_gradient_checkpointing_use_reentrant_true) def _image_features_get_expected_num_hidden_states(self, model_tester=None): # Vision model has its own num_hidden_layers; the base class would use the # text num_hidden_layers because vision_config is an object, not a dict. if model_tester is None: model_tester = self.model_tester return model_tester.get_vision_config().num_hidden_layers + 1 def _image_features_get_expected_num_attentions(self, model_tester=None): # Same reasoning as `_image_features_get_expected_num_hidden_states` above. if model_tester is None: model_tester = self.model_tester return model_tester.get_vision_config().num_hidden_layers @require_torch @slow class Step3p7ConversionMappingIntegrationTest(unittest.TestCase): """Validates the conversion mapping against the real checkpoint's `config.json` and safetensors headers (via `huggingface_hub.get_safetensors_metadata`), without downloading any weights.""" checkpoint = _REAL_CHECKPOINT def test_real_checkpoint_config(self): config = Step3p7Config.from_pretrained(self.checkpoint) self.assertEqual(config.model_type, "step3p7") text_config = config.text_config self.assertGreater(text_config.num_hidden_layers, 0) self.assertEqual(len(text_config.layer_types), text_config.num_hidden_layers) self.assertEqual(len(text_config.mlp_layer_types), text_config.num_hidden_layers) self.assertIn("sparse", text_config.mlp_layer_types) self.assertIn("sliding_attention", text_config.layer_types) # Real checkpoint uses a different head count for sliding vs. full-attention layers # (attention_other_setting) — this is the field `Step3p7Attention` reads to rebuild # q_proj/o_proj per layer type; regression-tested numerically below via real shapes. self.assertIsNotNone(text_config.num_sliding_attention_heads) self.assertNotEqual( text_config.num_sliding_attention_heads, text_config._getattr_without_heterogeneous_validation("num_attention_heads"), ) def test_real_checkpoint_weight_mapping_is_complete(self): """Every real-checkpoint weight key must rename to a key in our model (shape-checked for simple renames), except the checkpoint's trailing MTP layers, which this implementation deliberately doesn't model (see `Step3p7TextConfig`'s `num_nextn_predict_layers` docstring).""" import torch from huggingface_hub import get_safetensors_metadata from transformers.conversion_mapping import get_model_conversion_mapping from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key config = Step3p7Config.from_pretrained(self.checkpoint) with torch.device("meta"): model = Step3p7ForConditionalGeneration(config) meta_state_dict = model.state_dict() conversions = get_model_conversion_mapping(model) renamings = [c for c in conversions if isinstance(c, WeightRenaming)] converters = [c for c in conversions if isinstance(c, WeightConverter)] real_metadata = get_safetensors_metadata(self.checkpoint) real_shapes = { key: tuple(tensor_info.shape) for file_metadata in real_metadata.files_metadata.values() for key, tensor_info in file_metadata.tensors.items() } self.assertEqual(set(real_shapes), set(real_metadata.weight_map)) num_modeled_layers = config.text_config.num_hidden_layers shape_mismatches, unexpected_unmapped = [], [] matched_simple_renames = 0 for key, real_shape in real_shapes.items(): renamed_key, matched_converter_pattern = rename_source_key( key, renamings, converters, model.base_model_prefix, meta_state_dict ) target_key = renamed_key if renamed_key in meta_state_dict else key if target_key not in meta_state_dict: layer_match = re.search(r"model\.layers\.(\d+)\.", key) if layer_match and int(layer_match.group(1)) >= num_modeled_layers: continue # expected: trailing MTP layer this implementation doesn't model unexpected_unmapped.append(key) continue if matched_converter_pattern is not None: continue # Chunk/Concatenate: shape isn't directly comparable to the source tensor meta_shape = tuple(meta_state_dict[target_key].shape) if meta_shape != real_shape: shape_mismatches.append((key, target_key, real_shape, meta_shape)) else: matched_simple_renames += 1 self.assertEqual(unexpected_unmapped, [], f"Real checkpoint keys with no mapping: {unexpected_unmapped}") self.assertEqual(shape_mismatches, [], f"Shape mismatches (checkpoint, ours): {shape_mismatches}") # Sanity floor so a mapping that accidentally matches nothing doesn't slip through as "0 == 0". self.assertGreater(matched_simple_renames, 1000) @require_torch class Step3p7IntegrationTest(unittest.TestCase): model_id = "hf-internal-testing/tiny-random-step3p7" def tearDown(self): cleanup(torch_device, gc_collect=True) def _load_model(self): model = Step3p7ForConditionalGeneration.from_pretrained(self.model_id, dtype="float32", device_map="cpu") model.eval() return model @slow def test_text_generation(self): model = self._load_model() torch.manual_seed(0) input_ids = torch.randint(0, model.config.text_config.vocab_size - 1, (1, 8), device=model.device) with torch.no_grad(): output = model.generate(input_ids=input_ids, max_new_tokens=24, do_sample=False) generated_ids = output[0, input_ids.shape[1] :] EXPECTED_OUTPUT_TOKEN_IDS = [41, 56, 35, 26, 53, 63, 56, 35, 13, 39, 1, 37, 4, 37, 4, 23, 2] # fmt: skip self.assertEqual(generated_ids.tolist(), EXPECTED_OUTPUT_TOKEN_IDS) @slow def test_image_and_text_generation(self): model = self._load_model() # `input_ids[0, 3] = image_token_id` gives exactly one placeholder token, matching the # single merged image feature `get_image_features` returns for `num_local_patches=[0]` # (no local patches, just the main image). torch.manual_seed(1) input_ids = torch.randint(0, model.config.text_config.vocab_size - 1, (1, 8), device=model.device) input_ids[0, 3] = model.config.image_token_id pixel_values = torch.randn( 1, model.config.vision_config.num_channels, model.config.vision_config.image_size, model.config.vision_config.image_size, device=model.device, dtype=model.dtype, ) with torch.no_grad(): output = model.generate( input_ids=input_ids, pixel_values=pixel_values, num_local_patches=[0], max_new_tokens=24, do_sample=False, ) generated_ids = output[0, input_ids.shape[1] :] EXPECTED_OUTPUT_TOKEN_IDS = [37, 1, 10, 41, 33, 51, 46, 61, 51, 26, 8, 23, 38, 61, 36, 17, 51, 33, 38, 23, 2] # fmt: skip self.assertEqual(generated_ids.tolist(), EXPECTED_OUTPUT_TOKEN_IDS)