From a524344b78159f550b3b6b5f463f6d63c233aa73 Mon Sep 17 00:00:00 2001 From: Boian Petkantchin Date: Mon, 9 Dec 2024 15:57:33 +0000 Subject: [PATCH] Add CLI script exporting CLIP Toy model IREE test data This is required to have an easy way of exporting test data that will be used in IREE to guard against regressions. E.g. ``` python -m sharktank.models.clip.export_toy_text_model_iree_test_data \ --output-path-prefix=clip_toy_text_model ``` Refactor some of the existing tests to reuse the new export logic. --- sharktank/sharktank/models/clip/export.py | 25 ++- .../export_toy_text_model_iree_test_data.py | 29 ++++ sharktank/sharktank/models/clip/testing.py | 161 +++++++++++++++++- sharktank/sharktank/utils/io.py | 25 ++- sharktank/tests/models/clip/clip_test.py | 145 ++++++---------- 5 files changed, 280 insertions(+), 105 deletions(-) create mode 100644 sharktank/sharktank/models/clip/export_toy_text_model_iree_test_data.py diff --git a/sharktank/sharktank/models/clip/export.py b/sharktank/sharktank/models/clip/export.py index 3cae3f4c4..95dbdacad 100644 --- a/sharktank/sharktank/models/clip/export.py +++ b/sharktank/sharktank/models/clip/export.py @@ -11,8 +11,8 @@ CLIPEncoderLayer as HfCLIPEncoderLayer, CLIPEncoder as HfCLIPEncoder, ) -from os import PathLike import torch +from os import PathLike from ...types.theta import Theta, Dataset, torch_module_to_theta from ...layers.configs import ClipTextConfig @@ -50,9 +50,14 @@ def clip_text_model_to_dataset(model: ClipTextModel) -> Dataset: return Dataset(properties=model.config.to_properties(), root_theta=model.theta) +def export_clip_text_model_iree_parameters(model: ClipTextModel, output_path: PathLike): + dataset = clip_text_model_to_dataset(model) + dataset.save(output_path) + + def export_clip_text_model_dataset_from_hugging_face( - model_or_name_or_path: Union[str, PathLike, transformers.CLIPTextModel], - output_path: Union[str, PathLike], + model_or_name_or_path: Union[PathLike, transformers.CLIPTextModel], + output_path: PathLike, dtype: Optional[torch.dtype] = None, ): if isinstance(model_or_name_or_path, transformers.CLIPTextModel): @@ -99,3 +104,17 @@ def _( output = export(fxb, import_symbolic_shape_expressions=True) output.save_mlir(mlir_output_path) + + +def export_clip_text_model_to_iree( + model: ClipTextModel, + batch_sizes: list[int], + mlir_output_path: PathLike, + parameters_output_path: PathLike, +): + export_clip_text_model_iree_parameters(model, parameters_output_path) + export_clip_text_model_mlir( + model=parameters_output_path, + batch_sizes=batch_sizes, + mlir_output_path=mlir_output_path, + ) diff --git a/sharktank/sharktank/models/clip/export_toy_text_model_iree_test_data.py b/sharktank/sharktank/models/clip/export_toy_text_model_iree_test_data.py new file mode 100644 index 000000000..979bc3255 --- /dev/null +++ b/sharktank/sharktank/models/clip/export_toy_text_model_iree_test_data.py @@ -0,0 +1,29 @@ +# Copyright 2024 Advanced Micro Devices, Inc. +# +# Licensed under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +from argparse import ArgumentParser +from typing import Optional +from pathlib import Path + +from .testing import export_clip_toy_text_model_default_iree_test_data + + +def main(args: Optional[list[str]] = None): + parser = ArgumentParser( + description=( + "Export test data for toy-sized CLIP text model." + " This program MLIR, parameters sample input and expected output." + " Exports float32 and bfloat16 model variants." + " The expected output is always in float32 precision." + ) + ) + parser.add_argument("--output-dir", type=str, default=f"clip_toy_text_model") + args = parser.parse_args(args=args) + export_clip_toy_text_model_default_iree_test_data(Path(args.output_dir)) + + +if __name__ == "__main__": + main() diff --git a/sharktank/sharktank/models/clip/testing.py b/sharktank/sharktank/models/clip/testing.py index 87634c220..852da8a18 100644 --- a/sharktank/sharktank/models/clip/testing.py +++ b/sharktank/sharktank/models/clip/testing.py @@ -4,14 +4,167 @@ # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -from ...layers.configs.llm_configs import ClipTextConfig -from ...types.theta import Theta -from .export import hugging_face_clip_text_model_to_theta +import functools import torch +from os import PathLike, makedirs +from typing import Union, Optional +from copy import copy +from iree.turbine.aot.params import ParameterArchiveBuilder + +from ...layers.configs.llm_configs import ClipTextConfig +from .clip import ClipTextModel +from ...types.theta import Theta, Dataset +from ...types.tensors import dtype_to_serialized_short_name +from ...utils.io import save_tensor_as_irpa +from .export import ( + clip_text_model_to_dataset, + hugging_face_clip_text_model_to_theta, + export_clip_text_model_to_iree, +) +from ...transforms.dataset import set_float_dtype + + +def clip_toy_text_model_config(dtype: Optional[torch.dtype] = None) -> ClipTextConfig: + num_attention_heads = 5 + vocab_size = 11 + return ClipTextConfig( + vocab_size=vocab_size, + hidden_size=13 * num_attention_heads, + intermediate_size=7, + projection_dim=3, + num_attention_heads=num_attention_heads, + max_position_embeddings=17, + layer_norm_eps=1e-4, + num_hidden_layers=2, + bos_token_id=vocab_size - 2, + eos_token_id=vocab_size - 1, + dtype=dtype, + ) + + +def export_clip_toy_text_model_default_iree_test_data(output_dir: PathLike): + makedirs(output_dir, exist_ok=True) + + # We want to always export the same without interfering with RNG for the rest of + # the program. + rng_state = torch.get_rng_state() + torch.random.manual_seed(12345) + + reference_dtype = torch.float32 + target_dtypes = [torch.float32, torch.bfloat16] + target_iree_parameters_output_paths = [] + target_mlir_output_paths = [] + batch_size = 4 + for dtype in target_dtypes: + prefix = output_dir / f"{dtype_to_serialized_short_name(dtype)}" + target_iree_parameters_output_paths.append(f"{prefix}_parameters.irpa") + target_mlir_output_paths.append(f"{prefix}.mlir") + call_prefix = output_dir / f"forward_bs{batch_size}" + input_ids_output_path = f"{call_prefix}_arg0_input_ids.irpa" + expected_last_hidden_state_output_path = ( + f"{call_prefix}_expected_result0_last_hidden_state_" + f"{dtype_to_serialized_short_name(reference_dtype)}.irpa" + ) + export_clip_toy_text_model_iree_test_data( + reference_dtype=reference_dtype, + target_dtypes=target_dtypes, + batch_size=batch_size, + input_ids_output_path=input_ids_output_path, + expected_last_hidden_state_output_path=expected_last_hidden_state_output_path, + target_iree_parameters_output_paths=target_iree_parameters_output_paths, + target_mlir_output_paths=target_mlir_output_paths, + ) + + torch.set_rng_state(rng_state) + + +def export_clip_toy_text_model_iree_test_data( + reference_dtype: torch.dtype, + target_dtypes: list[torch.dtype], + batch_size: int, + target_iree_parameters_output_paths: list[PathLike], + target_mlir_output_paths: list[PathLike], + input_ids_output_path: PathLike, + expected_last_hidden_state_output_path: PathLike, +): + reference_config = clip_toy_text_model_config(reference_dtype) + input_ids = make_random_input_token_sequences( + batch_size=batch_size, config=reference_config + ) + reference_theta = make_clip_text_model_random_theta(reference_config) + reference_model = ClipTextModel(theta=reference_theta, config=reference_config) + for i, ( + target_dtype, + target_iree_parameters_output_path, + target_mlir_output_path, + ) in enumerate( + zip( + target_dtypes, + target_iree_parameters_output_paths, + target_mlir_output_paths, + strict=True, + ) + ): + current_input_ids_output_path = None + current_expected_last_hidden_state_output_path = None + if i == 0: + current_input_ids_output_path = input_ids_output_path + current_expected_last_hidden_state_output_path = ( + expected_last_hidden_state_output_path + ) + export_clip_text_model_iree_test_data( + reference_model=reference_model, + target_dtype=target_dtype, + input_ids=input_ids, + target_iree_parameters_output_path=target_iree_parameters_output_path, + target_mlir_output_path=target_mlir_output_path, + input_ids_output_path=current_input_ids_output_path, + expected_last_hidden_state_output_path=current_expected_last_hidden_state_output_path, + ) + + +def export_clip_text_model_iree_test_data( + reference_model: ClipTextModel, + target_dtype: torch.dtype, + input_ids: torch.LongTensor, + target_mlir_output_path: PathLike, + target_iree_parameters_output_path: PathLike, + input_ids_output_path: Optional[PathLike] = None, + expected_last_hidden_state_output_path: Optional[PathLike] = None, +): + batch_size = input_ids.shape[0] + reference_dataset = clip_text_model_to_dataset(reference_model) + target_config = copy(reference_model.config) + target_config.dtype = target_dtype + target_dataset = Dataset( + root_theta=reference_dataset.root_theta.transform( + functools.partial(set_float_dtype, dtype=target_dtype) + ), + properties=target_config.to_properties(), + ) + target_model = ClipTextModel(theta=target_dataset.root_theta, config=target_config) + export_clip_text_model_to_iree( + target_model, + batch_sizes=[batch_size], + mlir_output_path=target_mlir_output_path, + parameters_output_path=target_iree_parameters_output_path, + ) + + if input_ids_output_path is not None: + save_tensor_as_irpa(input_ids, input_ids_output_path) + + if expected_last_hidden_state_output_path is None: + return + + expected_last_hidden_state = reference_model(input_ids=input_ids)[ + "last_hidden_state" + ] + save_tensor_as_irpa( + expected_last_hidden_state, expected_last_hidden_state_output_path + ) def make_clip_text_model_random_theta(config: ClipTextConfig) -> Theta: - from transformers import CLIPTextConfig as HfCLIPTextConfig from transformers import CLIPTextModel as HfCLIPTextModel hf_config = config.to_hugging_face_clip_text_model_config() diff --git a/sharktank/sharktank/utils/io.py b/sharktank/sharktank/utils/io.py index ac2480846..12c9c002b 100644 --- a/sharktank/sharktank/utils/io.py +++ b/sharktank/sharktank/utils/io.py @@ -5,10 +5,10 @@ # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception from pathlib import Path +import torch +from os import PathLike -from iree.turbine.aot import ( - ParameterArchiveBuilder, -) +from iree.turbine.aot import ParameterArchiveBuilder, ParameterArchive class ShardedArchiveBuilder(ParameterArchiveBuilder): @@ -49,3 +49,22 @@ def path_for_rank(path: Path, rank: int): /tmp/foobar.rank0.irpa """ return path.with_suffix(f".rank{rank}{path.suffix}") + + +def save_tensor_as_irpa(tensor: torch.Tensor, path: PathLike): + """Save a single tensor into an IRPA file.""" + param_builder = ParameterArchiveBuilder() + param_builder.add_tensor("", tensor) + param_builder.save(path) + + +def load_irpa_as_tensor(tensor: torch.Tensor, path: PathLike, **kwargs): + """Load a tensor form an IRPA file that holds only one tensor.""" + params = ParameterArchive(path, **kwargs) + items = params.items() + if len(items) != 1: + raise ValueError( + f'Too many items {len(items)} in IRPA file "{path}".' + " Only a single tensor was expected." + ) + return items[0][1].as_tensor() diff --git a/sharktank/tests/models/clip/clip_test.py b/sharktank/tests/models/clip/clip_test.py index 99af4ba6f..704333c90 100644 --- a/sharktank/tests/models/clip/clip_test.py +++ b/sharktank/tests/models/clip/clip_test.py @@ -8,6 +8,7 @@ import functools import iree.compiler import os +from pathlib import Path from parameterized import parameterized from copy import copy import pytest @@ -47,18 +48,18 @@ test_prompts, ) from sharktank.models.clip.export import ( - export_clip_text_model_mlir, export_clip_text_model_dataset_from_hugging_face, hugging_face_clip_attention_to_theta, hugging_face_clip_encoder_layer_to_theta, hugging_face_clip_encoder_to_theta, hugging_face_clip_text_model_to_dataset, hugging_face_clip_text_model_to_theta, - clip_text_model_to_dataset, ) from sharktank.models.clip.testing import ( make_random_input_token_sequences, make_clip_text_model_random_theta, + export_clip_text_model_iree_test_data, + clip_toy_text_model_config, ) from sharktank.models.clip import ( ClipAttention, @@ -72,13 +73,15 @@ with_clip_data = pytest.mark.skipif("not config.getoption('with_clip_data')") -@pytest.mark.usefixtures("caching", "path_prefix") +@pytest.mark.usefixtures("path_prefix") class ClipTextIreeTest(TempDirTestBase): def setUp(self): super().setUp() torch.random.manual_seed(12345) if self.path_prefix is None: - self.path_prefix = f"{self._temp_dir}/" + self.path_prefix = self._temp_dir + else: + self.path_prefix = Path(self.path_prefix) @with_clip_data def testSmokeExportLargeF32FromHuggingFace(self): @@ -90,12 +93,20 @@ def testSmokeExportLargeF32FromHuggingFace(self): huggingface_repo_id, ).download() target_dtype_name = dtype_to_serialized_short_name(torch.float32) - target_model_path_prefix = f"{self.path_prefix}{huggingface_repo_id_as_path}_text_model_{target_dtype_name}" + target_model_path_prefix = ( + self.path_prefix + / f"{huggingface_repo_id_as_path}_text_model_{target_dtype_name}" + ) output_path = f"{target_model_path_prefix}.irpa" export_clip_text_model_dataset_from_hugging_face( huggingface_repo_id, output_path ) + def testSmokeExportToyIreeTestData(self): + from sharktank.models.clip.export_toy_text_model_iree_test_data import main + + main([f"--output-dir={self.path_prefix/'clip_toy_text_model'}"]) + @with_clip_data def testCompareLargeIreeF32AgainstTorchEagerF32(self): self.runTestCompareIreeAgainstPretrainedTorchEager( @@ -141,43 +152,31 @@ def runTestCompareIreeAgainstTorchEagerWithInputTokens( ) target_dtype_name = dtype_to_serialized_short_name(target_dtype) reference_model_path_prefix = ( - f"{self.path_prefix}{file_artifact_prefix_name}_{reference_dtype_name}" + self.path_prefix / f"{file_artifact_prefix_name}_{reference_dtype_name}" ) target_model_path_prefix = ( - f"{self.path_prefix}{file_artifact_prefix_name}_{target_dtype_name}" - ) - - target_config = copy(reference_model.config) - target_config.dtype = target_dtype - reference_dataset = clip_text_model_to_dataset(reference_model) - target_dataset = Dataset( - root_theta=reference_dataset.root_theta.transform( - functools.partial(set_float_dtype, dtype=target_config.dtype) - ), - properties=target_config.to_properties(), + self.path_prefix / f"{file_artifact_prefix_name}_{target_dtype_name}" ) parameters_path = f"{target_model_path_prefix}.irpa" - if not self.caching or not os.path.exists(parameters_path): - target_dataset.save(parameters_path) - - dataset = Dataset.load(parameters_path) - target_config = ClipTextConfig.from_properties(dataset.properties) input_args = OrderedDict([("input_ids", input_ids)]) batch_size = input_ids.shape[0] - mlir_path = f"{target_model_path_prefix}.mlir" - if not self.caching or not os.path.exists(mlir_path): - export_clip_text_model_mlir( - parameters_path, batch_sizes=[batch_size], mlir_output_path=mlir_path - ) + + export_clip_text_model_iree_test_data( + reference_model=reference_model, + target_dtype=target_dtype, + input_ids=input_ids, + target_mlir_output_path=mlir_path, + target_iree_parameters_output_path=parameters_path, + ) + iree_module_path = f"{target_model_path_prefix}.vmfb" - if not self.caching or not os.path.exists(iree_module_path): - iree.compiler.compile_file( - mlir_path, - output_file=iree_module_path, - extra_args=["--iree-hal-target-device=hip", "--iree-hip-target=gfx942"], - ) + iree.compiler.compile_file( + mlir_path, + output_file=iree_module_path, + extra_args=["--iree-hal-target-device=hip", "--iree-hip-target=gfx942"], + ) reference_result_dict = call_torch_module_function( module=reference_model, @@ -211,11 +210,11 @@ def runTestCompareIreeAgainstTorchEagerWithInputTokens( for i in range(len(expected_outputs)) ] - actual_last_hidden_states = actual_outputs[0] - expected_last_hidden_states = expected_outputs[0] + actual_last_hidden_state = actual_outputs[0] + expected_last_hidden_state = expected_outputs[0] assert_text_encoder_state_close( - actual_last_hidden_states, expected_last_hidden_states, atol + actual_last_hidden_state, expected_last_hidden_state, atol ) def runTestCompareRandomModelIreeAgainstTorch( @@ -243,21 +242,7 @@ def runTestCompareToyModelIreeAgainstTorch( self, reference_dtype: torch.dtype, target_dtype: torch.dtype, atol: float ): batch_size = 4 - num_attention_heads = 5 - vocab_size = 11 - reference_config = ClipTextConfig( - vocab_size=vocab_size, - hidden_size=13 * num_attention_heads, - intermediate_size=7, - projection_dim=3, - num_attention_heads=num_attention_heads, - max_position_embeddings=17, - layer_norm_eps=1e-4, - num_hidden_layers=2, - bos_token_id=vocab_size - 2, - eos_token_id=vocab_size - 1, - dtype=reference_dtype, - ) + reference_config = clip_toy_text_model_config(reference_dtype) file_artifact_prefix_name = "clip_text_model_toy" self.runTestCompareRandomModelIreeAgainstTorch( reference_config=reference_config, @@ -404,21 +389,9 @@ def testCompareEagerToySizedModelAgainstTransformers( ): torch.set_default_dtype(reference_dtype) batch_size = 19 - tgt_len = 23 - num_attention_heads = 5 vocab_size = 11 - reference_config = transformers.CLIPTextConfig( - vocab_size=vocab_size, - hidden_size=13 * num_attention_heads, - intermediate_size=7, - projection_dim=3, - num_attention_heads=num_attention_heads, - layer_norm_eps=1e-4, - num_hidden_layers=2, - final_layer_norm=1e-3, - bos_token_id=vocab_size - 2, - eos_token_id=vocab_size - 1, - ) + config = clip_toy_text_model_config() + reference_config = config.to_hugging_face_clip_text_model_config() reference_model = HfCLIPTextModel( reference_config, ) @@ -432,7 +405,9 @@ def testCompareEagerToySizedModelAgainstTransformers( ) model = ClipTextModel(theta, config) - input_ids = torch.randint(low=0, high=vocab_size, size=[batch_size, tgt_len]) + input_ids = torch.randint( + low=0, high=vocab_size, size=[batch_size, config.max_position_embeddings] + ) expected_outputs = reference_model(input_ids=input_ids) @@ -471,16 +446,10 @@ def testCompareEagerToySizedModelAgainstTransformers( ): torch.set_default_dtype(reference_dtype) batch_size = 19 - tgt_len = 23 + config = clip_toy_text_model_config() + reference_config = config.to_hugging_face_clip_text_model_config() + tgt_len = config.max_position_embeddings src_len = tgt_len - num_attention_heads = 2 - reference_config = transformers.CLIPTextConfig( - vocab_size=11, - hidden_size=13 * num_attention_heads, - intermediate_size=7, - projection_dim=3, - num_attention_heads=num_attention_heads, - ) reference_model = HfCLIPAttention( reference_config, ) @@ -495,7 +464,7 @@ def testCompareEagerToySizedModelAgainstTransformers( model = ClipAttention(theta, config) reference_hidden_states = make_rand_torch( - shape=[batch_size, tgt_len, reference_config.hidden_size], + shape=[batch_size, tgt_len, config.hidden_size], dtype=reference_dtype, ) reference_attention_mask = make_random_mask( @@ -551,17 +520,10 @@ def testCompareEagerToySizedModelAgainstTransformers( ): torch.set_default_dtype(reference_dtype) batch_size = 19 - tgt_len = 23 + config = clip_toy_text_model_config() + reference_config = config.to_hugging_face_clip_text_model_config() + tgt_len = config.max_position_embeddings src_len = tgt_len - num_attention_heads = 2 - reference_config = transformers.CLIPTextConfig( - vocab_size=11, - hidden_size=13 * num_attention_heads, - intermediate_size=7, - projection_dim=3, - num_attention_heads=num_attention_heads, - layer_norm_eps=1e-4, - ) reference_model = HfCLIPEncoderLayer( reference_config, ) @@ -634,15 +596,8 @@ def testCompareEagerToySizedModelAgainstTransformers( batch_size = 19 tgt_len = 23 src_len = tgt_len - num_attention_heads = 5 - reference_config = transformers.CLIPTextConfig( - vocab_size=11, - hidden_size=13 * num_attention_heads, - intermediate_size=7, - projection_dim=3, - num_attention_heads=num_attention_heads, - layer_norm_eps=1e-4, - num_hidden_layers=2, + reference_config = ( + clip_toy_text_model_config().to_hugging_face_clip_text_model_config() ) reference_model = HfCLIPEncoder( reference_config,