diff --git a/README.md b/README.md index 0c30f71..bcaca9f 100644 --- a/README.md +++ b/README.md @@ -37,79 +37,28 @@ model.print_trainable_parameters() # output: trainable params: 2359296 || all params: 1231940608 || trainable%: 0.19151053100118282 ``` +## Use Cases + +### Get comparable performance to full finetuning by adapting LLMs to downstream tasks using less computational resources + +### Parameter Efficient Tuning of Diffusion Models + +### Parameter Efficient Tuning of LLMs for RLHF components [ToDo] + +### Save compute and storage even for medium and small models + ## PET + 🤗 Accelerate -PET models work with 🤗 Accelerate out of the box. -For scaling to large models, you can leverage 🤗 Accelerate's PyTorch FSDP integration as shown below. -PyTorch FSDP shards parameters, gradients and optimizer states across data parallel workers which enables -large language models to fit on available hardware. -It also supports CPU offloading to further enable distributed training at scale. +PET models work with 🤗 Accelerate out of the box. Use 🤗 Accelerate for Distributed training on various hardware such as GPUs, Apple Silicon devices etc during training. +Use 🤗 Accelerate for inferencing on consumer hardware with small resources. -```python -from pet.utils.other import fsdp_auto_wrap_policy +### Example of PET model distributed training using 🤗 Accelerate -... - -if os.environ.get("ACCELERATE_USE_FSDP", None) is not None: - accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) - -model = accelerator.prepare(model) -``` - -Example of parameter efficient tuning with `mt0-xxl` base model using 🤗 Accelerate is provided in `~examples/pet_lora_seq2seq_accelerate_fsdp.py`. -1. First run `accelerate config --config_file fsdp_config.yaml` and answer the questionaire. -Below are the contents of the config file. -``` -command_file: null -commands: null -compute_environment: LOCAL_MACHINE -deepspeed_config: {} -distributed_type: FSDP -downcast_bf16: 'no' -dynamo_backend: 'NO' -fsdp_config: - fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP - fsdp_backward_prefetch_policy: BACKWARD_PRE - fsdp_offload_params: true - fsdp_sharding_strategy: 1 - fsdp_state_dict_type: FULL_STATE_DICT - fsdp_transformer_layer_cls_to_wrap: T5Block -gpu_ids: null -machine_rank: 0 -main_process_ip: null -main_process_port: null -main_training_function: main -megatron_lm_config: {} -mixed_precision: 'no' -num_machines: 1 -num_processes: 2 -rdzv_backend: static -same_network: true -tpu_name: null -tpu_zone: null -use_cpu: false -``` -2. run the below command to launch example script -``` -accelerate launch --config_file fsdp_config.yaml examples/pet_lora_seq2seq_accelerate_fsdp.py -``` +### Example of PET model inference using 🤗 Accelerate ## Models support matrix -### Sequence Classification -| Model | LoRA | Prefix Tuning | P-Tuning | Prompt Tuning | -| --------- | ---- | ---- | ---- | ---- | -| BERT | ✅ | ✅ | ✅ | ✅ | -| RoBERTa | ✅ | ✅ | ✅ | ✅ | -| GPT-2 | ✅ | ✅ | ✅ | ✅ | -| Bloom | ✅ | ✅ | ✅ | ✅ | -| OPT | ✅ | ✅ | ✅ | ✅ | -| GPT-Neo | ✅ | ✅ | ✅ | ✅ | -| GPT-J | ✅ | ✅ | ✅ | ✅ | -| Deberta | ✅ | | | | -| Deberta-v2 | ✅ | | | | - ### Causal Language Modeling | Model | LoRA | Prefix Tuning | P-Tuning | Prompt Tuning | | --------- | ---- | ---- | ---- | ---- | @@ -125,11 +74,135 @@ accelerate launch --config_file fsdp_config.yaml examples/pet_lora_seq2seq_accel | T5 | ✅ | ✅ | ✅ | ✅ | | BART | ✅ | ✅ | ✅ | ✅ | +### Sequence Classification +| Model | LoRA | Prefix Tuning | P-Tuning | Prompt Tuning | +| --------- | ---- | ---- | ---- | ---- | +| BERT | ✅ | ✅ | ✅ | ✅ | +| RoBERTa | ✅ | ✅ | ✅ | ✅ | +| GPT-2 | ✅ | ✅ | ✅ | ✅ | +| Bloom | ✅ | ✅ | ✅ | ✅ | +| OPT | ✅ | ✅ | ✅ | ✅ | +| GPT-Neo | ✅ | ✅ | ✅ | ✅ | +| GPT-J | ✅ | ✅ | ✅ | ✅ | +| Deberta | ✅ | | | | +| Deberta-v2 | ✅ | | | | + +### Token Classification +| Model | LoRA | Prefix Tuning | P-Tuning | Prompt Tuning | +| --------- | ---- | ---- | ---- | ---- | +| BERT | ✅ | ✅ | ✅ | ✅ | +| RoBERTa | ✅ | ✅ | ✅ | ✅ | +| GPT-2 | ✅ | ✅ | ✅ | ✅ | +| Bloom | ✅ | ✅ | ✅ | ✅ | +| OPT | ✅ | ✅ | ✅ | ✅ | +| GPT-Neo | ✅ | ✅ | ✅ | ✅ | +| GPT-J | ✅ | ✅ | ✅ | ✅ | +| Deberta | ✅ | | | | +| Deberta-v2 | ✅ | | | | + ## Caveats: -1. Doesn't work currently with DeeSpeed ZeRO Stage-3. Extending support with DeeSpeed ZeRO Stage-3 is in backlog. -2. When using `P_TUNING` or `PROMPT_TUNING` with `SEQ_2_SEQ` task, remember to remove the `num_virtual_token` virtual prompt predictions from the left side of the model outputs during evaluations. +1. Needs a workaround when using DeeSpeed ZeRO Stage-3 for training. However, it doesn't lead to any GPU memory savings. Plase refer [[REQUEST] efficiently deal with frozen weights during training](https://github.com/microsoft/DeepSpeed/issues/2615) issue on DeepSpeed repository. Example is provided in `~examples/pet_lora_seq2seq_accelerate_ds_zero3_offload.py`. + a. First run `accelerate config --config_file ds_zero3_config.yaml` and answer the questionaire. + Below are the contents of the config file. + ``` + command_file: null + commands: null + compute_environment: LOCAL_MACHINE + deepspeed_config: {} + distributed_type: FSDP + downcast_bf16: 'no' + dynamo_backend: 'NO' + fsdp_config: + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_backward_prefetch_policy: BACKWARD_PRE + fsdp_offload_params: true + fsdp_sharding_strategy: 1 + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_transformer_layer_cls_to_wrap: T5Block + gpu_ids: null + machine_rank: 0 + main_process_ip: null + main_process_port: null + main_training_function: main + megatron_lm_config: {} + mixed_precision: 'no' + num_machines: 1 + num_processes: 2 + rdzv_backend: static + same_network: true + tpu_name: null + tpu_zone: null + use_cpu: false + ``` + b. run the below command to launch example script + ``` + accelerate launch --config_file ds_zero3_config.yaml examples/pet_lora_seq2seq_accelerate_ds_zero3_offload.py + ``` + +2. Below is an example of using PyTorch FSDP for training. However, it doesn't lead to +any GPU memory savings. Please refer issue [[FSDP] FSDP with CPU offload consumes 1.65X more GPU memory when training models with most of the params frozen](https://github.com/pytorch/pytorch/issues/91165). + + ```python + from pet.utils.other import fsdp_auto_wrap_policy + + ... + + if os.environ.get("ACCELERATE_USE_FSDP", None) is not None: + accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) + + model = accelerator.prepare(model) + ``` + + Example of parameter efficient tuning with `mt0-xxl` base model using 🤗 Accelerate is provided in `~examples/pet_lora_seq2seq_accelerate_fsdp.py`. + a. First run `accelerate config --config_file fsdp_config.yaml` and answer the questionaire. + Below are the contents of the config file. + ``` + command_file: null + commands: null + compute_environment: LOCAL_MACHINE + deepspeed_config: {} + distributed_type: FSDP + downcast_bf16: 'no' + dynamo_backend: 'NO' + fsdp_config: + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_backward_prefetch_policy: BACKWARD_PRE + fsdp_offload_params: true + fsdp_sharding_strategy: 1 + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_transformer_layer_cls_to_wrap: T5Block + gpu_ids: null + machine_rank: 0 + main_process_ip: null + main_process_port: null + main_training_function: main + megatron_lm_config: {} + mixed_precision: 'no' + num_machines: 1 + num_processes: 2 + rdzv_backend: static + same_network: true + tpu_name: null + tpu_zone: null + use_cpu: false + ``` + b. run the below command to launch example script + ``` + accelerate launch --config_file fsdp_config.yaml examples/pet_lora_seq2seq_accelerate_fsdp.py + ``` + +3. When using `P_TUNING` or `PROMPT_TUNING` with `SEQ_2_SEQ` task, remember to remove the `num_virtual_token` virtual prompt predictions from the left side of the model outputs during evaluations. + +4. `P_TUNING` or `PROMPT_TUNING` doesn't support `generate` functionality of transformers bcause `generate` strictly requires `input_ids`/`decoder_input_ids` but +`P_TUNING`/`PROMPT_TUNING` appends soft prompt embeddings to `input_embeds` to create +new `input_embeds` to be given to the model. Therefore, `generate` doesn't support this yet. + +## Backlog: +1. Explore and possibly integrate `(IA)^3` and `UniPELT` +2. Add tests +3. Add more use cases and examples ## Citing 🤗 PET diff --git a/examples/pet_lora_seq2seq_accelerate_fsdp.py b/examples/pet_lora_seq2seq_accelerate_fsdp.py index a534f2c..9cc3a9e 100644 --- a/examples/pet_lora_seq2seq_accelerate_fsdp.py +++ b/examples/pet_lora_seq2seq_accelerate_fsdp.py @@ -78,7 +78,7 @@ def main(): num_training_steps=(len(train_dataloader) * num_epochs), ) - if accelerator.state.fsdp_plugin is not None: + if getattr(accelerator.state, "fsdp_plugin", None) is not None: accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) model, train_dataloader, eval_dataloader, optimizer, lr_scheduler = accelerator.prepare( diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py index b1034cf..7dd19f7 100644 --- a/src/pet/utils/save_and_load.py +++ b/src/pet/utils/save_and_load.py @@ -1,6 +1,7 @@ from loralib import lora_state_dict from .config import PETType +from ..mapping import get_pet_model def get_pet_model_state_dict(model, state_dict=None): @@ -44,3 +45,32 @@ def set_pet_model_state_dict(model, pet_model_state_dict): {"weight": pet_model_state_dict["prompt_embeddings"]}, strict=True ) return model + + +def pet_model_load_and_dispatch(model, pet_model_state_dict, pet_config, max_memory=None): + """ + Load the PET model state dict and dispatch the model to the correct device. + + Args: + model (:obj:`PETModel`): The Pre-trained base model which has already been sharded and dispatched + using `accelerate` functionalities. + pet_model_state_dict (:obj:`dict`): The state dict of the PET model. + max_memory (`Dict`, *optional*): + A dictionary device identifier to maximum memory. Will default to the maximum memory available for each GPU + and the available CPU RAM if unset. + """ + from accelerate import infer_auto_device_map, dispatch_model + from accelerate.hooks import remove_hook_from_submodules, AlignDevicesHook, add_hook_to_module + + remove_hook_from_submodules(model) + model = get_pet_model(model, pet_config) + model.print_trainable_parameters() + set_pet_model_state_dict(model, pet_model_state_dict) + device_map = infer_auto_device_map(model, max_memory=max_memory, no_split_module_classes=model._no_split_modules) + model = dispatch_model(model, device_map=device_map) + hook = AlignDevicesHook(io_same_device=True) + if model.pet_config.pet_type == PETType.LORA: + add_hook_to_module(model.base_model.model, hook) + else: + add_hook_to_module(model.base_model, hook) + return model