diff --git a/.dockerignore b/.dockerignore index f5e96dbfa..79dff1922 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1 +1,4 @@ -venv \ No newline at end of file +* +!dist +!README.md +!LICENSE diff --git a/.github/workflows/build-image.yml b/.github/workflows/build-image.yml new file mode 100644 index 000000000..6437ee2e6 --- /dev/null +++ b/.github/workflows/build-image.yml @@ -0,0 +1,32 @@ +on: + push: + branches: [main] + paths: + - "caikit_nlp" + - "README.md" + - "pyproject.toml" + - "Dockerfile" + + pull_request: + +name: Build Image + +jobs: + build-image: + name: Build Image + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - uses: actions/setup-python@v4 + with: + python-version: 3.9 + - name: Setup tox + run: | + pip install -U pip wheel + pip install tox + - name: Build wheel + run: | + tox -e build + - name: Build image + run: | + docker build -t caikit-nlp:latest . diff --git a/.github/workflows/publish-library.yml b/.github/workflows/publish-library.yml index ee94d5f39..9574feaf5 100644 --- a/.github/workflows/publish-library.yml +++ b/.github/workflows/publish-library.yml @@ -25,12 +25,11 @@ jobs: - uses: actions/checkout@v3 - name: Set up Python uses: actions/setup-python@v3 - - name: Release - env: - FLIT_PASSWORD: ${{ secrets.PYPI_TOKEN }} + - name: Build and check package run: | - pip install tox - RELEASE_VERSION=${GITHUB_REF#refs/*/} - RELEASE_VERSION=${RELEASE_VERSION#v*} - sed -i "s/^version = .*/version = \"${RELEASE_VERSION}\"/" pyproject.toml - tox -e publish + tox -e build,twinecheck + - name: Upload package + if: github.event_name == 'release' + uses: pypa/gh-action-pypi-publish@release/v1 + with: + password: ${{ secrets.PYPI_TOKEN }} diff --git a/.gitignore b/.gitignore index 4ff0e0cbb..d855873a1 100644 --- a/.gitignore +++ b/.gitignore @@ -166,3 +166,4 @@ prompt_prefixes sample_prompt transformers_cache generated_interfaces +/caikit_nlp/_version.py diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 000000000..be46ec603 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,45 @@ +FROM registry.access.redhat.com/ubi8/ubi-minimal:latest as builder + +RUN microdnf update -y && \ + microdnf install -y \ + git python39-pip && \ + pip3 install --upgrade --no-cache-dir pip && \ + microdnf clean all + +RUN python3 -m venv /opt/caikit/ + +ENV VIRTUAL_ENV=/opt/caikit +ENV PATH="$VIRTUAL_ENV/bin:$PATH" + +COPY dist/caikit_nlp*.whl /tmp/ +RUN pip install --no-cache /tmp/caikit_nlp*.whl && rm /tmp/caikit_nlp*.whl + + +FROM registry.access.redhat.com/ubi8/ubi-minimal:latest as deploy + +RUN microdnf update -y && \ + microdnf install -y \ + shadow-utils python39 && \ + microdnf clean all + +COPY --from=builder /opt/caikit /opt/caikit +COPY LICENSE /opt/caikit/ +COPY README.md /opt/caikit/ + +RUN groupadd --system caikit --gid 1001 && \ + adduser --system --uid 1001 --gid 0 --groups caikit \ + --home-dir /caikit --shell /sbin/nologin \ + --comment "Caikit User" caikit + +ENV VIRTUAL_ENV=/opt/caikit +ENV PATH="$VIRTUAL_ENV/bin:$PATH" + +USER caikit + +ENV RUNTIME_LIBRARY=caikit_nlp +# Optional: use `CONFIG_FILES` and the /caikit/ volume to explicitly provide a configuration file and models +# ENV CONFIG_FILES=/caikit/caikit.yml +VOLUME ["/caikit/"] +WORKDIR /caikit + +CMD ["python"] diff --git a/README.md b/README.md index f58882712..aab338512 100644 --- a/README.md +++ b/README.md @@ -16,20 +16,142 @@ Capabilities provided by `caikit-nlp`: | Tokenization | 1. `RegexSentenceSplitter` | 1. Demo purposes only | | Embedding | [COMING SOON] | [COMING SOON] | -### Getting Started +## Getting Started + +### Notebooks To help you quickly get started with using Caikit, we have prepared a [Jupyter notebook](examples/Caikit_Getting_Started.ipynb) that can be run in Google Colab. Caikit-nlp is a powerful library that leverages prompt tuning and fine-tuning to add NLP domain capabilities to caikit. +### Installation + +To install from git repo: + +```bash +python -m venv .venv +source .venv/bin/activate +pip install git+https://github.com/caikit/caikit-nlp +``` + +### Bootstrapping models + +`caikit_nlp` can use Hugging Face models, allowing for direct download and bootstrapping. + +For example, to use [google/flan-t5-small](https://huggingface.co/google/flan-t5-small): + +```python +import os +# The env var ALLOW_DOWNLOADS has to be set to allow model downloads before importing caikit_nlp +os.environ['ALLOW_DOWNLOADS'] = "1" + +import caikit_nlp + +model_name = "google/flan-t5-small" +model = caikit_nlp.text_generation.TextGeneration.bootstrap(model_name) +model.save(f"{model_name}-caikit") # optionally save the model +``` + +### Serving models + +To serve models, the following basic configuration can be used: + +```yaml +# config.yml +runtime: + library: caikit_nlp + local_models_dir: ./models + +log: + formatter: pretty # optional: log formatter is set to json by default +``` + +Start the server: + +```bash +env CONFIG_FILES=./config.yml python -m caikit.runtime +``` + +The model can now be queried at `localhost:8080` via http or at `localhost:8085` via grpc. + +For example, using the http server and using curl to send a POST request: + +```bash +curl --json '{ + "model_id": "flan-t5-small-caikit", + "inputs": "At what temperature does liquid Nitrogen boil?" +}' localhost:8080/api/v1/task/text-generation +``` + +We get the following response: + +```json +{ + "generated_text": "74 degrees F", + "generated_tokens": 5, + "finish_reason": "MAX_TOKENS", + "producer_id": { + "name": "Text Generation", + "version": "0.1.0" + }, + "input_token_count": 10, + "seed": null +} +``` + +All the available API endpoints and protos can be dumped using [`scripts/dump_apis.sh`](/scripts/dump_apis.sh). + +### Docker + +To build the docker image: + +```bash +python -m build --wheel +docker build -t caikit-nlp:latest . +``` + +A volume can be mounted at `/caikit` providing configuration and (optionally) models: + +```bash +mkdir -p caikit +$EDITOR caikit/config.yml # edit as required +cp -r ./caikit/models +docker run -e CONFIG_FILES=/caikit/config.yml -v $PWD/caikit/:/caikit -p 8080:8080 -p 8085:8085 python -m caikit.runtime +``` + +#### Serving with containers + +In order to start the serving runtime: + +```bash +docker run -e CONFIG_FILES=/caikit/config.yml \ + -v $PWD/caikit/:/caikit -p 8080:8080 -p 8085 \ + python -m caikit.runtime +``` + +Assuming the standard configuration with port `8080` for the http server and `8085` for the grpc server. + +### Configuration + +Configuration can be provided via environment variables or by providing a yaml configuration file thanks to [`alchemy-config`](https://github.com/IBM/alchemy-config). + +For example, to set the caikit runtime, setting `RUNTIME_LIBRARY=caikit_nlp` via environment variables or providing the following yaml configuration is equivalent. + +```yaml +# config.yml +runtime: + library: caikit_nlp +``` + +For configuration options see `caikit_nlp`'s example config: [`config.yml`](/caikit_nlp/config/config.yml) or `caikit`'s example [`caikit.yml`](https://github.com/caikit/caikit/blob/main/caikit/config/config.yml). -### Contributing +## Contributing We welcome contributions from the community! If you would like to contribute to `caikit-nlp`, please read the guidelines in the main project's [CONTRIBUTING.md](CONTRIBUTING.md) file. It includes information on submitting bug reports, feature requests, and pull requests. Make sure to follow our coding standards, [code of conduct](code-of-conduct.md), [security standards](https://github.com/caikit/community/blob/main/SECURITY.md), and documentation guidelines to streamline the contribution process. -### License +## License This project is licensed under the [ASFv2 License](LICENSE). -### Glossary +## Glossary A list of terms that either may be unfamiliar or that have nebulous definitions based on who and where you hear them, defined for how they are used/thought of in the `caikit`/`caikit-nlp` project: @@ -41,7 +163,7 @@ Prompt tuning - learning soft prompts. This is different from prompt engineering The important difference between fine tuning and capabilities like prompt tuning/multi-taskprompt tuning is that the latter doesn't change the base model's weights at all. So when you run inference for prompt tuned models, you can have n prompts to 1 base model, and just inject the prompt tensors you need when they're requested instead of having _n_ separate fine-tuned models. -### Runtime Performance Benchmarking +## Runtime Performance Benchmarking [Runtime Performance Benchmarking](./benchmarks/README.md) for tuning various models. diff --git a/caikit_nlp/__init__.py b/caikit_nlp/__init__.py index 0c9c020f6..4f71857a9 100644 --- a/caikit_nlp/__init__.py +++ b/caikit_nlp/__init__.py @@ -29,6 +29,7 @@ from .data_model import * from .modules import * from .resources import * +from .version import __version__, __version_tuple__ # Configure the library with library-specific configuration file CONFIG_PATH = os.path.realpath( diff --git a/caikit_nlp/config/config.yml b/caikit_nlp/config/config.yml index a0e41ac2e..92a72d9ba 100644 --- a/caikit_nlp/config/config.yml +++ b/caikit_nlp/config/config.yml @@ -21,11 +21,20 @@ torch_dtype: float32 # Path of folder that will contain all the source prompts source_prompt_base: "" +# Path for searching base models from +base_models_dir: "" + # Whether or not to purge TGIS prompts on model deletion unload_tgis_prompt_artifacts: false # Torchrun elastic launch configuration, e.g., for fine tuning on multiple GPUs master_addr: localhost master_port: 29550 +training_data_limit: + __default__: -1 + # Configuration for PeftPromptTuning module + 6655831b-960a-4dc5-8df4-867026e2cd41: + add_model_name_here: 10000 + runtime: library: caikit_nlp diff --git a/caikit_nlp/modules/text_classification/sequence_classification.py b/caikit_nlp/modules/text_classification/sequence_classification.py index 8a88f489f..a485c59b1 100644 --- a/caikit_nlp/modules/text_classification/sequence_classification.py +++ b/caikit_nlp/modules/text_classification/sequence_classification.py @@ -21,8 +21,8 @@ import torch # First Party +from caikit.core.exceptions import error_handler from caikit.core.modules import ModuleBase, ModuleLoader, ModuleSaver, module -from caikit.core.toolkit import error_handler from caikit.interfaces.nlp.data_model import ClassificationResult, ClassificationResults from caikit.interfaces.nlp.tasks import TextClassificationTask import alog diff --git a/caikit_nlp/modules/text_generation/peft_config.py b/caikit_nlp/modules/text_generation/peft_config.py index 57c9fc581..9222a1d01 100644 --- a/caikit_nlp/modules/text_generation/peft_config.py +++ b/caikit_nlp/modules/text_generation/peft_config.py @@ -15,6 +15,7 @@ # Standard from enum import Enum import os +import re # Third Party from peft import MultitaskPromptTuningInit @@ -44,6 +45,8 @@ log = alog.use_channel("PFT_CNFG_TLKT") error = error_handler.get(log) +SOURCE_DIR_VALIDATION_REGEX = re.compile(r"^[-a-zA-Z_0-9\/]+") + class TuningType(str, Enum): PROMPT_TUNING = "PROMPT_TUNING" @@ -56,6 +59,20 @@ class TuningType(str, Enum): def resolve_base_model(base_model, cls, torch_dtype): if isinstance(base_model, str): + + error.value_check( + "", + re.fullmatch(SOURCE_DIR_VALIDATION_REGEX, base_model), + "invalid characters in base_model name", + ) + if get_config().base_models_dir: + + base_model_full_path = os.path.join( + get_config().base_models_dir, base_model + ) + if os.path.exists(base_model_full_path): + base_model = base_model_full_path + model_config = AutoConfig.from_pretrained( base_model, local_files_only=not get_config().allow_downloads ) @@ -200,4 +217,4 @@ def get_peft_config( output_model_types=output_model_types, ) - return task_type, output_model_types, peft_config, tuning_type + return task_type, output_model_types, peft_config, tuning_type \ No newline at end of file diff --git a/caikit_nlp/modules/text_generation/peft_prompt_tuning.py b/caikit_nlp/modules/text_generation/peft_prompt_tuning.py index 599571d7f..200d60370 100644 --- a/caikit_nlp/modules/text_generation/peft_prompt_tuning.py +++ b/caikit_nlp/modules/text_generation/peft_prompt_tuning.py @@ -13,14 +13,15 @@ # limitations under the License. """This module contains prompt tuning through PEFT""" # Standard -from datetime import datetime from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union import gc import json import os +import tempfile # Third Party -from accelerate import Accelerator +from datasets import Dataset +from datasets import IterableDataset as TransformersIterableDataset from peft import ( MultitaskPromptTuningConfig, PeftConfig, @@ -30,23 +31,17 @@ TaskType, get_peft_model, ) -from torch.optim import AdamW -from torch.utils.data import DataLoader -from tqdm import tqdm -from transformers import ( - AutoModelForCausalLM, - DataCollatorForLanguageModeling, - default_data_collator, -) +from transformers import AutoModelForCausalLM, default_data_collator from transformers.models.auto.tokenization_auto import AutoTokenizer -from transformers.optimization import get_linear_schedule_with_warmup import numpy as np import torch +import transformers # First Party +from caikit import get_config from caikit.core.data_model import DataStream +from caikit.core.exceptions import error_handler from caikit.core.modules import ModuleBase, ModuleConfig, ModuleSaver, module -from caikit.core.toolkit import error_handler from caikit.interfaces.nlp.data_model import ( ClassificationTrainRecord, GeneratedTextResult, @@ -62,12 +57,7 @@ PromptOutputModelType, TuningConfig, ) -from ...resources.pretrained_model import ( - HFAutoCausalLM, - HFAutoSeq2SeqLM, - PretrainedModelBase, -) -from ...toolkit.data_stream_wrapper import SimpleIterableStreamWrapper +from ...resources.pretrained_model import HFAutoCausalLM, HFAutoSeq2SeqLM from ...toolkit.data_type_utils import get_torch_dtype, str_to_torch_dtype from ...toolkit.task_specific_utils import convert_to_generation_record from ...toolkit.text_generation.model_run_utils import ( @@ -75,6 +65,15 @@ generate_text_func, generate_text_func_stream, ) +from ...toolkit.text_generation.training_utils import ( + ALLOWED_TRAINING_ARGS, + collect_trainer_arguments, + infer_max_steps, + launch_training, + preprocess_function +) +from ...toolkit.torch_run import get_torch_elastic_launch_config +from ...toolkit.trainer_utils import validate_training_data from ...toolkit.verbalizer_utils import render_verbalizer from .peft_config import TuningType, get_peft_config, resolve_base_model @@ -107,6 +106,7 @@ class PeftPromptTuning(ModuleBase): # TuningType.LORA: PeftType.LORA, } + RANDOM_SEED = 73 supported_resources = [HFAutoCausalLM, HFAutoSeq2SeqLM] ################################ Constructor / Destructor ##################################### @@ -160,28 +160,26 @@ def run( min_new_tokens: Optional[int] = 0, truncate_input_tokens: Optional[int] = 0, decoding_method: Optional[str] = "GREEDY", - top_k: Optional[int] = 0, - top_p: Optional[float] = 1.0, - typical_p: Optional[float] = 1.0, - temperature: Optional[float] = 1.0, - seed: Optional[np.uint64] = None, - repetition_penalty: Optional[float] = 1.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + temperature: Optional[float] = None, + repetition_penalty: Optional[float] = None, max_time: Optional[float] = None, exponential_decay_length_penalty: Optional[ Union[Tuple[int, float], ExponentialDecayLengthPenalty] ] = None, stop_sequences: Optional[List[str]] = None, + seed: Optional[np.uint64] = None, ) -> GeneratedTextResult: + f""" + Run the full text generation model. + Args: + {GENERATE_FUNCTION_ARGS} + Returns: + GeneratedTextResult + Generated text result produced by PEFT / Transformers. """ - Run the full text generation model. - Args: - {} - Returns: - GeneratedTextResult - Generated text result produced by PEFT / Transformers. - """.format( - GENERATE_FUNCTION_ARGS - ) verbalized_text = render_verbalizer(self.verbalizer, {"input": text}) @@ -220,19 +218,19 @@ def run_stream_out( min_new_tokens=0, truncate_input_tokens: Optional[int] = 0, decoding_method: Optional[str] = "GREEDY", - top_k: Optional[int] = 0, - top_p: Optional[float] = 0.0, - typical_p: Optional[float] = 0.0, - temperature: Optional[float] = 1.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + temperature: Optional[float] = None, seed: Optional[np.uint64] = None, - repetition_penalty: Optional[float] = 0.0, + repetition_penalty: Optional[float] = None, max_time: Optional[float] = None, exponential_decay_length_penalty: Optional[ Union[Tuple[int, float], ExponentialDecayLengthPenalty] ] = None, stop_sequences: Optional[List[str]] = None, ) -> Iterable[GeneratedTextStreamResult]: - """Run the text generation model with output streaming + f"""Run the text generation model with output streaming NOTE: This implementation is marked as WIP since the API for HuggingFace streamer classes at time of implementation is still @@ -240,13 +238,11 @@ def run_stream_out( Ref. https://huggingface.co/docs/transformers/v4.30.0/generation_strategies#streaming Args: - {} + {GENERATE_FUNCTION_ARGS} Returns: Iterable[GeneratedTextStreamResult] - """.format( - GENERATE_FUNCTION_ARGS - ) + """ # Apply the verbalizer to our text string verbalized_text = render_verbalizer(self.verbalizer, {"input": text}) @@ -295,9 +291,10 @@ def train( batch_size: Optional[int] = 8, max_source_length: Optional[int] = 256, max_target_length: Optional[int] = 128, - accumulate_steps: Optional[int] = 32, + accumulate_steps: Optional[int] = None, torch_dtype: Optional[str] = None, # TODO: Optional[Union[torch.dtype, str]] silence_progress_bars: Optional[bool] = True, + seed: int = RANDOM_SEED, **kwargs, ) -> "PeftPromptTuning": """Run prompt tuning (vanilla or MPT) through PEFT on a CausalLM or Seq2seq model @@ -333,8 +330,8 @@ def train( Max length of input sequences being considered. Default: 256. max_target_length: int Max length of target sequences being predicted. Default: 128. - accumulate_steps: int - Number of steps to use for gradient accumulation. Default: 1. + accumulate_steps: int (DEPRECATED) + Optional, number of steps to use for gradient accumulation. Default: None. torch_dtype: str TODO: Optional[Union[torch.dtype, str]] Data type to use for training/inference of the underlying text generation model. @@ -343,17 +340,46 @@ def train( underpinning the resource will be converted in place to the correct torch dtype. silence_progress_bars: bool Silences TQDM progress bars at train time. Default: True. + seed: int + Integer to be used as random seed for training. Returns: PeftPromptTuning Instance of this class with tuned prompt vectors. """ + error.value_check( + "", len(train_stream) > 0, "train_stream cannot be empty" + ) - # HACK - These things can't be passed through the train API currently + if accumulate_steps: + log.warning( + "", + "accumulate_steps parameter is DEPRECATED and will be removed in future. This parameter is also not getting used internally anymore") - metric = kwargs.get("metric") + # Configure random seed + transformers.set_seed(seed) + # NOTE: Following can be uncommented to allow full determinism + # but it can have impact on performance. + # transformers.enable_full_determinism(seed) + + torch_dtype = get_torch_dtype(torch_dtype) + # NOTE: We are not support "metrics" at the moment + + # Coerce the passed model into a resource; if we have one, this is a noop + # TODO: When splitting up this mono-module, use the configured resource + # type of the concrete class to bootstrap base_model = resolve_base_model(base_model, cls, torch_dtype) + # Enable gradient checkpointing on base model + # PeftModel checks if the base_model has gradient checkpointing + # enabled and then configures the tensors it creates with appropriate + # setting. If we do not enable this, then we will get `tensor 0` requires + # grad error, where `tensor 0` is created by peft + # base_model.model.gradient_checkpointing_enable() + base_model_name = base_model._model_name + # Get config of the base model + base_model_config = base_model.get_config() + task_type, output_model_types, peft_config, tuning_type = get_peft_config( tuning_type, tuning_config, @@ -363,26 +389,20 @@ def train( verbalizer, ) - # Coerce the passed model into a resource; if we have one, this is a noop - # TODO: When splitting up this mono-module, use the configured resource - # type of the concrete class to bootstrap - torch_dtype = get_torch_dtype(torch_dtype) + # Check if data is within limit allowed for this module and model + validate_training_data( + train_stream, + base_model_name, + cls.MODULE_ID, + ) train_stream = train_stream.map(convert_to_generation_record) if val_stream: - val_stream = val_stream.map(convert_to_generation_record) + error.value_check( + "", len(val_stream) > 0, "val_stream cannot be empty" + ) - # Convert our datastreams -> data loaders by disguising them as PyTorch iterable datasets - train_dataloader, val_dataloader = cls.create_dataloaders_from_stream( - base_model=base_model, - task_type=task_type, - train_stream=train_stream, - verbalizer=verbalizer, - validation_stream=val_stream or None, - batch_size=batch_size, - max_source_length=max_source_length, - max_target_length=max_target_length, - ) + val_stream = val_stream.map(convert_to_generation_record) log.debug("Peft config [%s]", peft_config) # FIXME: Should only do following line for causal LM (and bloomz?) - check that is the case @@ -394,24 +414,126 @@ def train( # Convert our Peft model (not just the underlying # transformers model) to the right underlying type. device = cls._get_device(device) - cls.convert_peft_model_to_type(device, peft_model, torch_dtype) - - training_loss_tracker = cls._execute_train_loop( - peft_model, - num_epochs, - train_dataloader, - device, - eval_dataloader=val_dataloader, - metric=metric, - learning_rate=learning_rate, + + ## Generate data loader from stream + training_dataset: Union[ + Dataset, TransformersIterableDataset + ] = preprocess_function( + base_model=base_model, + train_stream=train_stream, tokenizer=base_model.tokenizer, - accumulate_steps=accumulate_steps, - silence_progress_bars=silence_progress_bars, - torch_dtype=torch_dtype, + max_source_length=max_source_length, + max_target_length=max_target_length, + shuffle=True, + use_iterable_dataset=False, + random_seed=cls.RANDOM_SEED, + task_ids=0, ) - # Get config of the base model - base_model_config = base_model.get_config() + # Filter **training_arguments to only process allowed ones + filtered_training_arguments = { + k: v for k, v in kwargs.items() if k in ALLOWED_TRAINING_ARGS + } + + extra_training_args = set(kwargs.keys()).difference( + filtered_training_arguments.keys() + ) + + if extra_training_args: + log.warning( + "", + f"{extra_training_args} parameter(s) not allowed by \ + {cls.__name__} currently and will be ignored!", + ) + + if num_epochs < 1: + log.warning( + "", + f"Number of epochs configured is {num_epochs} which is less than minimum 1. \ + No training will be performed", + ) + + return PeftPromptTuning( + tokenizer=base_model.tokenizer, + model=peft_model, + base_model_config=base_model_config, + base_model_name=base_model_name, + verbalizer=verbalizer, + task_type=task_type, + tuning_type=tuning_type, + output_model_types=output_model_types, + training_metadata={"loss": []}, + ) + + processing_configuration = {} + + # Conditionally enable sharding if multiple GPUs available + if torch.cuda.is_available() and torch.cuda.device_count() > 1: + processing_configuration = { + "fsdp": "full_shard offload auto_wrap", + "fsdp_config": { + # NOTE: Every transformers model has `_no_split_modules` property that can be + # leveraged to identify the layers to split. This seems to be a decent + # "default" behavior unless we want to optimize further. We will start with + # this generic approach, since it allows us to handle variety + # of models and iterate on it, based on what we encounter. + "fsdp_transformer_layer_cls_to_wrap": base_model._model._no_split_modules, + # We need to use the original parameters for peft because we have mixed values + # for require_grads in our parametes, which otherwise breaks layer flattening + # in FSDP. + "use_orig_params": "true", + }, + } + + # Open an intermediate checkpoint directory until we've bootstrapped + # our model or we've early exited (if epochs < 1) + with tempfile.TemporaryDirectory() as checkpoint_dir: + + training_args = collect_trainer_arguments( + torch_dtype, + checkpoint_dir, + batch_size, + num_epochs, + cls.RANDOM_SEED, + learning_rate, + max_steps=infer_max_steps(num_epochs, batch_size, training_dataset), + silence_progress_bars=silence_progress_bars, + # NOTE: following can override above arguments in order + **filtered_training_arguments, + **processing_configuration, + ) + + if torch.cuda.is_available(): + # NOTE: torch distributed can hang if run on CPUs, + # to avoid that, specially for unit tests, we are only + # running below when GPUs are available + launch_config = get_torch_elastic_launch_config( + get_config().master_addr, + get_config().master_port, + ) + + training_loss_history = torch.distributed.launcher.api.elastic_launch( + launch_config, launch_training + )( + peft_model, + training_dataset, + training_args, + checkpoint_dir, + base_model, + ) + # NOTE: We are currently only storing the loss information from + # rank 0, i.e main process. training_loss_history is dictionary containing + # rank of the process as key + training_loss_history = training_loss_history[0] + else: + # Do not do FSDP things + training_loss_history = launch_training( + peft_model, + training_dataset, + training_args, + checkpoint_dir, + base_model, + ) # Remove _name_or_path field as a model can be # saved in different location but still same @@ -432,7 +554,7 @@ def train( task_type=task_type, tuning_type=tuning_type, output_model_types=output_model_types, - training_metadata=training_loss_tracker, + training_metadata={"loss": training_loss_history}, # TODO: Export other training params to model as well ) @@ -692,83 +814,6 @@ def get_exportable_prompt_vectors( return prompt_dict - @classmethod - def create_dataloaders_from_stream( - cls, - base_model: "caikit_nlp.resources.pretrained_model.base.PretrainedModelBase", - task_type: str, - train_stream: DataStream[GenerationTrainRecord], - verbalizer: str, - batch_size: int, - max_source_length: int, - max_target_length: int, - validation_stream: Union[DataStream[GenerationTrainRecord], None] = None, - collate_fn: Callable = None, - ) -> Tuple[DataLoader]: - """Build PyTorch data loaders around training and (optionally) evaluation DataStreams. - - Args: - base_model: caikit_nlp.resources.pretrained_model.base.PretrainedModelBase - Base resource model used for underlying generation. - task_type: str - Str indicating which task is being accomplished; currently used for determining - tokenization / preprocessing behavior. - train_stream: DataStream[GenerationTrainRecord] - Data to be used for training the prompt vectors of the generation model. - verbalizer: str - Verbalizer template with which we will render text at both train & inference time. - batch_size: int - Batch size to be used for train/eval data loaders. - max_source_length: int - Maximum length to be used for tokenized sequences. - max_target_length: int - Max length of target sequences being predicted. - validation_stream: Union[DataStream[GenerationTrainRecord], None] - Data to be used for validation throughout the train process or None. - collate_fn: Callable - Function to be used for forming batches via lists of dataset inputs. - - Returns: - Tuple[torch.utils.data.DataLoader] - Training & evaluation datastreams for the provided data, respectively. If no - validation_stream is provided, the returned loader for validation_stream will - be None. - """ - if collate_fn is None: - # collate_fn -> pads and maps our inputs to PyTorch vectors - collate_fn = cls._get_collate_fn(base_model.tokenizer, task_type) - - # Grab the data loaders for this task. - # NOTE: Currently we do not expose the buffer size and we - # default to loading the whole dataset into memory - train_dataloader = cls._get_data_loaders_from_stream( - base_model, - train_stream, - base_model.tokenizer, - batch_size, - collate_fn, - verbalizer, - max_source_length, - max_target_length, - shuffle=True, - ) - if validation_stream is not None: - val_dataloader = cls._get_data_loaders_from_stream( - base_model, - validation_stream, - base_model.tokenizer, - batch_size, - collate_fn, - verbalizer, - max_source_length, - max_target_length, - shuffle=False, - ) - else: - val_dataloader = None - - return train_dataloader, val_dataloader - @classmethod def create_hf_tuning_config( cls, @@ -890,251 +935,10 @@ def _get_collate_fn(tokenizer: AutoTokenizer, task_type: str) -> Callable: Callable collate_fn to be used for processing batches from our datasets. """ - if task_type == "CAUSAL_LM": - return DataCollatorForLanguageModeling( - tokenizer=tokenizer, - return_tensors="pt", - mlm=False, - ) + # HACK: Do NOT use the causal LM collator (for now) because + # want to set labels ourselves. TODO: centralize collator management. return default_data_collator - @staticmethod - def _get_data_loaders_from_stream( - base_model: PretrainedModelBase, - train_stream: DataStream[GenerationTrainRecord], - tokenizer: AutoTokenizer, - batch_size: int, - collate_fn: Callable, - verbalizer: str, - max_source_length: int, - max_target_length: int, - shuffle: bool, - ) -> DataLoader: - """Get the data loaders for train / evaluation. - Args: - base_model: caikit_nlp.resources.pretrained_model.base.PretrainedModelBase - Base resource model used for underlying generation. - train_stream: DataStream[GenerationTrainRecord] - Data to be used for training the prompt vectors of the generation model. - tokenizer: AutoTokenizer - Model tokenizer to be used in preprocessing, i.e., when we iterate over our data. - batch_size: int - Batch sized to be used when building the DataLoader around the stream. - collate_fn: Callable - Function to be used for forming batches via lists of dataset inputs. - verbalizer: str - Verbalizer template to be used for formatting data. This template may use brackets - to indicate where fields from the data model TrainGenerationRecord must be rendered. - max_source_length: int - Max length of sequences being considered. - max_target_length: int - Max length of target sequences being predicted. - shuffle: bool - Indicates whether or not the stream should reshuffle upon reentry. - - Returns: - torch.utils.data.DataLoader - DataLoader to be used for training / evaluating the stream data. - """ - ( - tokenize_function, - requires_unwrapping, - ) = base_model.build_task_tokenize_closure( - tokenizer, max_source_length, max_target_length, verbalizer, task_ids=0 - ) - mapped_stream = train_stream.map(tokenize_function) - if requires_unwrapping: - mapped_stream = mapped_stream.flatten() - wrapped_stream = SimpleIterableStreamWrapper(mapped_stream, shuffle=shuffle) - dataloader = DataLoader( - wrapped_stream, collate_fn=collate_fn, batch_size=batch_size - ) - - return dataloader - - @classmethod - def _execute_train_loop( - cls, - model: PeftModel, - num_epochs: int, - train_dataloader: DataLoader, - device: str, - eval_dataloader: Union[DataLoader, None] = None, - metric: Optional[Callable] = None, - learning_rate: int = 1e-3, - tokenizer: Union[AutoTokenizer, None] = None, - accumulate_steps: int = 1, - silence_progress_bars: bool = True, - torch_dtype: "torch.dtype" = torch.float32, - ) -> None: - """Execute the core training logic for training the prompt vectors on the frozen model. - Note that this is done by reference. - - Args: - model: PeftModel - Underlying model being leveraged for text generation via prompt tuning. - num_epochs: int - Number of epochs to train. - train_dataloader: torch.utils.data.DataLoader - DataLoader to be used for loading training data. - device: str - Device to be used for training the model. - eval_dataloader: Union[DataLoader, None]. - DataLoader to be used for loading eval data or None. - metric: Union[Callable, None] - Function to be used for evaluating data if an eval data loader is provided. - Default: None. - learning_rate: float - Learning rate to be used while tuning prompt vectors. Default: 1e-3. - tokenizer: Union[AutoTokenizer, None] - Tokenizer for default evaluation; only used if no metric is provided and we have - an eval dataloader. - TODO - remove this can likely be removed. - accumulate_steps: int - Number of steps to use for gradient accumulation. Default: 1. - silence_progress_bars: bool - Silences TQDM progress bars. Default: True - torch_dtype: torch.dtype - Dtype to be used for training. Default: torch.float32 - - Returns: - training_metadata: Dict - Metadata computed during training - """ - optimizer = AdamW(params=model.parameters(), lr=learning_rate) - lr_scheduler = get_linear_schedule_with_warmup( - optimizer=optimizer, - num_warmup_steps=0, - num_training_steps=(len(train_dataloader) * num_epochs), - ) - - # Enable gradient checkpointing - model.gradient_checkpointing_enable() - - if torch_dtype == torch.float16: - mixed_precision = "fp16" - elif ( - torch.cuda.is_available() - and torch.cuda.is_bf16_supported() - and torch_dtype == torch.bfloat16 - ): - mixed_precision = "bf16" - else: - mixed_precision = "no" - - accelerator = Accelerator( - gradient_accumulation_steps=accumulate_steps, - device_placement=True, - mixed_precision=mixed_precision, - ) - - # Disable cache for training - model.config.use_cache = False - - # Below would send all the data and model to - # configured device and convert them to required dtypes - model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( - model, - optimizer, - train_dataloader, - lr_scheduler, - ) - - training_loss_tracker = [] - - for epoch in range(num_epochs): - model.train() - total_loss = 0 - tqdm_loader = tqdm(train_dataloader, disable=silence_progress_bars) - for batch in tqdm_loader: - - tqdm_loader.set_description("Epoch: {}".format(epoch)) - - # TODO Can this dict comprehension always replace "batch.to(device)" for us? - try: - with accelerator.accumulate(model): - outputs = model(**batch) - loss = outputs.loss - total_loss += loss.detach().float() - accelerator.backward(loss) - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - except torch.cuda.OutOfMemoryError: - error( - "", - MemoryError("Not enough memory available for training!"), - ) - - log.info("", {"loss": float(loss), "epoch": epoch}) - # Below is added to be propagated and stored as training_metadata - training_loss_tracker.append( - { - "epoch": epoch, - "value": float(loss), - "timestamp": datetime.isoformat(datetime.now()), - } - ) - - if eval_dataloader is not None: - model.eval() - - if metric is not None: - for _, batch in enumerate( - tqdm(eval_dataloader, disable=silence_progress_bars) - ): - batch.to(device) - with torch.no_grad(): - outputs = model(**batch) - predictions = outputs.logits.argmax(dim=-1) - references = batch["labels"] - metric.add_batch( - predictions=predictions, - references=references, - ) - eval_metric = metric.compute() - - log.info("epoch %s: %s", epoch, eval_metric) - else: - eval_loss = 0 - # TODO Can we get away with not maintaining eval_preds? - eval_preds = [] - for _, batch in enumerate( - tqdm(eval_dataloader, disable=silence_progress_bars) - ): - batch = {k: v.to(device) for k, v in batch.items()} - with torch.no_grad(): - outputs = model(**batch) - loss = outputs.loss - eval_loss += loss.detach().float() - - if tokenizer is not None: - eval_preds.extend( - tokenizer.batch_decode( - torch.argmax(outputs.logits, -1) - .detach() - .cpu() - .numpy(), - skip_special_tokens=True, - ) - ) - - eval_epoch_loss = eval_loss / len(train_dataloader) - eval_ppl = torch.exp(eval_epoch_loss) - train_epoch_loss = total_loss / len(eval_dataloader) - train_ppl = torch.exp(train_epoch_loss) - log.debug( - "epoch %s: %s %s %s %s", - epoch, - train_ppl, - train_epoch_loss, - eval_ppl, - eval_epoch_loss, - ) - - error.value_check("", len(training_loss_tracker) == num_epochs) - return {"loss": training_loss_tracker} - @classmethod def _filter_params_for_prompt_config(cls, prompt_config, params): """Utility function to filter out required parameters for prompt_config diff --git a/caikit_nlp/modules/text_generation/peft_tgis_remote.py b/caikit_nlp/modules/text_generation/peft_tgis_remote.py index fa21e34cf..b53e7ef5c 100644 --- a/caikit_nlp/modules/text_generation/peft_tgis_remote.py +++ b/caikit_nlp/modules/text_generation/peft_tgis_remote.py @@ -24,8 +24,8 @@ # First Party from caikit.config import get_config from caikit.core import ModuleBase, ModuleConfig, ModuleSaver, modules +from caikit.core.exceptions import error_handler from caikit.core.module_backends import BackendBase, backend_types -from caikit.core.toolkit import error_handler from caikit.interfaces.nlp.data_model import ( GeneratedTextResult, GeneratedTextStreamResult, @@ -175,6 +175,7 @@ def save(self, model_path: str): } ) + # pylint: disable=duplicate-code @TextGenerationTask.taskmethod() def run( self, @@ -183,29 +184,27 @@ def run( min_new_tokens: Optional[int] = 0, truncate_input_tokens: Optional[int] = 0, decoding_method: Optional[str] = "GREEDY", - top_k: Optional[int] = 0, - top_p: Optional[float] = 1.0, - typical_p: Optional[float] = 1.0, - temperature: Optional[float] = 1.0, - seed: Optional[np.uint64] = None, - repetition_penalty: Optional[float] = 1.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + temperature: Optional[float] = None, + repetition_penalty: Optional[float] = None, max_time: Optional[float] = None, exponential_decay_length_penalty: Optional[ Union[Tuple[int, float], ExponentialDecayLengthPenalty] ] = None, stop_sequences: Optional[List[str]] = None, + seed: Optional[np.uint64] = None, preserve_input_text: bool = False, ) -> GeneratedTextResult: - """Run inference against the model running in TGIS. + f"""Run inference against the model running in TGIS. Args: - {} + {GENERATE_FUNCTION_TGIS_ARGS} Returns: GeneratedTextResult Generated text result produced by TGIS. - """.format( - GENERATE_FUNCTION_TGIS_ARGS - ) + """ error.value_check( "", @@ -239,12 +238,12 @@ def run_stream_out( min_new_tokens: Optional[int] = 0, truncate_input_tokens: Optional[int] = 0, decoding_method: Optional[str] = "GREEDY", - top_k: Optional[int] = 0, - top_p: Optional[float] = 1.0, - typical_p: Optional[float] = 1.0, - temperature: Optional[float] = 1.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + temperature: Optional[float] = None, seed: Optional[np.uint64] = None, - repetition_penalty: Optional[float] = 1.0, + repetition_penalty: Optional[float] = None, max_time: Optional[float] = None, exponential_decay_length_penalty: Optional[ Union[Tuple[int, float], ExponentialDecayLengthPenalty] @@ -252,15 +251,13 @@ def run_stream_out( stop_sequences: Optional[List[str]] = None, preserve_input_text: bool = False, ) -> Iterable[GeneratedTextStreamResult]: - """Run output stream inferencing against the model running in TGIS + f"""Run output stream inferencing against the model running in TGIS Args: - {} + {GENERATE_FUNCTION_TGIS_ARGS} Returns: Iterable[GeneratedTextStreamResult] - """.format( - GENERATE_FUNCTION_TGIS_ARGS - ) + """ error.value_check( "", diff --git a/caikit_nlp/modules/text_generation/text_generation_local.py b/caikit_nlp/modules/text_generation/text_generation_local.py index 3b2ff535f..84827ad66 100644 --- a/caikit_nlp/modules/text_generation/text_generation_local.py +++ b/caikit_nlp/modules/text_generation/text_generation_local.py @@ -14,22 +14,23 @@ # Standard -from typing import Optional, Union +from typing import Any, Dict, Optional, Union import gc +import json import os import tempfile # Third Party from datasets import Dataset from datasets import IterableDataset as TransformersIterableDataset -from transformers import AutoConfig, AutoTokenizer +from transformers import AutoConfig import torch # First Party from caikit import get_config from caikit.core.data_model import DataStream +from caikit.core.exceptions import error_handler from caikit.core.modules import ModuleBase, ModuleConfig, ModuleSaver, module -from caikit.core.toolkit import error_handler from caikit.interfaces.nlp.data_model import GeneratedTextResult from caikit.interfaces.nlp.tasks import TextGenerationTask import alog @@ -46,12 +47,21 @@ GENERATE_FUNCTION_ARGS, generate_text_func, ) +from ...toolkit.text_generation.training_utils import ( + ALLOWED_TRAINING_ARGS, + collect_trainer_arguments, + infer_max_steps, + launch_training, + preprocess_function, +) from ...toolkit.torch_run import get_torch_elastic_launch_config log = alog.use_channel("TXT_GEN") error = error_handler.get(log) +TRAINING_LOSS_LOG_FILENAME = "training_logs.jsonl" + # pylint: disable=too-many-lines,too-many-instance-attributes @module( id="f9181353-4ccf-4572-bd1e-f12bcda26792", @@ -65,28 +75,6 @@ class TextGeneration(ModuleBase): RANDOM_SEED = 73 supported_resources = [HFAutoCausalLM, HFAutoSeq2SeqLM] - # Below list is taken from - # https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments - allowed_training_args = { - "weight_decay", - "adam_beta1", - "adam_beta2", - "adam_epsilon", - "max_grad_norm", - "lr_scheduler_type", - "warmup_ratio", - "warmup_steps", - "use_ipex", - "disable_tqdm", - "label_names", - "optim", - "optim_args", - "group_by_length", - "dataloader_pin_memory", - "gradient_checkpointing", - "full_determinism", - } - def __init__( self, model_name: str, @@ -95,6 +83,7 @@ def __init__( sep_token: Optional[str] = None, eos_token: Optional[str] = None, pad_token: Optional[str] = None, + training_metadata: Union[Dict[str, Any], None] = None, ): super().__init__() @@ -106,6 +95,9 @@ def __init__( self._sep_token = sep_token self._eos_token = eos_token self._pad_token = pad_token + self.training_metadata = ( + training_metadata if training_metadata is not None else {} + ) # pylint: disable=duplicate-code def __del__(self): @@ -173,7 +165,7 @@ def train( num_epochs: int = 5, accumulate_steps: int = 32, random_seed: int = RANDOM_SEED, - lr: float = 2e-5, + learning_rate: float = 2e-5, use_iterable_dataset: bool = True, **kwargs, ) -> "TextGeneration": @@ -201,7 +193,7 @@ def train( Number of epochs to tune the model. Default: 20. accumulate_steps: int Number of steps to use for gradient accumulation. Default: 1. - lr: float + learning_rate: float Learning rate to be used while tuning model. Default: 2e-5. use_iterable_dataset: bool Indicates whether or not we should load the full dataset into memory @@ -217,6 +209,9 @@ def train( TextGeneration Instance of this class with fine-tuned models. """ + error.value_check( + "", len(train_stream) > 0, "train_stream cannot be empty" + ) torch_dtype = get_torch_dtype(torch_dtype) @@ -255,6 +250,18 @@ def train( error.type_check("", PretrainedModelBase, base_model=base_model) + if num_epochs < 1: + log.warning( + "", + f"Number of epochs configured is {num_epochs} which is less than minimum 1. \ + No training will be performed", + ) + + return TextGeneration( + model_name=base_model._model_name, + model=base_model, + ) + # TODO; For now, we don't support iterable datasets for causal LM, the reason # being that the dynamically padded collator breaks with FSDP/torch run when # multiple devices are used. This appears to be because each process fetches @@ -277,7 +284,7 @@ def train( ## Generate data loader from stream training_dataset: Union[ Dataset, TransformersIterableDataset - ] = cls._preprocess_function( + ] = preprocess_function( base_model=base_model, train_stream=train_stream, tokenizer=base_model.tokenizer, @@ -285,30 +292,16 @@ def train( max_target_length=max_target_length, shuffle=True, use_iterable_dataset=use_iterable_dataset, + random_seed=cls.RANDOM_SEED, ) - ### Dtype based processing - # NOTE: Following is not exhaustive list of all parameters - # for all dtypes - if torch_dtype == torch.float16: - dtype_based_params = { - "fp16": True, - } - elif torch_dtype == torch.bfloat16: - dtype_based_params = { - "bf16": True, - } - else: - # default to float32 - dtype_based_params = {} - ## TODO: Add automatic sharding selection based on number of parameters # in base model ## TODO: Fetch trainer from resource # Filter **training_arguments to only process allowed ones filtered_training_arguments = { - k: v for k, v in kwargs.items() if k in cls.allowed_training_args + k: v for k, v in kwargs.items() if k in ALLOWED_TRAINING_ARGS } extra_training_args = set(kwargs.keys()).difference( @@ -341,50 +334,19 @@ def train( # Open an intermediate checkpoint directory until we've bootstrapped # our model or we've early exited (if epochs < 1) with tempfile.TemporaryDirectory() as checkpoint_dir: - training_args = { - "output_dir": checkpoint_dir, - "per_device_train_batch_size": batch_size, - "per_device_eval_batch_size": batch_size, - "num_train_epochs": num_epochs, - "seed": random_seed, - # NOTE: We have disabled evaluation for now - "do_eval": False, - "do_train": True, - "learning_rate": lr, - "weight_decay": 0.01, - "save_total_limit": 3, - "push_to_hub": False, - "no_cuda": not torch.cuda.is_available(), # Default - "remove_unused_columns": True, - "dataloader_pin_memory": False, - "gradient_accumulation_steps": accumulate_steps, - "gradient_checkpointing": True, - # NOTE: This is explicitly set to false since it will - # negatively impact the performance - "full_determinism": False, - # Required for iterable dataset - "max_steps": cls.infer_max_steps( - num_epochs, batch_size, training_dataset - ), - # Some interesting parameters: - "auto_find_batch_size": True, + training_args = collect_trainer_arguments( + torch_dtype, + checkpoint_dir, + batch_size, + num_epochs, + random_seed, + learning_rate, + accumulate_steps, + max_steps=infer_max_steps(num_epochs, batch_size, training_dataset), # NOTE: following can override above arguments in order **filtered_training_arguments, **processing_configuration, - **dtype_based_params, - } - - if num_epochs < 1: - log.warning( - "", - f"Number of epochs configured is {num_epochs} which is less than minimum 1. \ - No training will be performed", - ) - - return TextGeneration( - model_name=base_model._model_name, - model=base_model, - ) + ) launch_config = get_torch_elastic_launch_config( get_config().master_addr, @@ -395,11 +357,16 @@ def train( # NOTE: torch distributed can hang if run on CPUs, # to avoid that, specially for unit tests, we are only # running below when GPUs are available - torch.distributed.launcher.api.elastic_launch( - launch_config, cls._launch_training + training_loss_history = torch.distributed.launcher.api.elastic_launch( + launch_config, launch_training )(base_model, training_dataset, training_args, checkpoint_dir) + + # NOTE: We are currently only storing the loss information from + # rank 0, i.e main process. training_loss_history is dictionary containing + # rank of the process as key + training_loss_history = training_loss_history[0] else: - cls._launch_training( + training_loss_history = launch_training( base_model, training_dataset, training_args, checkpoint_dir ) @@ -423,6 +390,7 @@ def train( sep_token=model.tokenizer.sep_token or None, eos_token=model.tokenizer.eos_token or None, pad_token=model.tokenizer.pad_token or None, + training_metadata={"loss": training_loss_history}, ) @classmethod @@ -485,6 +453,24 @@ def save(self, model_path): base_model_dirname=artifacts_dir, ) + training_loss_filename = TRAINING_LOSS_LOG_FILENAME + + saver.update_config({"training_logs": training_loss_filename}) + + # We are currently only saving logs containing loss in jsonl format + if "loss" in self.training_metadata: + loss_log_lines = self.training_metadata.get("loss") + error.type_check("", list, loss_log_lines=loss_log_lines) + with open( + os.path.join(model_path, training_loss_filename), + "w", + encoding="utf-8", + ) as f: + for loss_log in loss_log_lines: + loss_log = {"name": "loss", "data": loss_log} + json.dump(loss_log, f) + f.write("\n") + def run( self, text: str, @@ -492,25 +478,23 @@ def run( min_new_tokens: Optional[int] = 0, truncate_input_tokens: Optional[int] = 0, decoding_method: Optional[str] = "GREEDY", - top_k: Optional[int] = 0, - top_p: Optional[float] = 0.0, - typical_p: Optional[float] = 0.0, - temperature: Optional[float] = 1.0, - repetition_penalty: Optional[float] = 0.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + temperature: Optional[float] = None, + repetition_penalty: Optional[float] = None, max_time: Optional[float] = None, **kwargs, ) -> GeneratedTextResult: + f""" + Run the full text generation model. + Args: + {GENERATE_FUNCTION_ARGS} + Returns: + GeneratedTextResult + Generated text result produced by the model. """ - Run the full text generation model. - Args: - {} - Returns: - GeneratedTextResult - Generated text result produced by the model. - """.format( - GENERATE_FUNCTION_ARGS - ) # TODO: Beam search currently not supported @@ -533,94 +517,3 @@ def run( max_time=max_time, **kwargs, ) - - ################################## Private Functions ###################################### - - @staticmethod - def _preprocess_function( - base_model: PretrainedModelBase, - train_stream: DataStream[GenerationTrainRecord], - tokenizer: AutoTokenizer, - max_source_length: int, - max_target_length: int, - shuffle: bool, - use_iterable_dataset: bool, - ): - """Pre-process each example to get it prepared for training.""" - dataset_type = TransformersIterableDataset if use_iterable_dataset else Dataset - log.debug("Loading dataset class: [%s]", dataset_type.__name__) - fn_kwargs = { - "tokenizer": tokenizer, - "max_source_length": max_source_length, - "max_target_length": max_target_length, - } - dataset = dataset_type.from_generator( - get, gen_kwargs={"train_stream": train_stream} - ) - mapped_dataset = dataset.map( - base_model.tokenize_function, - fn_kwargs=fn_kwargs, - batched=base_model.REQUIRES_TOKEN_UNWRAPPING, - # Drop the input / output columns; we need to do this for dimensions to play - # happily when operating on batched inputs for causal language modeling. - remove_columns=["input", "output"], - ) - - if shuffle: - log.debug("Shuffling the dataset") - return mapped_dataset.shuffle(seed=TextGeneration.RANDOM_SEED) - - return mapped_dataset - - @staticmethod - def _launch_training( - base_model, training_dataset, training_args, checkpoint_dir - ) -> None: - """Utility function to wrap trainer and execute training""" - - trainer = base_model.get_trainer( - train_dataset=training_dataset, **training_args - ) - - # Start training via Trainer.train function - trainer.train() - - # save the model temporarily and reload it - # this is done, since otherwise the model might be distributed in different - # devices, in which case its better to use trainer's `prediction_step` - # functions, but then, they don't always give API similar to `generate` - # and thus cause incompatibilities in `run` function - trainer.save_state() - trainer.save_model(checkpoint_dir) - - # save tokenizer explicitly - base_model.tokenizer.save_pretrained(checkpoint_dir) - - @staticmethod - def infer_max_steps( - num_epochs: int, - batch_size: int, - training_dataset: Union[Dataset, TransformersIterableDataset], - ): - # Calculate the number of samples that we have - if isinstance(training_dataset, Dataset): - data_len = len(training_dataset) - else: - data_len = 0 - for _ in training_dataset: - data_len += 1 - # Figure out how many batches we'll have per epoch - num_batches = data_len // batch_size - # Assume drop_last=False; in general, this doesn't really matter. - # We mostly do this to avoid strange behavior when the dataset - # size is smaller than the batch size. - if num_batches != (data_len * batch_size): - num_batches += 1 - num_steps = num_batches * num_epochs - log.debug("Number of inferred steps: [%s]", num_steps) - return num_steps - - -def get(train_stream): - for data in train_stream: - yield {"input": data.input, "output": data.output} diff --git a/caikit_nlp/modules/text_generation/text_generation_tgis.py b/caikit_nlp/modules/text_generation/text_generation_tgis.py index b537889bb..0a75ce054 100644 --- a/caikit_nlp/modules/text_generation/text_generation_tgis.py +++ b/caikit_nlp/modules/text_generation/text_generation_tgis.py @@ -21,9 +21,9 @@ import numpy as np # First Party +from caikit.core.exceptions import error_handler from caikit.core.module_backends import BackendBase, backend_types from caikit.core.modules import ModuleBase, ModuleConfig, ModuleSaver, module -from caikit.core.toolkit import error_handler from caikit.interfaces.nlp.data_model import ( GeneratedTextResult, GeneratedTextStreamResult, @@ -201,6 +201,7 @@ def save(self, model_path: str): } ) + # pylint: disable=duplicate-code @TextGenerationTask.taskmethod() def run( self, @@ -209,29 +210,27 @@ def run( min_new_tokens: Optional[int] = 0, truncate_input_tokens: Optional[int] = 0, decoding_method: Optional[str] = "GREEDY", - top_k: Optional[int] = 0, - top_p: Optional[float] = 1.0, - typical_p: Optional[float] = 1.0, - temperature: Optional[float] = 1.0, - seed: Optional[np.uint64] = None, - repetition_penalty: Optional[float] = 1.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + temperature: Optional[float] = None, + repetition_penalty: Optional[float] = None, max_time: Optional[float] = None, exponential_decay_length_penalty: Optional[ Union[Tuple[int, float], ExponentialDecayLengthPenalty] ] = None, stop_sequences: Optional[List[str]] = None, + seed: Optional[np.uint64] = None, preserve_input_text: bool = False, ) -> GeneratedTextResult: - """Run inference against the model running in TGIS. + f"""Run inference against the model running in TGIS. Args: - {} + {GENERATE_FUNCTION_TGIS_ARGS} Returns: GeneratedTextResult Generated text result produced by TGIS. - """.format( - GENERATE_FUNCTION_TGIS_ARGS - ) + """ if self._model_loaded: return self.tgis_generation_client.unary_generate( @@ -260,12 +259,12 @@ def run_stream_out( min_new_tokens: Optional[int] = 0, truncate_input_tokens: Optional[int] = 0, decoding_method: Optional[str] = "GREEDY", - top_k: Optional[int] = 0, - top_p: Optional[float] = 1.0, - typical_p: Optional[float] = 1.0, - temperature: Optional[float] = 1.0, + top_k: Optional[int] = None, + top_p: Optional[float] = None, + typical_p: Optional[float] = None, + temperature: Optional[float] = None, seed: Optional[np.uint64] = None, - repetition_penalty: Optional[float] = 1.0, + repetition_penalty: Optional[float] = None, max_time: Optional[float] = None, exponential_decay_length_penalty: Optional[ Union[Tuple[int, float], ExponentialDecayLengthPenalty] @@ -273,15 +272,13 @@ def run_stream_out( stop_sequences: Optional[List[str]] = None, preserve_input_text: bool = False, ) -> Iterable[GeneratedTextStreamResult]: - """Run output stream inferencing for text generation module. + f"""Run output stream inferencing for text generation module. Args: - {} + {GENERATE_FUNCTION_TGIS_ARGS} Returns: Iterable[GeneratedTextStreamResult] - """.format( - GENERATE_FUNCTION_TGIS_ARGS - ) + """ if self._model_loaded: return self.tgis_generation_client.stream_generate( diff --git a/caikit_nlp/modules/token_classification/filtered_span_classification.py b/caikit_nlp/modules/token_classification/filtered_span_classification.py index ded05969d..15ea432a8 100644 --- a/caikit_nlp/modules/token_classification/filtered_span_classification.py +++ b/caikit_nlp/modules/token_classification/filtered_span_classification.py @@ -21,6 +21,7 @@ import os # First Party +from caikit.core.exceptions import error_handler from caikit.core.modules import ( ModuleBase, ModuleConfig, @@ -28,7 +29,6 @@ ModuleSaver, module, ) -from caikit.core.toolkit import error_handler from caikit.interfaces.nlp.data_model import ( TokenClassificationResult, TokenClassificationResults, @@ -87,7 +87,7 @@ def __init__( error.type_check("", ModuleBase, tokenizer=tokenizer) error.value_check( "", - tokenizer.TASK_CLASS == TokenizationTask, + TokenizationTask in type(tokenizer).tasks, "tokenizer does not implement TokenizationTask", ) error.type_check( @@ -99,18 +99,24 @@ def __init__( error.type_check_all( "", str, allow_none=True, labels_to_output=labels_to_output ) - classification_task = classifier.TASK_CLASS + classification_tasks = type(classifier).tasks + tasks_intersection = [i for i in classification_tasks if i in ALLOWED_TASKS] error.value_check( "", - classification_task in ALLOWED_TASKS, + any(tasks_intersection), f"classifier does not implement one of required tasks: {ALLOWED_TASKS}", ) + error.value_check( + "", + len(tasks_intersection) == 1, + f"classifier should implement only one task in: {ALLOWED_TASKS}", + ) self.lang = lang self.tokenizer = tokenizer self.classifier = classifier self.default_threshold = default_threshold self.labels_to_output = labels_to_output - self.classification_task = classification_task + self.classification_task = tasks_intersection[0] ################################## API functions ############################################# diff --git a/caikit_nlp/modules/tokenization/regex_sentence_splitter.py b/caikit_nlp/modules/tokenization/regex_sentence_splitter.py index d9f578eed..5b8adf31b 100644 --- a/caikit_nlp/modules/tokenization/regex_sentence_splitter.py +++ b/caikit_nlp/modules/tokenization/regex_sentence_splitter.py @@ -18,8 +18,8 @@ import re # First Party +from caikit.core.exceptions import error_handler from caikit.core.modules import ModuleBase, ModuleConfig, ModuleSaver, module -from caikit.core.toolkit import error_handler from caikit.interfaces.nlp.data_model import Token, TokenizationResults from caikit.interfaces.nlp.tasks import TokenizationTask import alog diff --git a/caikit_nlp/resources/pretrained_model/base.py b/caikit_nlp/resources/pretrained_model/base.py index 3f6aa4984..2217a2617 100644 --- a/caikit_nlp/resources/pretrained_model/base.py +++ b/caikit_nlp/resources/pretrained_model/base.py @@ -15,7 +15,7 @@ # Standard from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import Callable, List, Optional, Tuple, Type, Union +from typing import Callable, Dict, List, Optional, Tuple, Type, Union import json import os @@ -24,6 +24,7 @@ from transformers import ( AutoTokenizer, DataCollatorWithPadding, + PreTrainedTokenizerBase, Trainer, TrainingArguments, ) @@ -33,18 +34,36 @@ # First Party from caikit import get_config from caikit.core.data_model import DataStream +from caikit.core.exceptions import error_handler from caikit.core.modules import ModuleBase, ModuleConfig, ModuleSaver -from caikit.core.toolkit import error_handler import alog # Local from ...data_model import GenerationTrainRecord, PromptOutputModelType from ...toolkit.data_type_utils import get_torch_dtype, str_to_torch_dtype +from ...toolkit.trainer_utils import log_step log = alog.use_channel("HFRBAS") error = error_handler.get(log) +class LoggingTrainer(Trainer): + def log(self, logs: Dict[str, float]) -> None: + """ + Log `logs` on the various objects watching training. + + Subclass and override this method to inject custom behavior. + + Args: + logs (`Dict[str, float]`): + The values to log. + """ + self.state = log_step(self.state, logs) + self.control = self.callback_handler.on_log( + self.args, self.state, self.control, logs + ) + + class PretrainedModelBase(ABC, ModuleBase): """Common abstractions and requirements for pretrained model resources""" @@ -132,7 +151,7 @@ def bootstrap( Args: model_name (str) The name/path of the HF sequence classifier model - tokenizer_name (Optional[str] + tokenizer_name (Optional[Union[str, PreTrainedTokenizerBase]]) The name/path of the HF tokenizer model (matches model_name if not given) or an instance of a loaded tokenizer. NOTE: If a loaded tokenizer is provided, and it doesn't have @@ -151,42 +170,46 @@ def bootstrap( model HFAutoSequenceClassifier The loaded resource model """ - # Default tokenizer to model if downloading with a model name rather - # than a path - error.value_check( - "", - not os.path.isdir(model_name) or tokenizer_name is not None, - "Must provide path to tokenizer if model_name is a path", - ) + torch_dtype = get_torch_dtype(torch_dtype) - if tokenizer_name is None: - tokenizer_name = model_name - if not os.path.isdir(tokenizer_name) and tokenizer_name != model_name: - log.warning( - "Bootstrapping with mismatched tokenizer (%s) / model (%s)", - tokenizer_name, - model_name, - ) - # Figure out the right padding side based on the name of the HF model - # NOTE: This matches models whose name includes the left-pad types as a - # substring and not just as an exact match. - if padding_side is None: - padding_side = ( - "left" - if any(k in model_name for k in cls._LEFT_PAD_MODEL_TYPES) - else "right" + # Check if we passed the tokenizer directly; for now, we keep + # the arg name tokenizer_name for compatibility reasons + if isinstance(tokenizer_name, PreTrainedTokenizerBase): + log.debug("Bootstrapping with in-memory tokenizer") + tokenizer = tokenizer_name + + else: + if tokenizer_name is None: + tokenizer_name = model_name + log.info("Loading tokenizer from model directory") + + if not os.path.isdir(tokenizer_name) and tokenizer_name != model_name: + log.warning( + "Bootstrapping with mismatched tokenizer (%s) / model (%s)", + tokenizer_name, + model_name, + ) + + # Figure out the right padding side based on the name of the HF model + # NOTE: This matches models whose name includes the left-pad types as a + # substring and not just as an exact match. + if padding_side is None: + padding_side = ( + "left" + if any(k in model_name for k in cls._LEFT_PAD_MODEL_TYPES) + else "right" + ) + + # Load the tokenizer and set up the pad token if needed + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name, + local_files_only=not get_config().allow_downloads, + padding_side=padding_side, + # We can't disable use_fast otherwise unit test fails + # use_fast=False, ) - # Load the tokenizer and set up the pad token if needed - tokenizer = AutoTokenizer.from_pretrained( - tokenizer_name, - local_files_only=not get_config().allow_downloads, - padding_side=padding_side, - # We can't disable use_fast otherwise unit test fails - # use_fast=False, - ) - if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id @@ -256,6 +279,7 @@ def get_trainer( train_dataset: IterableDataset, eval_dataset: Union[IterableDataset, None] = None, optimizers=(None, None), + model=None, **kwargs, ): """ @@ -281,7 +305,12 @@ def get_trainer( "eval_dataset": eval_dataset, } - return Trainer(self._model, training_args, **trainer_arguments) + # If extra model is provided, we will configure trainer + # with that model + if model: + return LoggingTrainer(model, training_args, **trainer_arguments) + + return LoggingTrainer(self._model, training_args, **trainer_arguments) def _get_data_collator(self, **kwargs): """Function to return appropriate data collator based on resource. diff --git a/caikit_nlp/resources/pretrained_model/hf_auto_causal_lm.py b/caikit_nlp/resources/pretrained_model/hf_auto_causal_lm.py index e5c403955..fc09734e1 100644 --- a/caikit_nlp/resources/pretrained_model/hf_auto_causal_lm.py +++ b/caikit_nlp/resources/pretrained_model/hf_auto_causal_lm.py @@ -16,8 +16,7 @@ """ # Standard from collections.abc import Mapping -from copy import copy -from typing import Callable, List, Tuple, Union +from typing import List, Union # Third Party from transformers import ( @@ -26,11 +25,12 @@ DataCollatorForLanguageModeling, ) from transformers.models.auto import modeling_auto +import torch # First Party from caikit.core.data_model import DataStream +from caikit.core.exceptions import error_handler from caikit.core.modules import module -from caikit.core.toolkit import error_handler import alog # Local @@ -66,7 +66,10 @@ def tokenize_function( max_target_length: int, verbalizer: Union[None, str] = None, task_ids: Union[None, int] = None, - ) -> DataStream[BatchEncoding]: + use_seq2seq_approach: bool = True, + chunk_size: int = 128, + drop_remainder: bool = False, + ) -> Union[DataStream[BatchEncoding], BatchEncoding]: """Tokenization function to be used for causallm training; this function consumes a GenerationTrainRecord object and applies the verbalizer to it followed by the model tokenizer. Due to the nature of our training data with src/target seqs, @@ -86,11 +89,23 @@ def tokenize_function( Verbalizer to be rendered into each text. task_ids: Union[None, int] Task IDs to be used for multiprompt tuning. + use_seq2seq_approach: bool + Indicates whether or not we should use a sequence style approach + or use chunking parameters. + chunk_size: int + unsigned int value to be used for chunk size. + Only used if use_seq2seq_approach=True. + drop_remainder: bool + Whether or not to keep the residual as an extra chunk if the + total number of tokens is not divisible by the chunk size. + Only used if use_seq2seq_approach=True. Returns: - DataStream[transformers.tokenization_utils_base.BatchEncoding] - stream of encoded tokenization output corresponding to the input example. + Union[DataStream[BatchEncoding], BatchEncoding] + stream of encoded tokenization output corresponding to the input example + or a single batch encoding object containing 1+ tokenized results. """ + ### Things common to all Causal LM tokenization approaches # Extract the source & target from our provided inputs source, target = cls.decompose_example_io(example) # Determine if our mapped inputs are in batched mode or not @@ -104,50 +119,31 @@ def tokenize_function( source = ( source if verbalizer is None else render_verbalizer(verbalizer, example) ) - - source_ids = tokenizer(source, max_length=max_source_length, truncation=True) - target_ids = tokenizer(target, max_length=max_target_length, truncation=True) - - # Force everything to a list of batch encodings; for non-batch mode, this just - # puts it into a list. For batch mode, we get a list of batch encodings, - # allowing us to standardize subsequent processing a bit. - source_ids, num_target_samples = cls._force_to_batch_encoding_list( - source_ids, target_ids, batched_mode, task_ids + # Treat this as a seq2seq type problem. Note that this implementation is different + # from the seq2seq tokenization function even though it is conceptually similar due + # to sequence length / padding requirements assumed internally by causal LMs. + if use_seq2seq_approach: + return cls._causal_lm_padding_as_seq2seq( + tokenizer=tokenizer, + source=source, + target=target, + max_source_length=max_source_length, + max_target_length=max_target_length, + task_ids=task_ids, + ) + # Do causal language model chunking + return cls._causal_lm_as_chunked( + tokenizer=tokenizer, + source=source, + target=target, + max_source_length=max_source_length, + max_target_length=max_target_length, + batched_mode=batched_mode, + task_ids=task_ids, + chunk_size=chunk_size, + drop_remainder=drop_remainder, ) - def build_generator_func( - source_ids: BatchEncoding, num_target_samples: int - ) -> Callable: - """Builds a generator that can be applied to a single batch encoding and its - corresponding original number of target samples. - - source_ids: BatchEncoding - Source ID to generate different samples from. - num_target_samples: int - Number of target IDs; used for attention mask creation. - """ - - def single_generator_func(): - for idx in range(num_target_samples): - ret_source_ids = copy(source_ids) - ret_source_ids["attention_mask"] = cls._get_attention_mask( - source_ids, - idx, - num_target_samples, - ) - yield ret_source_ids - - return single_generator_func - - if not batched_mode: - return DataStream(build_generator_func(source_ids, num_target_samples)) - streams = [ - DataStream(build_generator_func(s_ids, n_target_samples)) - for s_ids, n_target_samples in zip(source_ids, num_target_samples) - ] - encoding_keys = source_ids[0].keys() - return cls._collapse_streams_into_encoding(streams, encoding_keys) - def _get_data_collator(self, **kwargs) -> "transformers.DataCollator": """Function to return appropriate data collator based on resource. @@ -158,6 +154,10 @@ def _get_data_collator(self, **kwargs) -> "transformers.DataCollator": NOTE: If mlm (masked language modeling) is not passed in kwargs, this function will automatically set it to `False`. + FIXME: This should be consolidated with what is in the prompt tuning + module, which currently does its own collator management outside of the + resource classes. + Args: **kwargs: All the keyword arguments passed to this function @@ -165,6 +165,7 @@ def _get_data_collator(self, **kwargs) -> "transformers.DataCollator": applicable to implemented data collator. Returns: transformers.DataCollator + Collator to be used for causal language modeling. """ applicable_args = ["mlm", "pad_to_multiple_of"] @@ -177,13 +178,92 @@ def _get_data_collator(self, **kwargs) -> "transformers.DataCollator": tokenizer=self._tokenizer, return_tensors="pt", **collator_kwargs ) + ### Tokenization strategy implementations + # Chunked causal language modeling + @classmethod + def _causal_lm_as_chunked( + cls, + tokenizer: "AutoTokenizer", + source: str, + target: str, + max_source_length: int, + max_target_length: int, + batched_mode: bool, + task_ids: Union[None, int], + chunk_size: int, + drop_remainder: bool, + ) -> Union[DataStream[BatchEncoding], BatchEncoding]: + """Given a source and target string, build the chunked concatenated sequence and formulate + the batch encoded chunks for the sequence. If running in batch mode, the chunks will be + collapsed into a single batch encoding for the whole sequence. Otherwise, each chunk will + placed in its own BatchEncoding and encapsulated within a datastream. + + Args: + tokenizer: AutoTokenizer + Tokenizer object to be applied to input records. + source: str + Raw source string. + target: str + Raw target string. + max_source_length: int + Maximum length for input sequences. + max_target_length: int + Maximum length for output sequences. + batched_mode: bool + Whether or not we should produce a stream of encodings or a single + encoding representing all of the chunked sequence. + task_ids: Union[None, int] + Task IDs to be used for multiprompt tuning. + chunk_size: int + unsigned int value to be used for chunk size. + drop_remainder: bool + Whether or not to keep the residual as an extra chunk if the + total number of tokens is not divisible by the chunk size. + + Returns: + Union[DataStream[BatchEncoding], BatchEncoding] + Encoded chunked sequence as a stream or batch encoding object. + """ + source_ids = tokenizer(source, max_length=max_source_length, truncation=True) + target_ids = tokenizer(target, max_length=max_target_length, truncation=True) + + # Force everything to a list of batch encodings; for non-batch mode, this just + # puts it into a list. For batch mode, we get a list of batch encodings, + # allowing us to standardize subsequent processing a bit. + # + # For example, given chunk size 2, we might have something like: + # [ + # {'input_ids': [31, 48], 'attention_mask': [1, 1]}, + # {'input_ids': [47, 1], 'attention_mask': [1, 1]}, + # ... + # ] + # (where the above objects are batch encodings, which are a subclass of dict) + source_id_chunks = cls._force_to_batch_encoding_list_of_chunks( + source_ids, target_ids, batched_mode, task_ids, chunk_size, drop_remainder + ) + + def generator_func(): + for chunk in source_id_chunks: + yield chunk + + chunk_stream = DataStream(generator_func) + # If it's batch mode, collapse down into one encoding batch object + if batched_mode: + return cls._collapse_stream_into_encoding(chunk_stream) + # Otherwise just produce the stream to be chained + # NOTE: it might be a good idea to deprecate this to force standardization + # onto using batch encodings the way that they are intended to be + return chunk_stream + @staticmethod - def _force_to_batch_encoding_list( + def _force_to_batch_encoding_list_of_chunks( source_ids: BatchEncoding, target_ids: BatchEncoding, batch_mode: bool, task_ids: Union[None, int], - ) -> Tuple[Union[BatchEncoding, List[BatchEncoding]], Union[int, List[int]]]: + chunk_size: int, + drop_remainder: bool, + ) -> List[BatchEncoding]: """Forces our inputs into either a single batch encoding (if we aren't running in batch mode), or a list of Batch Encodings. I.e., a list of dicts instead of a dict of lists. The primary reason that we do this is to allow us to easily map a common generator @@ -198,19 +278,29 @@ def _force_to_batch_encoding_list( Whether or not we are processing a batch. task_ids: Union[None, int] Optional task IDs for MPT to be propagated to produced encodings. + chunk_size: int + unsigned int value to be used for chunk size. + drop_remainder: bool + Whether or not to keep the residual as an extra chunk if the + total number of tokens is not divisible by the chunk size. Returns: - Tuple[Union[BatchEncoding, List[BatchEncoding]], Union[int, List]] + List[BatchEncoding] + List of batch encodings, each of which encapsulates the contents + of a single chunk. """ if not batch_mode: - source_ids["input_ids"] = source_ids.input_ids + target_ids.input_ids - source_ids["task_ids"] = task_ids - num_target_samples = len(target_ids.input_ids) - return source_ids, num_target_samples + HFAutoCausalLM._concatenate_encodings(source_ids, target_ids) + chunks = HFAutoCausalLM._split_encoding_into_chunks( + encoding=source_ids, + chunk_size=chunk_size, + drop_remainder=drop_remainder, + task_ids=task_ids, + ) + return chunks # Otherwise we need to expand the dict along its keys, # mapping all of its encapsulated objects to new items. encodings = [] - num_target_samples = [] id_keys = source_ids.keys() key = None error.value_check( @@ -218,49 +308,93 @@ def _force_to_batch_encoding_list( source_ids.keys(), "Source ID batch encoding must have keys", ) + for batch_idx in range(len(source_ids.input_ids)): new_encoding = BatchEncoding() for key in id_keys: - if key == "input_ids": - new_encoding[key] = ( - source_ids[key][batch_idx] + target_ids[key][batch_idx] - ) - else: - new_encoding[key] = source_ids[key][batch_idx] - num_target_samples.append(len(target_ids[key][batch_idx])) - new_encoding["task_ids"] = task_ids - encodings.append(new_encoding) - return encodings, num_target_samples + new_encoding[key] = ( + source_ids[key][batch_idx] + target_ids[key][batch_idx] + ) + chunks = HFAutoCausalLM._split_encoding_into_chunks( + encoding=new_encoding, + chunk_size=chunk_size, + drop_remainder=drop_remainder, + task_ids=task_ids, + ) + # Chunks are held as a list of lists + encodings += chunks + return encodings @staticmethod - def _get_attention_mask( - source_ids: BatchEncoding, idx: int, num_target_samples: int - ) -> List[int]: - """Get the attention mask for a given target token from some source encoding. + def _concatenate_encodings(left: BatchEncoding, right: BatchEncoding) -> None: + """Given two batch encodings, combine their entries into a single encoding. Args: - source_ids: BatchEncoding - Source encoding that requires an attention mask. - idx: int - Index of the output token we attend up to. - num_target_samples: int - Length of the original target seequence being considered. + left: BatchEncoding + Encoding representing left sequence, which will be updated in place. + Corresponds to source. + right: BatchEncoding + Encoding representing right sequence, which will be stacked onto the left + encoding. Corresponds to target. + """ + for k in left.keys(): + left[k].extend(right[k]) + + @staticmethod + def _split_encoding_into_chunks( + encoding: BatchEncoding, + chunk_size: int, + drop_remainder: bool, + task_ids: Union[None, int], + ) -> List[BatchEncoding]: + """Fetch the chunked batch encoding objects from the concatenated encoding. + + Args: + encoding: BatchEncoding + BatchEncoding holding the concatenated source/target for one example. + chunk_size: int + unsigned int value to be used for chunk size. + drop_remainder: bool + Whether or not to keep the residual as an extra chunk if the + total number of tokens is not divisible by the chunk size. + task_ids: Union[None, int] + Optional task IDs for MPT to be propagated to produced encodings. Returns: - List[int] - Binary attention mask. + List[BatchEncoding] + List of encodings, where each encoding represents one chunk. """ - return ( - source_ids["attention_mask"] - + [1] * (idx + 1) - + [0] * (num_target_samples - idx - 1) - ) + chunked_encodings = [] + # all encoding keys have the same length list values; we just use input ids + tok_len = len(encoding["input_ids"]) + # Build a batch encoding for every chunk; for each data, + # use the slice for all keys inside of the source_encoding. + if tok_len >= chunk_size: + slice_len = (tok_len // chunk_size) * chunk_size + # If we have a remainder and we don't want to drop it, add a new chunk + if not drop_remainder and slice_len != tok_len: + slice_len += chunk_size + # We just have one big chunk + else: + slice_len = tok_len + chunked_encodings = [ + BatchEncoding( + data={ + k: v[chunk_num : chunk_num + chunk_size] + for k, v in encoding.items() + } + ) + for chunk_num in range(0, slice_len, chunk_size) + ] + for enc in chunked_encodings: + enc["task_ids"] = task_ids + return chunked_encodings @staticmethod - def _collapse_streams_into_encoding( - streams: List[DataStream[BatchEncoding]], encoding_keys: "dict_keys" + def _collapse_stream_into_encoding( + stream: DataStream[BatchEncoding], ) -> BatchEncoding: - """Given a list of streams of batch encodings, collapse them back into + """Given a stream batch encodings, collapse them back into one encoding, i.e., the return value of the batch encoding. Args: @@ -271,14 +405,105 @@ def _collapse_streams_into_encoding( Returns: BatchEncoding - Collapsed batch encoding to be returned from tokenizatino func. + Collapsed batch encoding to be returned from tokenization func. """ + encoding_keys = None new_encoding = BatchEncoding() - for k in encoding_keys: - new_encoding[k] = [] # Now build the individual lists lists for each entry - for stream in streams: - for enc in stream: + for enc in stream: + # Initialize the existing keys in the new encoding + if encoding_keys is None: + encoding_keys = enc.keys() for k in encoding_keys: - new_encoding[k].append(enc[k]) + new_encoding[k] = [] + for k in encoding_keys: + new_encoding[k].append(enc[k]) return new_encoding + + # Causal language modeling as a sequence to sequence problem + @staticmethod + def _causal_lm_padding_as_seq2seq( + tokenizer: "AutoTokenizer", + source: str, + target: str, + max_source_length: int, + max_target_length: int, + task_ids: Union[None, int], + ) -> BatchEncoding: + """Tokenize the example as a seq2seq type problem; this is conceptually similar to + what seq2seq tokenization is doing, but some care needs be taken to ensure the labels + are the same length as the input sequence because of the shifting mechanism implemented + in most causal language models. + + Collator compatability is extremely important here; because we are setting the labels + directly, we should NOT use the causal lm collator, otherwise it will clobber it with a + shifted input sequence. + + Args: + tokenizer: AutoTokenizer + Tokenizer object to be applied to input records. + source: str + Raw source string. + target: str + Raw target string. + max_source_length: int + Maximum length for input sequences. + max_target_length: int + Maximum length for output sequences. + task_ids: Union[None, int] + Optional task IDs for MPT to be propagated to produced encodings. + Returns: + BatchEncoding + BatchEncoding object corresponding to this example, where the input_ids, + attention_mask, and labels all have the same length, i.e., + [max_source_length + max_target_length + 1]. + """ + IGNORE_ID = -100 + # ID of the token to append after our target string; this should generally be pad / EOS + FINAL_TOK_ID = tokenizer.eos_token_id + max_concat_length = max_source_length + max_target_length + 1 + + # Truncate based on max source or max target length before considering as a joined sequence + model_inputs = tokenizer(source, truncation=True, max_length=max_source_length) + labels = tokenizer(target, truncation=True, max_length=max_target_length + 1) + + # Combine the source + target strings into the source input IDs + # This makes the source and target the same length, and then masks the source out of the + # target IDs, and updates the length of the attention vector to be evenly spread on the + # whole combined sequence + sample_input_ids = model_inputs["input_ids"] + label_input_ids = labels["input_ids"] + [FINAL_TOK_ID] + model_inputs["input_ids"] = sample_input_ids + label_input_ids + labels["input_ids"] = [IGNORE_ID] * len(sample_input_ids) + label_input_ids + model_inputs["attention_mask"] = [1] * len(model_inputs["input_ids"]) + # Now we have to update everything to be the max length of the tokenizer, then pad & + # ensure all of the padded stuff we have added has attention weights of 0. + sample_input_ids = model_inputs[ + "input_ids" + ] # NOTE - combined source + target + + + label_input_ids = labels["input_ids"] + model_inputs = tokenizer.pad( + model_inputs, padding="max_length", max_length=max_concat_length + ) + + if tokenizer.padding_side.lower() == "left": + labels["input_ids"] = [IGNORE_ID] * ( + max_concat_length - len(sample_input_ids) + ) + label_input_ids + else: + labels["input_ids"] = label_input_ids + [IGNORE_ID] * ( + max_concat_length - len(sample_input_ids) + ) + + model_inputs["input_ids"] = torch.tensor( + model_inputs["input_ids"][:max_concat_length] + ) + model_inputs["attention_mask"] = torch.tensor( + model_inputs["attention_mask"][:max_concat_length] + ) + + labels["input_ids"] = torch.tensor(labels["input_ids"][:max_concat_length]) + model_inputs["labels"] = labels["input_ids"] + model_inputs["task_ids"] = task_ids + return model_inputs diff --git a/caikit_nlp/resources/pretrained_model/hf_auto_seq2seq_lm.py b/caikit_nlp/resources/pretrained_model/hf_auto_seq2seq_lm.py index 4a6c9ceb6..089dcfd9c 100644 --- a/caikit_nlp/resources/pretrained_model/hf_auto_seq2seq_lm.py +++ b/caikit_nlp/resources/pretrained_model/hf_auto_seq2seq_lm.py @@ -16,7 +16,7 @@ """ # Standard from collections.abc import Mapping -from typing import List, Union +from typing import Dict, List, Union # Third Party from torch.utils.data import IterableDataset @@ -29,12 +29,13 @@ from transformers.models.auto import modeling_auto # First Party +from caikit.core.exceptions import error_handler from caikit.core.modules import module -from caikit.core.toolkit import error_handler import alog # Local from ...data_model import GenerationTrainRecord, PromptOutputModelType +from ...toolkit.trainer_utils import log_step from ...toolkit.verbalizer_utils import render_verbalizer from .base import PretrainedModelBase @@ -44,6 +45,23 @@ IGNORE_ID = -100 +class LoggingTrainer(Seq2SeqTrainer): + def log(self, logs: Dict[str, float]) -> None: + """ + Log `logs` on the various objects watching training. + + Subclass and override this method to inject custom behavior. + + Args: + logs (`Dict[str, float]`): + The values to log. + """ + self.state = log_step(self.state, logs) + self.control = self.callback_handler.on_log( + self.args, self.state, self.control, logs + ) + + @module( id="6759e891-287b-405b-bd8b-54a4a4d51c25", name="HF Transformers Auto Seq2Seq LM", @@ -80,6 +98,7 @@ def get_trainer( train_dataset: IterableDataset, eval_dataset: Union[IterableDataset, None] = None, optimizers=(None, None), + model=None, **kwargs ): """ @@ -110,7 +129,12 @@ def get_trainer( # "generation_max_length": max_target_length, } - return Seq2SeqTrainer(self._model, training_args, **trainer_arguments) + # If extra model is provided, we will configure trainer + # with that model + if model: + return LoggingTrainer(model, training_args, **trainer_arguments) + + return LoggingTrainer(self._model, training_args, **trainer_arguments) def _get_data_collator(self, **kwargs): """Function to return appropriate data collator based on resource. diff --git a/caikit_nlp/toolkit/data_stream_wrapper.py b/caikit_nlp/toolkit/data_stream_wrapper.py index c486561bd..3d42a46c4 100644 --- a/caikit_nlp/toolkit/data_stream_wrapper.py +++ b/caikit_nlp/toolkit/data_stream_wrapper.py @@ -21,7 +21,7 @@ from torch.utils.data import IterableDataset # First Party -from caikit.core.toolkit import error_handler +from caikit.core.exceptions import error_handler import alog log = alog.use_channel("PEFT_PROMPT") diff --git a/caikit_nlp/toolkit/data_type_utils.py b/caikit_nlp/toolkit/data_type_utils.py index 959aa9383..7ee9c228a 100644 --- a/caikit_nlp/toolkit/data_type_utils.py +++ b/caikit_nlp/toolkit/data_type_utils.py @@ -20,7 +20,7 @@ # First Party from caikit import get_config -from caikit.core.toolkit import error_handler +from caikit.core.exceptions import error_handler import alog log = alog.use_channel("DATA_UTIL") diff --git a/caikit_nlp/toolkit/task_specific_utils.py b/caikit_nlp/toolkit/task_specific_utils.py index 0d42fd350..e4d197dba 100644 --- a/caikit_nlp/toolkit/task_specific_utils.py +++ b/caikit_nlp/toolkit/task_specific_utils.py @@ -13,7 +13,7 @@ # limitations under the License. # First Party -from caikit.core.toolkit import error_handler +from caikit.core.exceptions import error_handler from caikit.interfaces.nlp.data_model import ClassificationTrainRecord import alog diff --git a/caikit_nlp/toolkit/text_generation/__init__.py b/caikit_nlp/toolkit/text_generation/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/caikit_nlp/toolkit/text_generation/model_run_utils.py b/caikit_nlp/toolkit/text_generation/model_run_utils.py index d5ece8da9..d51a56a9d 100644 --- a/caikit_nlp/toolkit/text_generation/model_run_utils.py +++ b/caikit_nlp/toolkit/text_generation/model_run_utils.py @@ -19,14 +19,15 @@ # Third Party from peft.peft_model import PeftModel -from transformers import StoppingCriteria, TextStreamer +from transformers import AutoModel, AutoTokenizer, StoppingCriteria, TextStreamer import numpy as np import torch # First Party from caikit.core.data_model.producer import ProducerId -from caikit.core.toolkit.errors import error_handler +from caikit.core.exceptions import error_handler from caikit.interfaces.nlp.data_model import ( + FinishReason, GeneratedTextResult, GeneratedTextStreamResult, TokenStreamDetails, @@ -34,7 +35,7 @@ import alog # Local -from ...data_model import ExponentialDecayLengthPenalty +from caikit_nlp.data_model import ExponentialDecayLengthPenalty log = alog.use_channel("RUN_UTILS") error = error_handler.get(log) @@ -81,9 +82,6 @@ The value used to modulate the next token probabilities. Only applicable when decoding_method is SAMPLING. Default: 1.0 - means disabled - equivalent to 1.0 - seed: numpy.uint64 - Random seed to control sampling. Only applicable when decoding_method - is SAMPLING. Default: None repetition_penalty: float The more a token is used within generation the more it is penalized to not be picked in successive generation passes. @@ -100,6 +98,9 @@ of exponential decay stop_sequences: List[str] List of strings to be used as stopping criteria + seed: numpy.uint64 + Random seed to control sampling. Only applicable when decoding_method + is SAMPLING. Default: None """ @@ -111,6 +112,7 @@ def on_finalized_text(self, text: str, stream_end: bool = False): class SequenceStoppingCriteria(StoppingCriteria): + # pylint: disable-next=super-init-not-called # false positive: StoppingCriteria is an abc and has no __init__ def __init__(self, target_sequence_ids): self.target_sequence_ids = target_sequence_ids @@ -130,10 +132,10 @@ def __iter__(self): def generate_text_func( - model, - tokenizer, + model: "Union[PeftModel, AutoModel]", + tokenizer: "AutoTokenizer", producer_id: ProducerId, - eos_token: str, + eos_token: Optional[str], text: str, max_new_tokens: Optional[int] = 20, min_new_tokens: Optional[int] = 0, @@ -190,7 +192,7 @@ def generate_text_func( ) inputs = {k: v.to(model.device) for k, v in tok_tensors.items()} - input_token_count = len(tok_tensors) + input_token_count = tok_tensors["input_ids"].size(1) gen_optional_params = __process_gen_args( tokenizer, @@ -233,12 +235,20 @@ def generate_text_func( for g in generate_ids ] - if generate_ids[0][-1].item() == eos_token: - finish_reason = "EOS_TOKEN" - elif generate_ids.size(1) - 1 == max_new_tokens: - finish_reason = "MAX_TOKENS" + if (eos_token and tokenizer.decode(generate_ids[0, -1].item()) == eos_token) or ( + generate_ids[0, -1] == tokenizer.eos_token_id + ): + finish_reason = FinishReason.EOS_TOKEN + elif ("stopping_criteria" in gen_optional_params) and ( + gen_optional_params["stopping_criteria"]( + generate_ids, + None, # scores, unused by SequenceStoppingCriteria + ) + ): + finish_reason = FinishReason.STOP_SEQUENCE else: - finish_reason = "OTHER" + finish_reason = FinishReason.MAX_TOKENS + return GeneratedTextResult( generated_tokens=token_count, generated_text=preds[0], diff --git a/caikit_nlp/toolkit/text_generation/tgis_utils.py b/caikit_nlp/toolkit/text_generation/tgis_utils.py index 4a4d9f529..50728f0db 100644 --- a/caikit_nlp/toolkit/text_generation/tgis_utils.py +++ b/caikit_nlp/toolkit/text_generation/tgis_utils.py @@ -17,7 +17,7 @@ from typing import Iterable # First Party -from caikit.core.toolkit import error_handler +from caikit.core.exceptions import error_handler from caikit.interfaces.nlp.data_model import ( GeneratedTextResult, GeneratedTextStreamResult, @@ -121,19 +121,29 @@ def validate_inf_params( ) error.value_check( - "", temperature >= 0.05, "temperature must be >= 0.05" + "", + not temperature or temperature >= 0.05, + "temperature must be >= 0.05", ) error.value_check( - "", top_p > 0.0 and top_p <= 1.0, "top_p must be > 0.0 and <= 1.0" + "", + not top_p or 0 < top_p <= 1.0, + "top_p must be > 0.0 and <= 1.0", ) - error.value_check("", top_k >= 0, "top_k must be strictly positive") + error.value_check( + "", not top_k or top_k >= 0, "top_k must be strictly positive" + ) - error.value_check("", typical_p <= 1.0, "typical_p must be <= 1.0") + error.value_check( + "", not typical_p or typical_p <= 1.0, "typical_p must be <= 1.0" + ) error.value_check( - "", repetition_penalty > 0.0, "repetition_penalty must be > 0.0" + "", + not repetition_penalty or repetition_penalty > 0.0, + "repetition_penalty must be > 0.0", ) if exponential_decay_length_penalty: @@ -150,16 +160,19 @@ def validate_inf_params( ) if decoding_method == "GREEDY" and ( - temperature != 1 or top_k != 0 or top_p != 1 or seed + temperature not in (1, None) + or top_k not in (0, None) + or top_p not in (1, None) + or seed ): raise ValueError( - "sampling parameters (temperature/top_k/top_p/typical_p/seed) aren't applicable in greedy decoding mode" + "sampling parameters (temperature/top_k/top_p/typical_p/seed) aren't " + "applicable in greedy decoding mode" ) def get_params( preserve_input_text, - eos_token, max_new_tokens, min_new_tokens, truncate_input_tokens, @@ -177,8 +190,6 @@ def get_params( """Get generation parameters Args: - eos_token: str - A special token representing the end of a sentence. {} """.format( GENERATE_FUNCTION_TGIS_ARGS @@ -205,7 +216,7 @@ def get_params( token_ranks=True, ) stopping = generation_pb2.StoppingCriteria( - stop_sequences=stop_sequences or [eos_token] if eos_token else None, + stop_sequences=stop_sequences, max_new_tokens=max_new_tokens, min_new_tokens=min_new_tokens, time_limit_millis=int(max_time * 1000) if max_time else None, @@ -313,7 +324,6 @@ def unary_generate( params = get_params( preserve_input_text=preserve_input_text, - eos_token=self.eos_token, max_new_tokens=max_new_tokens, min_new_tokens=min_new_tokens, truncate_input_tokens=truncate_input_tokens, @@ -423,7 +433,6 @@ def stream_generate( params = get_params( preserve_input_text=preserve_input_text, - eos_token=self.eos_token, max_new_tokens=max_new_tokens, min_new_tokens=min_new_tokens, truncate_input_tokens=truncate_input_tokens, diff --git a/caikit_nlp/toolkit/text_generation/training_utils.py b/caikit_nlp/toolkit/text_generation/training_utils.py new file mode 100644 index 000000000..244d2c550 --- /dev/null +++ b/caikit_nlp/toolkit/text_generation/training_utils.py @@ -0,0 +1,244 @@ +# Copyright The Caikit Authors +# +# 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. +"""Utility script that contains logic for training""" + +# Standard +from typing import List, Optional, Union + +# Third Party +from datasets import Dataset +from datasets import IterableDataset as TransformersIterableDataset +from transformers import AutoTokenizer +import torch + +# First Party +from caikit.core.data_model import DataStream +from caikit.core.toolkit import error_handler +import alog + +# Local +from ...data_model import GenerationTrainRecord +from ...resources.pretrained_model import PretrainedModelBase + +log = alog.use_channel("TXTGEN_TRN_UTLS") +error = error_handler.get(log) + +# Below list is taken from +# https://huggingface.co/docs/transformers/main/en/main_classes/trainer#transformers.TrainingArguments +ALLOWED_TRAINING_ARGS = { + "weight_decay", + "adam_beta1", + "adam_beta2", + "adam_epsilon", + "max_grad_norm", + "lr_scheduler_type", + "warmup_ratio", + "warmup_steps", + "use_ipex", + "disable_tqdm", + "label_names", + "optim", + "optim_args", + "group_by_length", + "dataloader_pin_memory", + "gradient_checkpointing", + "full_determinism", +} + +# Create trainer arguments +def collect_trainer_arguments( + torch_dtype, + output_dir, + batch_size, + num_epochs, + random_seed, + learning_rate, + max_steps, + silence_progress_bars=True, + **kwargs +): + """Utility function to return processed HF Trainer argument dictionary""" + + # NOTE: Following is not exhaustive list of all parameters + # for all dtypes + if torch_dtype == torch.float16: + dtype_based_params = { + "fp16": True, + } + elif torch_dtype == torch.bfloat16: + dtype_based_params = { + "bf16": True, + } + else: + # default to float32 + dtype_based_params = {} + + return { + # trainer settings + "output_dir": output_dir, + # NOTE: We have disabled evaluation for now + "do_eval": False, + "do_train": True, + "no_cuda": not torch.cuda.is_available(), + # NOTE: This is explicitly set to false since it will + # negatively impact the performance + "full_determinism": False, + # logging configuration + "logging_strategy": "steps", + "logging_steps": 1, # logging at every step + "disable_tqdm": silence_progress_bars, + # computation configurations + "seed": random_seed, + "per_device_train_batch_size": batch_size, + "per_device_eval_batch_size": batch_size, + "num_train_epochs": num_epochs, + "learning_rate": learning_rate, + "weight_decay": 0.01, + "save_total_limit": 3, + "gradient_checkpointing": True, + # huggingface configurations + "push_to_hub": False, + # dataset configurations + "remove_unused_columns": True, + "dataloader_pin_memory": False, + # Required for iterable dataset + "max_steps": max_steps, + # others + "auto_find_batch_size": True, + **dtype_based_params, + **kwargs, + } + + +def preprocess_function( + base_model: PretrainedModelBase, + train_stream: DataStream[GenerationTrainRecord], + tokenizer: AutoTokenizer, + max_source_length: int, + max_target_length: int, + shuffle: bool, + use_iterable_dataset: bool, + random_seed: int, + task_ids: Optional[List[int]] = None, +): + """Pre-process each example to get it prepared for training.""" + dataset_type = TransformersIterableDataset if use_iterable_dataset else Dataset + log.debug("Loading dataset class: [%s]", dataset_type.__name__) + fn_kwargs = { + "tokenizer": tokenizer, + "max_source_length": max_source_length, + "max_target_length": max_target_length, + } + if task_ids is not None: + fn_kwargs["task_ids"] = task_ids + + # TODO: Add check for empty training stream + dataset = dataset_type.from_generator( + get_record, gen_kwargs={"train_stream": train_stream} + ) + mapped_dataset = dataset.map( + base_model.tokenize_function, + fn_kwargs=fn_kwargs, + # For now, we hardcode to False, since causal LM chunking is not exposed yet + batched=False, + # batched=base_model.REQUIRES_TOKEN_UNWRAPPING, + # Drop the input / output columns; we need to do this for dimensions to play + # happily when operating on batched inputs for causal language modeling. + remove_columns=["input", "output"], + ) + + if shuffle: + log.debug("Shuffling the dataset") + return mapped_dataset.shuffle(seed=random_seed) + + return mapped_dataset + + +def launch_training( + base_model, + training_dataset, + training_args, + checkpoint_dir, + caikit_resource=None, + tokenizer=None, +) -> None: + """Utility function to wrap trainer and execute training""" + # If we have a caikit resource, grab the trainer through it + if caikit_resource is not None: + trainer = caikit_resource.get_trainer( + train_dataset=training_dataset, model=base_model, **training_args + ) + else: + # If trainer is not provided fetch it from base_model + if hasattr(base_model, "get_trainer"): + trainer = base_model.get_trainer( + train_dataset=training_dataset, **training_args + ) + else: + error("", "could not resolve trainer. Check base model type!") + + # Start training via Trainer.train function + result = trainer.train() + + # Log the output of the training. This will include stats about training + log.info("", "Training completed. Summary: {}".format(result)) + + # save the model temporarily and reload it + # this is done, since otherwise the model might be distributed in different + # devices, in which case its better to use trainer's `prediction_step` + # functions, but then, they don't always give API similar to `generate` + # and thus cause incompatibilities in `run` function + trainer.save_state() + trainer.save_model(checkpoint_dir) + + # save tokenizer explicitly + if hasattr(base_model, "tokenizer"): + base_model.tokenizer.save_pretrained(checkpoint_dir) + elif tokenizer: + tokenizer.save_pretrained(checkpoint_dir) + else: + log.warning("", "Cannot save tokenizer as not available to train function.") + + # Below will return log history but launch will automatically attach rank to it. + # if started in distributed fashion + return trainer.state.log_history + + +def infer_max_steps( + num_epochs: int, + batch_size: int, + training_dataset: Union[Dataset, TransformersIterableDataset], +): + # Calculate the number of samples that we have + if isinstance(training_dataset, Dataset): + data_len = len(training_dataset) + else: + data_len = 0 + for _ in training_dataset: + data_len += 1 + # Figure out how many batches we'll have per epoch + num_batches = data_len // batch_size + # Assume drop_last=False; in general, this doesn't really matter. + # We mostly do this to avoid strange behavior when the dataset + # size is smaller than the batch size. + if num_batches != (data_len * batch_size): + num_batches += 1 + num_steps = num_batches * num_epochs + log.debug("Number of inferred steps: [%s]", num_steps) + return num_steps + + +def get_record(train_stream): + for data in train_stream: + yield {"input": data.input, "output": data.output} diff --git a/caikit_nlp/toolkit/trainer_utils.py b/caikit_nlp/toolkit/trainer_utils.py new file mode 100644 index 000000000..3ce41d581 --- /dev/null +++ b/caikit_nlp/toolkit/trainer_utils.py @@ -0,0 +1,84 @@ +# Copyright The Caikit Authors +# +# 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. +"""Contains toolkit functionality for huggingface Trainer""" +# Standard +from datetime import datetime + +# Third Party +import torch + +# First Party +from caikit import get_config +from caikit.core.data_model import DataStream +from caikit.core.exceptions import error_handler +import alog + +log = alog.use_channel("TRNR_UTILS") +error = error_handler.get(log) + + +def validate_training_data(train_stream: DataStream, model_name: str, module_id: str): + + global_default = get_config().training_data_limit.__default__ + module_default = ( + get_config() + .training_data_limit.get(module_id, {}) + .get("__default__", global_default) + ) + + max_num_examples = ( + get_config() + .training_data_limit.get(module_id, {}) + .get(model_name, module_default) + ) + + if max_num_examples > 0: + error.value_check( + "", + len(train_stream) <= max_num_examples, + "Number of examples larger than maximum number of examples allowed for this model", + ) + + +def log_step(state, logs): + if state.epoch is not None: + logs["epoch"] = round(state.epoch, 2) + + # Get Rank + if torch.distributed.is_initialized(): + rank = torch.distributed.get_rank() + else: + rank = 0 + + if "loss" in logs: + if state.epoch is not None: + logs["epoch"] = round(state.epoch, 2) + + log.debug( + "process rank: {} loss: {} step: {}".format( + rank, float(logs["loss"]), state.global_step + ) + ) + output = { + "epoch": float(logs["epoch"]), + "step": state.global_step, + "value": float(logs["loss"]), + "timestamp": datetime.isoformat(datetime.now()), + } + state.log_history.append(output) + else: + output = {**logs, **{"step": state.global_step}} + state.log_history.append(output) + + return state diff --git a/caikit_nlp/toolkit/verbalizer_utils.py b/caikit_nlp/toolkit/verbalizer_utils.py index f0a888c45..86b3dc2e3 100644 --- a/caikit_nlp/toolkit/verbalizer_utils.py +++ b/caikit_nlp/toolkit/verbalizer_utils.py @@ -15,7 +15,7 @@ import re # First Party -from caikit.core.toolkit import error_handler +from caikit.core.exceptions import error_handler import alog log = alog.use_channel("VERBALIZER_UTIL") diff --git a/caikit_nlp/version.py b/caikit_nlp/version.py new file mode 100644 index 000000000..e88d411d4 --- /dev/null +++ b/caikit_nlp/version.py @@ -0,0 +1,7 @@ +# pylint: disable=unused-import +try: + # Local + from ._version import __version__, __version_tuple__ +except ImportError: + __version__ = "unknown" + version_tuple = (0, 0, __version__) diff --git a/examples/evaluate_model.py b/examples/evaluate_model.py index e74601dfe..902690df6 100644 --- a/examples/evaluate_model.py +++ b/examples/evaluate_model.py @@ -65,11 +65,25 @@ def parse_args() -> argparse.Namespace: help="JSON file to dump raw source / target texts to.", default="model_preds.json", ) + parser.add_argument( + "--max_new_tokens", + help="Maximum number of new tokens to be generated", + type=int, + default=20, + ) + parser.add_argument( + "--truncate_input_tokens", + help="Number of allowed input tokens (no truncation=0)", + type=int, + default=0, + ) args = parser.parse_args() return args -def get_model_preds_and_references(model, validation_stream): +def get_model_preds_and_references( + model, validation_stream, max_new_tokens, truncate_input_tokens +): """Given a model & a validation stream, run the model against every example in the validation stream and compare the outputs to the target/output sequence. @@ -79,6 +93,10 @@ def get_model_preds_and_references(model, validation_stream): validation_stream: DataStream[GenerationTrainRecord] Validation stream with labeled targets that we want to compare to our model's predictions. + max_new_tokens: int + Max number of new tokens to be generated, i.e., output limit + truncate_input_tokens: int + Number of allowed input tokens, i.e., input limit Returns: Tuple(List) @@ -90,7 +108,11 @@ def get_model_preds_and_references(model, validation_stream): for datum in tqdm(validation_stream): # Local .run() currently prepends the input text to the generated string; # Ensure that we're just splitting the first predicted token & beyond. - raw_model_text = model.run(datum.input).generated_text + raw_model_text = model.run( + datum.input, + max_new_tokens=max_new_tokens, + truncate_input_tokens=truncate_input_tokens, + ).generated_text parse_pred_text = raw_model_text.split(datum.input)[-1].strip() model_preds.append(parse_pred_text) targets.append(datum.output) @@ -120,15 +142,15 @@ def export_model_preds(preds_file, predictions, validation_stream, verbalizer): """ pred_objs = [] for pred, record in zip(predictions, validation_stream): - src, target = record.input, record.output - pred_objs.append( - { - "source": record.input, - "target": record.output, - "predicted_target": pred, - "verbalized_source": render_verbalizer(verbalizer, record), - } - ) + res = { + "source": record.input, + "target": record.output, + "predicted_target": pred, + } + if verbalizer is not None: + res["verbalized_source"] = render_verbalizer(verbalizer, record) + pred_objs.append(res) + with open(preds_file, "w") as jfile: json.dump(pred_objs, jfile, indent=4, sort_keys=True) @@ -153,14 +175,19 @@ def export_model_preds(preds_file, predictions, validation_stream, verbalizer): # Run the data through the model; save the predictions & references print_colored("Getting model predictions...") - predictions, references = get_model_preds_and_references(model, validation_stream) + predictions, references = get_model_preds_and_references( + model, validation_stream, args.max_new_tokens, args.truncate_input_tokens + ) print_colored( "Exporting model preds, source, verbalized source, and ground truth targets to {}".format( args.preds_file ) ) export_model_preds( - args.preds_file, predictions, validation_stream, model.verbalizer + args.preds_file, + predictions, + validation_stream, + getattr(model, "verbalizer", None), ) for metric_func in metric_funcs: diff --git a/examples/run_peft_tuning.py b/examples/run_peft_tuning.py index 305bd7797..52c42b96c 100644 --- a/examples/run_peft_tuning.py +++ b/examples/run_peft_tuning.py @@ -395,7 +395,9 @@ def show_experiment_configuration(args, dataset_info, model_type) -> None: train_stream = subsample_stream(train_stream, args.num_shots) # Init the resource & Build the tuning config from our dataset/arg info print_colored("[Loading the base model resource...]") - base_model = model_type.bootstrap(args.model_name, tokenizer_name=args.model_name) + base_model = model_type.bootstrap( + args.model_name, tokenizer_name=args.model_name, torch_dtype=args.torch_dtype + ) tuning_config = build_tuning_config(args, dataset_info) # Then actually train the model & save it print_colored("[Starting the training...]") @@ -408,7 +410,7 @@ def show_experiment_configuration(args, dataset_info, model_type) -> None: max_target_length=args.max_target_length, tuning_type=args.tuning_type, num_epochs=args.num_epochs, - lr=args.learning_rate, + learning_rate=args.learning_rate, batch_size=args.batch_size, verbalizer=dataset_info.verbalizer, silence_progress_bars=not args.verbose, diff --git a/pyproject.toml b/pyproject.toml index 139ad5a6a..e99ce0b0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,11 +1,11 @@ [build-system] -requires = ["flit_core >=3.2,<4"] -build-backend = "flit_core.buildapi" +requires = [ + "setuptools>=60", + "setuptools-scm>=8.0"] [project] name = "caikit-nlp" -# Not the actual current version: overwritten by CI -version = "0.0.1" +dynamic = ["version"] description = "Caikit NLP" license = {text = "Apache-2.0"} readme = "README.md" @@ -14,7 +14,7 @@ classifiers=[ "License :: OSI Approved :: Apache Software License" ] dependencies = [ - "caikit[runtime-grpc,runtime-http]>=0.18.1,<0.21.0", + "caikit[runtime-grpc,runtime-http]>=0.24.0,<0.25.0", "caikit-tgis-backend>=0.1.17,<0.2.0", # TODO: loosen dependencies "accelerate>=0.22.0", @@ -32,8 +32,16 @@ dependencies = [ # which broke caikit-nlp build. peft hasn't released newer version yet, so to get # the build fix, we pulling peft from main branch commit. In future, we will pull PEFT from # pypi - "peft@git+https://github.com/huggingface/peft.git#8c17d556a8fe9522e10d73d7bd3fad46a6ecae14" + "peft@git+https://github.com/huggingface/peft.git@8c17d556a8fe9522e10d73d7bd3fad46a6ecae14" ] +[tool.setuptools.packages.find] +exclude = ["tests", "tests.*"] +namespaces = false + + +[tool.setuptools_scm] +version_file = "caikit_nlp/_version.py" + [project.urls] Source = "https://github.com/caikit/caikit-nlp" diff --git a/setup_requirements.txt b/setup_requirements.txt index f4de449b0..207774426 100644 --- a/setup_requirements.txt +++ b/setup_requirements.txt @@ -1,2 +1,2 @@ tox>=4.4.2,<5 -build>=0.10.0,<1.0 \ No newline at end of file +build>=0.10.0,<2.0 \ No newline at end of file diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py index 6ddfd7b6e..8e4ed9f57 100644 --- a/tests/fixtures/__init__.py +++ b/tests/fixtures/__init__.py @@ -36,6 +36,18 @@ SEQ2SEQ_LM_MODEL = os.path.join(TINY_MODELS_DIR, "T5ForConditionalGeneration") +dummy_train_stream = caikit.core.data_model.DataStream.from_iterable( + [ + caikit_nlp.data_model.GenerationTrainRecord( + input="@foo what a cute dog!", output="no complaint" + ), + caikit_nlp.data_model.GenerationTrainRecord( + input="@bar this is the worst idea ever.", output="complaint" + ), + ] +) + + @pytest.fixture() def set_cpu_device(request): """Fixture to set default cuda device. @@ -108,14 +120,14 @@ def models_cache_dir(request): ### Fixtures for grabbing a randomly initialized model to test interfaces against ## Causal LM -@pytest.fixture +@pytest.fixture(scope="session") def causal_lm_train_kwargs(): """Get the kwargs for a valid train call to a Causal LM.""" model_kwargs = { "base_model": HFAutoCausalLM.bootstrap( model_name=CAUSAL_LM_MODEL, tokenizer_name=CAUSAL_LM_MODEL ), - "train_stream": caikit.core.data_model.DataStream.from_iterable([]), + "train_stream": dummy_train_stream, "num_epochs": 0, "tuning_config": caikit_nlp.data_model.TuningConfig( num_virtual_tokens=8, prompt_tuning_init_text="hello world" @@ -124,7 +136,7 @@ def causal_lm_train_kwargs(): return model_kwargs -@pytest.fixture +@pytest.fixture(scope="session") def causal_lm_dummy_model(causal_lm_train_kwargs): """Train a Causal LM dummy model.""" return caikit_nlp.modules.text_generation.PeftPromptTuning.train( @@ -132,7 +144,7 @@ def causal_lm_dummy_model(causal_lm_train_kwargs): ) -@pytest.fixture +@pytest.fixture(scope="session") def saved_causal_lm_dummy_model(causal_lm_dummy_model): """Give a path to a saved dummy model that can be loaded""" with tempfile.TemporaryDirectory() as workdir: @@ -142,14 +154,14 @@ def saved_causal_lm_dummy_model(causal_lm_dummy_model): ## Seq2seq -@pytest.fixture +@pytest.fixture(scope="session") def seq2seq_lm_train_kwargs(): """Get the kwargs for a valid train call to a Causal LM.""" model_kwargs = { "base_model": HFAutoSeq2SeqLM.bootstrap( model_name=SEQ2SEQ_LM_MODEL, tokenizer_name=SEQ2SEQ_LM_MODEL ), - "train_stream": caikit.core.data_model.DataStream.from_iterable([]), + "train_stream": dummy_train_stream, "num_epochs": 0, "tuning_config": caikit_nlp.data_model.TuningConfig( num_virtual_tokens=16, prompt_tuning_init_text="hello world" @@ -158,7 +170,7 @@ def seq2seq_lm_train_kwargs(): return model_kwargs -@pytest.fixture +@pytest.fixture(scope="session") def seq2seq_lm_dummy_model(seq2seq_lm_train_kwargs): """Train a Seq2Seq LM dummy model.""" return caikit_nlp.modules.text_generation.PeftPromptTuning.train( diff --git a/tests/model_management/test_tgis_auto_finder.py b/tests/model_management/test_tgis_auto_finder.py index 6638e8659..13a73da7f 100644 --- a/tests/model_management/test_tgis_auto_finder.py +++ b/tests/model_management/test_tgis_auto_finder.py @@ -254,6 +254,7 @@ def test_bad_tgis_connection(): } }, "test_connections": True, + "connect_timeout": 5, }, } ], diff --git a/tests/modules/text_generation/test_peft_config.py b/tests/modules/text_generation/test_peft_config.py index 2680ffb14..7f037c6f5 100644 --- a/tests/modules/text_generation/test_peft_config.py +++ b/tests/modules/text_generation/test_peft_config.py @@ -7,12 +7,21 @@ # Local from caikit_nlp.data_model import TuningConfig -from caikit_nlp.modules.text_generation.peft_config import TuningType, get_peft_config +from caikit_nlp.modules.text_generation import TextGeneration +from caikit_nlp.modules.text_generation.peft_config import ( + TuningType, + get_peft_config, + resolve_base_model, +) +from caikit_nlp.resources.pretrained_model import HFAutoSeq2SeqLM from tests.fixtures import ( + SEQ2SEQ_LM_MODEL, + TINY_MODELS_DIR, causal_lm_dummy_model, causal_lm_train_kwargs, seq2seq_lm_dummy_model, seq2seq_lm_train_kwargs, + temp_config, ) @@ -63,3 +72,34 @@ def test_get_peft_config(train_kwargs, dummy_model, request): assert peft_config.task_type == dummy_resource.TASK_TYPE assert peft_config.prompt_tuning_init == tuning_config.prompt_tuning_init_method assert peft_config.prompt_tuning_init_text == tuning_config.prompt_tuning_init_text + + +def test_resolve_model_with_invalid_path_raises(): + """Test passing invalid path to resolve_model function raises""" + + invalid_base_model = "path/../../important" + with pytest.raises(ValueError): + resolve_base_model(invalid_base_model, None, "foo") + + +def test_resolve_model_with_valid_folder_path(): + """Test passing valid folder path to resolve_model function works""" + + model = resolve_base_model(SEQ2SEQ_LM_MODEL, TextGeneration, "float32") + + assert isinstance(model, HFAutoSeq2SeqLM) + + +def test_resolve_model_works_preloaded_model(): + + base_model = HFAutoSeq2SeqLM.bootstrap(SEQ2SEQ_LM_MODEL) + resolved_model = resolve_base_model(base_model, TextGeneration, "float32") + assert isinstance(resolved_model, HFAutoSeq2SeqLM) + + +def test_resolve_model_with_different_base_path_works(): + + base_model_name = "T5ForConditionalGeneration" + with temp_config(base_models_dir=TINY_MODELS_DIR): + resolved_model = resolve_base_model(base_model_name, TextGeneration, "float32") + assert isinstance(resolved_model, HFAutoSeq2SeqLM) diff --git a/tests/modules/text_generation/test_peft_prompt_tuning.py b/tests/modules/text_generation/test_peft_prompt_tuning.py index 1d600305b..4989ad8d4 100644 --- a/tests/modules/text_generation/test_peft_prompt_tuning.py +++ b/tests/modules/text_generation/test_peft_prompt_tuning.py @@ -33,6 +33,7 @@ seq2seq_lm_dummy_model, seq2seq_lm_train_kwargs, set_cpu_device, + temp_config, ) import caikit_nlp @@ -69,13 +70,13 @@ def test_save_log_loss_file(causal_lm_dummy_model): """Ensure saving a model saves the log loss file""" with tempfile.TemporaryDirectory() as model_dir: causal_lm_dummy_model.save(model_dir, save_base_model=False) - assert os.path.isfile( - os.path.join( - model_dir, - caikit_nlp.modules.text_generation.peft_prompt_tuning.TRAINING_LOSS_LOG_FILENAME, - ) + file_path = os.path.join( + model_dir, + caikit_nlp.modules.text_generation.peft_prompt_tuning.TRAINING_LOSS_LOG_FILENAME, ) + assert os.path.isfile(file_path) + def test_run_model(causal_lm_dummy_model): """Ensure that we can run a model and get the right type out.""" @@ -88,17 +89,20 @@ def test_run_stream_out_model(causal_lm_dummy_model): pred_stream = causal_lm_dummy_model.run_stream_out("This text doesn't matter") assert isinstance(pred_stream, Iterable) for pred in pred_stream: - print(pred) assert isinstance(pred, GeneratedTextStreamResult) -def test_verbalizer_rendering(causal_lm_dummy_model): +def test_verbalizer_rendering(causal_lm_dummy_model, monkeypatch): """Ensure that our model renders its verbalizer text correctly before calling tokenizer.""" # Mock the tokenizer; we want to make sure its inputs are rendered properly - causal_lm_dummy_model.tokenizer = mock.Mock( - side_effect=RuntimeError("Tokenizer is a mock!"), - # Set eos token property to be attribute of tokenizer - eos_token="", + monkeypatch.setattr( + causal_lm_dummy_model, + "tokenizer", + mock.Mock( + side_effect=RuntimeError("Tokenizer is a mock!"), + # Set eos token property to be attribute of tokenizer + eos_token="", + ), ) input_text = "This text doesn't matter" causal_lm_dummy_model.verbalizer = " | {{input}} |" @@ -143,7 +147,7 @@ def test_train_model(causal_lm_train_kwargs, set_cpu_device): ), ] ), - "torch_dtype": torch.bfloat16, + "torch_dtype": torch.float32, "device": "cpu", } causal_lm_train_kwargs.update(patch_kwargs) @@ -201,7 +205,7 @@ def test_train_model_classification_record(causal_lm_train_kwargs, set_cpu_devic ), ] ), - "torch_dtype": torch.bfloat16, + "torch_dtype": torch.float32, "device": "cpu", } causal_lm_train_kwargs.update(patch_kwargs) @@ -216,7 +220,7 @@ def test_train_model_classification_record(causal_lm_train_kwargs, set_cpu_devic def test_prompt_output_types(causal_lm_train_kwargs): - # Try training a model with outpout_model_types set to a list of strings + # Try training a model with output_model_types set to a list of strings patch_kwargs = { "num_epochs": 1, "verbalizer": "Tweet text : {{input}} Label : ", @@ -230,7 +234,7 @@ def test_prompt_output_types(causal_lm_train_kwargs): ), ] ), - "torch_dtype": torch.bfloat16, + "torch_dtype": torch.float32, "device": "cpu", "tuning_config": caikit_nlp.data_model.TuningConfig( num_virtual_tokens=8, @@ -257,6 +261,19 @@ def test_prompt_output_types(causal_lm_train_kwargs): assert model +def test_error_empty_stream(causal_lm_train_kwargs): + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable([]), + } + causal_lm_train_kwargs.update(patch_kwargs) + with pytest.raises(ValueError): + caikit_nlp.modules.text_generation.PeftPromptTuning.train( + **causal_lm_train_kwargs + ) + + ### Implementation details # These tests can probably be removed and tested directly through .save() once # full seq2seq support is completed and verified. @@ -399,3 +416,240 @@ def test_run_exponential_decay_len_penatly_object(causal_lm_dummy_model): exponential_decay_length_penalty=penalty, ) assert isinstance(pred, GeneratedTextResult) + + +def test_train_with_data_validation_raises(causal_lm_train_kwargs, set_cpu_device): + """Check if we are able to throw error for when number of examples are more than configured limit""" + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable( + [ + ClassificationTrainRecord( + text="@foo what a cute dog!", labels=["no complaint"] + ), + ClassificationTrainRecord( + text="@bar this is the worst idea ever.", labels=["complaint"] + ), + ] + ), + "torch_dtype": torch.bfloat16, + "device": "cpu", + } + causal_lm_train_kwargs.update(patch_kwargs) + + model_name = causal_lm_train_kwargs["base_model"]._model_name + module = caikit_nlp.modules.text_generation.PeftPromptTuning + with temp_config(training_data_limit={module.MODULE_ID: {model_name: 1}}): + with pytest.raises(ValueError): + module.train(**causal_lm_train_kwargs) + + +def test_train_with_data_validation_success(causal_lm_train_kwargs, set_cpu_device): + """Check if we are able to train successfully if training data is within limits""" + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable( + [ + ClassificationTrainRecord( + text="@foo what a cute dog!", labels=["no complaint"] + ), + ClassificationTrainRecord( + text="@bar this is the worst idea ever.", labels=["complaint"] + ), + ] + ), + "torch_dtype": torch.bfloat16, + "device": "cpu", + } + causal_lm_train_kwargs.update(patch_kwargs) + + model_name = causal_lm_train_kwargs["base_model"]._model_name + module = caikit_nlp.modules.text_generation.PeftPromptTuning + with temp_config(training_data_limit={module.MODULE_ID: {model_name: 2}}): + + model = module.train(**causal_lm_train_kwargs) + assert model + + +def test_train_with_non_existent_limit_success(causal_lm_train_kwargs, set_cpu_device): + """Check if we are able to train successfully if training data limit doesn't exist for particular model""" + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable( + [ + ClassificationTrainRecord( + text="@foo what a cute dog!", labels=["no complaint"] + ) + ] + ), + "torch_dtype": torch.bfloat16, + "device": "cpu", + } + causal_lm_train_kwargs.update(patch_kwargs) + + model_name = causal_lm_train_kwargs["base_model"]._model_name + module = caikit_nlp.modules.text_generation.PeftPromptTuning + with temp_config(training_data_limit={module.MODULE_ID: {"foo": 2}}): + + model = module.train(**causal_lm_train_kwargs) + assert model + + +def test_train_with_no_limit_for_module(causal_lm_train_kwargs, set_cpu_device): + """Check if we are able to train successfully if training data limit doesn't exist prompt tuning module""" + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable( + [ + ClassificationTrainRecord( + text="@foo what a cute dog!", labels=["no complaint"] + ) + ] + ), + "torch_dtype": torch.bfloat16, + "device": "cpu", + } + causal_lm_train_kwargs.update(patch_kwargs) + + model_name = causal_lm_train_kwargs["base_model"]._model_name + module = caikit_nlp.modules.text_generation.PeftPromptTuning + with temp_config(training_data_limit={}): + + model = module.train(**causal_lm_train_kwargs) + assert model + + +def test_train_module_level_data_validation_raises( + causal_lm_train_kwargs, set_cpu_device +): + """Check if train raises with module level default configuration + if training data is within limits and model config is not provided + """ + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable( + [ + ClassificationTrainRecord( + text="@foo what a cute dog!", labels=["no complaint"] + ), + ClassificationTrainRecord( + text="@bar this is the worst idea ever.", labels=["complaint"] + ), + ] + ), + "torch_dtype": torch.bfloat16, + "device": "cpu", + } + causal_lm_train_kwargs.update(patch_kwargs) + + module = caikit_nlp.modules.text_generation.PeftPromptTuning + with temp_config( + training_data_limit={module.MODULE_ID: {"__default__": 1, "foo": 2}} + ): + with pytest.raises(ValueError): + module.train(**causal_lm_train_kwargs) + + +def test_train_module_level_data_validation_success( + causal_lm_train_kwargs, set_cpu_device +): + """Check if we are able to train successfully with module level default configuration + if training data is within limits and model config present + """ + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable( + [ + ClassificationTrainRecord( + text="@foo what a cute dog!", labels=["no complaint"] + ), + ClassificationTrainRecord( + text="@bar this is the worst idea ever.", labels=["complaint"] + ), + ] + ), + "torch_dtype": torch.bfloat16, + "device": "cpu", + } + causal_lm_train_kwargs.update(patch_kwargs) + + model_name = causal_lm_train_kwargs["base_model"]._model_name + module = caikit_nlp.modules.text_generation.PeftPromptTuning + with temp_config( + training_data_limit={module.MODULE_ID: {"__default__": 1, model_name: 2}} + ): + + model = module.train(**causal_lm_train_kwargs) + assert model + + +def test_train_global_default_data_validation_raises( + causal_lm_train_kwargs, set_cpu_device +): + """Check if train raises with global default configuration + if training data is within limits and model config is not provided + """ + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable( + [ + ClassificationTrainRecord( + text="@foo what a cute dog!", labels=["no complaint"] + ), + ClassificationTrainRecord( + text="@bar this is the worst idea ever.", labels=["complaint"] + ), + ] + ), + "torch_dtype": torch.bfloat16, + "device": "cpu", + } + causal_lm_train_kwargs.update(patch_kwargs) + + module = caikit_nlp.modules.text_generation.PeftPromptTuning + with temp_config( + training_data_limit={"__default__": 1, module.MODULE_ID: {"foo": 2}} + ): + with pytest.raises(ValueError): + module.train(**causal_lm_train_kwargs) + + +def test_train_global_default_data_validation_success( + causal_lm_train_kwargs, set_cpu_device +): + """Check if we are able to train successfully with global default configuration + if training data is within limits and model config is present + """ + patch_kwargs = { + "num_epochs": 1, + "verbalizer": "Tweet text : {{input}} Label : ", + "train_stream": caikit.core.data_model.DataStream.from_iterable( + [ + ClassificationTrainRecord( + text="@foo what a cute dog!", labels=["no complaint"] + ), + ClassificationTrainRecord( + text="@bar this is the worst idea ever.", labels=["complaint"] + ), + ] + ), + "torch_dtype": torch.bfloat16, + "device": "cpu", + } + causal_lm_train_kwargs.update(patch_kwargs) + + model_name = causal_lm_train_kwargs["base_model"]._model_name + module = caikit_nlp.modules.text_generation.PeftPromptTuning + with temp_config( + training_data_limit={"__default__": 1, module.MODULE_ID: {model_name: 2}} + ): + + model = module.train(**causal_lm_train_kwargs) + assert model diff --git a/tests/resources/test_pretrained_model.py b/tests/resources/test_pretrained_model.py index d2f1f033c..056a9a8ef 100644 --- a/tests/resources/test_pretrained_model.py +++ b/tests/resources/test_pretrained_model.py @@ -10,6 +10,7 @@ # Third Party from datasets import IterableDataset as TransformersIterableDataset +from torch.utils.data import DataLoader import pytest import torch import transformers @@ -83,6 +84,17 @@ def test_boostrap_causal_lm_download_enabled(mock_tok_from_pretrained, temp_cach assert kwargs["local_files_only"] == False +def test_boostrap_model_path(models_cache_dir): + """Ensure that we can bootstrap works with loading a model from local directory""" + # If we have an empty cachedir & do allow downloads, we should be able to init happily + base_model = HFAutoCausalLM.bootstrap( + model_name=CAUSAL_LM_MODEL, + ) + assert isinstance(base_model, HFAutoCausalLM) + assert base_model.MODEL_TYPE is transformers.AutoModelForCausalLM + assert base_model.TASK_TYPE == "CAUSAL_LM" + + ### Tests for tokenization behaviors SAMPLE_TRAINING_DATA = caikit.core.data_model.DataStream.from_iterable( [ @@ -91,7 +103,8 @@ def test_boostrap_causal_lm_download_enabled(mock_tok_from_pretrained, temp_cach ] ) - +# Causal LM tokenization strategies +### 1. Tests for Causal LM tokenization chunking def test_causal_lm_tokenize_func_contains_wrapped_stream(models_cache_dir): """Ensure the Causal LM tokenize func produces a wrapped stream that can be flattened.""" causal_lm = HFAutoCausalLM.bootstrap( @@ -102,6 +115,7 @@ def test_causal_lm_tokenize_func_contains_wrapped_stream(models_cache_dir): max_source_length=100, max_target_length=100, verbalizer="{{input}}", + use_seq2seq_approach=False, ) map_stream = SAMPLE_TRAINING_DATA.map(tok_func) # Since tok_func for causal lm creates a datastream, we should get a stream @@ -116,13 +130,21 @@ def test_causal_lm_tokenize_func_contains_wrapped_stream(models_cache_dir): ) -def test_causal_lm_tok_output_correctness(models_cache_dir): - """Validate the correctness of the attention mask for the language modeling objective.""" +# Key cases here are: +# 1 - simplest and minimal case +# 3 - because the concat sequence is length 17, so we have a remainder +# 100 - which is much larger than the concatenated seq and should yield one chunk +@pytest.mark.parametrize( + "chunk_size,drop_remainder", + [(1, True), (1, False), (3, True), (3, False), (100, True), (100, False)], +) +def test_causal_lm_tok_output_correctness(models_cache_dir, chunk_size, drop_remainder): + """Validate the tokenized results for the chunked language modeling objective.""" causal_lm = HFAutoCausalLM.bootstrap( model_name=CAUSAL_LM_MODEL, tokenizer_name=CAUSAL_LM_MODEL ) sample = GenerationTrainRecord( - input="This len does not matter", output="but this one does!" + input="Hello world", output="How are you doing today?!" ) (tok_func, _) = causal_lm.build_task_tokenize_closure( tokenizer=causal_lm.tokenizer, @@ -130,35 +152,138 @@ def test_causal_lm_tok_output_correctness(models_cache_dir): max_target_length=100, verbalizer="{{input}}", task_ids=0, + use_seq2seq_approach=False, + chunk_size=chunk_size, + drop_remainder=drop_remainder, ) input_tok = causal_lm.tokenizer.encode(sample.input) output_tok = causal_lm.tokenizer.encode(sample.output) + concat_tok = input_tok + output_tok tok_stream = tok_func(sample) # Ensure we get one token per output in our stream assert isinstance(tok_stream, caikit.core.data_model.DataStream) - assert len(tok_stream) == len(output_tok) - for idx, tok_sample in enumerate(tok_stream): - # We expect by default, everything is in order, and each attention mask grows the tokens - # we attend to in the target by one, until we are paying attention to the whole sequence. - expected_target_mask = torch.tensor( - ([1] * (idx + 1)) + [0] * (len(output_tok) - idx - 1) - ) - actual_target_mask = torch.tensor( - tok_sample["attention_mask"][-len(output_tok) :] - ) - assert bool(torch.all(expected_target_mask == actual_target_mask)) - # Check the source mask; we should always attend to the whole source sequence - actual_source_mask = torch.tensor( - tok_sample["attention_mask"][: len(input_tok)] + # Figure out how many chunks we should have, including if we have a remainder + has_remainder = False + if len(concat_tok) > chunk_size: + num_expected_chunks = len(concat_tok) // chunk_size + # Should only care about the remainder if we are not dropping it + if num_expected_chunks * chunk_size != len(concat_tok) and not drop_remainder: + has_remainder = True + else: + num_expected_chunks = 1 + chunk_size = len(concat_tok) + tok_list = list(tok_stream) + assert len(tok_list) == num_expected_chunks + has_remainder + # Check all full chunks. Note that we always attend to everything + for idx in range(num_expected_chunks): + assert len(tok_list[idx]["attention_mask"]) == chunk_size + assert len(tok_list[idx]["input_ids"]) == chunk_size + assert all(atn == 1 for atn in tok_list[idx]["attention_mask"]) + assert tok_list[idx]["task_ids"] == 0 + # Check the remainder; lists should be the same length, but less than the chunk size + if has_remainder: + remainder = tok_list[-1] + assert len(remainder["attention_mask"]) == len(remainder["input_ids"]) + assert len(remainder["input_ids"]) < chunk_size + assert all(atn == 1 for atn in remainder["attention_mask"]) + + +def test_causal_lm_batch_tokenization(models_cache_dir): + """Ensure that we can batch process causal lm inputs correctly.""" + causal_lm = HFAutoCausalLM.bootstrap( + model_name=CAUSAL_LM_MODEL, tokenizer_name=CAUSAL_LM_MODEL + ) + train_stream = DataStream.from_iterable( + [ + GenerationTrainRecord(input="hello there", output="world"), + GenerationTrainRecord(input="how", output="today"), + ] + ) + fn_kwargs = { + "tokenizer": causal_lm.tokenizer, + "max_source_length": 10, + "max_target_length": 10, + "use_seq2seq_approach": False, + } + # Create an iterable dataset by batching... + def get(train_stream): + for data in train_stream: + yield {"input": data.input, "output": data.output} + + dataset = TransformersIterableDataset.from_generator( + get, gen_kwargs={"train_stream": train_stream} + ) + batched_dataset = dataset.map( + causal_lm.tokenize_function, + fn_kwargs=fn_kwargs, + batched=True, + remove_columns=["input", "output"], + ) + + # Do the same thing with no batching via tokenize closure + unwrapping + tok_func = causal_lm.build_task_tokenize_closure(**fn_kwargs)[0] + mapped_indiv_stream = train_stream.map(tok_func).flatten() + for indiv_res, batched_res in zip(mapped_indiv_stream, batched_dataset): + # All keys should match (input ids, attention mask) + assert indiv_res.keys() == batched_res.keys() + # And all of their values should be the same + for k in indiv_res: + assert indiv_res[k] == batched_res[k] + + +### 2. Tests for causal LM framed as a seq2seq problem +# NOTE: For these tests, we should be careful to always test left and right padding +@pytest.mark.parametrize( + "padding_side", + ["left", "right"], +) +def test_causal_lm_as_a_sequence_problem_no_truncation(models_cache_dir, padding_side): + causal_lm = HFAutoCausalLM.bootstrap( + model_name=CAUSAL_LM_MODEL, tokenizer_name=CAUSAL_LM_MODEL + ) + sample = GenerationTrainRecord( + input="Hello world", output="How are you doing today?!" + ) + max_lengths = 20 + # First, build the output we expect for left / right respectively... + input_tok = causal_lm.tokenizer.encode(sample.input) + output_tok = causal_lm.tokenizer.encode(sample.output) + [ + causal_lm.tokenizer.eos_token_id + ] + concat_res = input_tok + output_tok + masked_res = ([-100] * len(input_tok)) + output_tok + + # This must true because otherwise no padding was needed, e.g., truncation + assert len(input_tok) < max_lengths + assert len(output_tok) < (max_lengths + 1) + pads_needed = (1 + 2 * max_lengths) - len(concat_res) + if causal_lm.tokenizer.padding_side.lower() == "left": + expected_input_ids = torch.tensor( + [causal_lm.tokenizer.pad_token_id] * pads_needed + concat_res ) - assert bool(torch.all(torch.tensor([1] * len(input_tok)) == actual_source_mask)) - # Also, the number of tokens we attend to should be the sum of toks in input/output - assert (len(actual_target_mask) + len(actual_source_mask)) == len( - tok_sample["attention_mask"] + expected_attn_mask = torch.tensor([0] * pads_needed + [1] * len(concat_res)) + expected_labels = torch.tensor([-100] * pads_needed + masked_res) + else: + expected_input_ids = torch.tensor( + concat_res + [causal_lm.tokenizer.pad_token_id] * pads_needed ) - # Ensure we support MPT - assert hasattr(tok_sample, "task_ids") - assert tok_sample["task_ids"] == 0 + expected_attn_mask = torch.tensor([1] * len(concat_res) + [0] * pads_needed) + expected_labels = torch.tensor(masked_res + [-100] * pads_needed) + + # Now build the analogous tokenizer closure and compare the tensors + (tok_func, _) = causal_lm.build_task_tokenize_closure( + tokenizer=causal_lm.tokenizer, + max_source_length=max_lengths, + max_target_length=max_lengths, + verbalizer="{{input}}", + task_ids=0, + use_seq2seq_approach=True, + ) + tok_res = tok_func(sample) + assert tok_res["task_ids"] == 0 + assert torch.all(tok_res["input_ids"] == expected_input_ids) + assert torch.all(tok_res["attention_mask"] == expected_attn_mask) + assert torch.all(tok_res["labels"] == expected_labels) ### Tests for Seq2Seq tokenization @@ -211,43 +336,80 @@ def test_seq2seq_tok_output_correctness(models_cache_dir): assert tok_sample["task_ids"] == 0 -def test_causal_lm_batch_tokenization(models_cache_dir): - """Ensure that we can batch process causal lm inputs correctly.""" - causal_lm = HFAutoCausalLM.bootstrap( - model_name=CAUSAL_LM_MODEL, tokenizer_name=CAUSAL_LM_MODEL - ) +### Tests for collator compatability +# These tests should validate that we can use our tokenization function to +# build torch loaders around datasets using different collators. +# TODO: Expand to cover transformer datasets, i.e., what is produced by +# text gen preprocessing functions. For now, they only check the minimal +# case with the default data collator. +@pytest.mark.parametrize( + "collator_fn", + [transformers.default_data_collator], +) +def test_loader_can_batch_list_of_seq2seq_outputs(collator_fn): + # Build the dataset train_stream = DataStream.from_iterable( [ - GenerationTrainRecord(input="hello there", output="world"), - GenerationTrainRecord(input="how", output="today"), + GenerationTrainRecord(input="hello world", output="how are you today?"), + GenerationTrainRecord(input="goodbye", output="world"), + GenerationTrainRecord(input="good morning", output="have a good day"), + GenerationTrainRecord(input="good night", output="have nice dreams"), ] ) - fn_kwargs = { - "tokenizer": causal_lm.tokenizer, - "max_source_length": 10, - "max_target_length": 10, - } - # Create an iterable dataset by batching... - def get(train_stream): - for data in train_stream: - yield {"input": data.input, "output": data.output} - - dataset = TransformersIterableDataset.from_generator( - get, gen_kwargs={"train_stream": train_stream} + seq2seq = HFAutoSeq2SeqLM.bootstrap( + model_name=SEQ2SEQ_LM_MODEL, tokenizer_name=SEQ2SEQ_LM_MODEL ) - batched_dataset = dataset.map( - causal_lm.tokenize_function, - fn_kwargs=fn_kwargs, - batched=True, - remove_columns=["input", "output"], + (tok_func, _) = seq2seq.build_task_tokenize_closure( + tokenizer=seq2seq.tokenizer, + max_source_length=20, + max_target_length=20, + verbalizer="{{input}}", + task_ids=0, + ) + tok_results = [tok_func(x) for x in list(train_stream)] + dl = DataLoader( + tok_results, + shuffle=False, + batch_size=2, + collate_fn=collator_fn, ) + # Loader should create 2 batches + loader_list = list(dl) + assert len(loader_list) == 2 - # Do the same thing with no batching via tokenize closure + unwrapping - tok_func = causal_lm.build_task_tokenize_closure(**fn_kwargs)[0] - mapped_indiv_stream = train_stream.map(tok_func).flatten() - for indiv_res, batched_res in zip(mapped_indiv_stream, batched_dataset): - # All keys should match (input ids, attention mask) - assert indiv_res.keys() == batched_res.keys() - # And all of their values should be the same - for k in indiv_res: - assert indiv_res[k] == batched_res[k] + +@pytest.mark.parametrize( + "collator_fn", + [transformers.default_data_collator], +) +def test_loader_can_batch_list_of_causal_lm_outputs(collator_fn): + # Build the dataset + train_stream = DataStream.from_iterable( + [ + GenerationTrainRecord(input="hello world", output="how are you today?"), + GenerationTrainRecord(input="goodbye", output="world"), + GenerationTrainRecord(input="good morning", output="have a good day"), + GenerationTrainRecord(input="good night", output="have nice dreams"), + ] + ) + causal_lm = HFAutoCausalLM.bootstrap( + model_name=CAUSAL_LM_MODEL, tokenizer_name=CAUSAL_LM_MODEL + ) + (tok_func, _) = causal_lm.build_task_tokenize_closure( + tokenizer=causal_lm.tokenizer, + max_source_length=20, + max_target_length=20, + verbalizer="{{input}}", + task_ids=0, + use_seq2seq_approach=True, + ) + tok_results = [tok_func(x) for x in list(train_stream)] + dl = DataLoader( + tok_results, + shuffle=False, + batch_size=2, + collate_fn=collator_fn, + ) + # Loader should create 2 batches + loader_list = list(dl) + assert len(loader_list) == 2 diff --git a/tests/toolkit/text_generation/test_model_run_utils.py b/tests/toolkit/text_generation/test_model_run_utils.py new file mode 100644 index 000000000..8835089b9 --- /dev/null +++ b/tests/toolkit/text_generation/test_model_run_utils.py @@ -0,0 +1,45 @@ +# Third Party +import pytest + +# First Party +from caikit.core.data_model.producer import ProducerId +from caikit.interfaces.nlp.data_model import GeneratedTextResult + +# Local +from caikit_nlp.toolkit.text_generation.model_run_utils import generate_text_func +from tests.fixtures import ( + causal_lm_dummy_model, + causal_lm_train_kwargs, + seq2seq_lm_dummy_model, + seq2seq_lm_train_kwargs, +) + + +@pytest.mark.parametrize( + "model_fixture", ["seq2seq_lm_dummy_model", "causal_lm_dummy_model"] +) +@pytest.mark.parametrize( + "serialization_method,expected_type", + [ + ("to_dict", dict), + ("to_json", str), + ("to_proto", GeneratedTextResult._proto_class), + ], +) +def test_generate_text_func_serialization_json( + request, + model_fixture, + serialization_method, + expected_type, +): + model = request.getfixturevalue(model_fixture) + generated_text = generate_text_func( + model=model.model, + tokenizer=model.tokenizer, + producer_id=ProducerId("TextGeneration", "0.1.0"), + eos_token="<\n>", + text="What is the boiling point of liquid Nitrogen?", + ) + + serialized = getattr(generated_text, serialization_method)() + assert isinstance(serialized, expected_type) diff --git a/tox.ini b/tox.ini index c117700b8..ed361a283 100644 --- a/tox.ini +++ b/tox.ini @@ -15,7 +15,7 @@ passenv = LOG_FORMATTER LOG_THREAD_ID LOG_CHANNEL_WIDTH -commands = pytest --cov=caikit_nlp --cov-report=term --cov-report=html {posargs:tests} +commands = pytest --durations=42 --cov=caikit_nlp --cov-report=term --cov-report=html {posargs:tests} ; Unclear: We probably want to test wheel packaging ; But! tox will fail when this is set and _any_ interpreter is missing @@ -34,21 +34,17 @@ description = lint with pylint deps = pylint>=2.16.2,<3.0 commands = pylint caikit_nlp -[testenv:publish] -description = publish wheel to pypi -deps = flit==3.8 -passenv = - FLIT_PASSWORD -setenv = - FLIT_USERNAME = __token__ -commands = flit publish -skip_install = True - [testenv:build] description = build wheel -deps = flit==3.8 -passenv = - FLIT_PASSWORD -setenv = - FLIT_USERNAME = __token__ -commands = flit build +deps = + build + setuptools +commands = python -m build +skip_install = True + +[testenv:twinecheck] +description = check wheel +deps = + twine +commands = twine check dist/* +skip_install = True