Single GPU Billion-scale Model Training via Parameter-Efficient Finetuning#

Open In Colab Open In SageMaker Studio Lab

As pointed out by a recent paper from Stanford Institute for Human-Centered Artificial Intelligence, AI is undergoing a paradigm shift with the rise of “foundation models”, i.e., giant models that are trained on a diverse collection of datasets generally in a self-supervised way. These foundation models, which are the key of AutoMM, can be easily adapted to down-stream applications. However, as the size of these foundation models grows, finetuning these models becomes increasingly difficult. Following is a figure from the Microsoft research blog that demonstrates the trend:

Scaling of foundation models

The goal of AutoMM is to help anyone solve machine learning problems via open source foundation models, including these giant models. To finetune these large-scale models, we adopt the recently popularized parameter-efficient finetuning technique. The idea is to either finetune a small subset of the weights in the foundation model (e.g., BitFit), or adding a tiny tunable structure on top of the fixed backbone (e.g., Prompt Tuning, LoRA, Adapter, MAM Adapter, IA^3). These techniques can effectively reduce the peak memory usage and model training time, while maintaining the performance.

In this tutorial, we introduce how to apply parameter-efficient finetuning in MultiModalPredictor. We first introduce how to adopt the "ia3_bias" algorithm for parameter-efficient finetuning. Afterwards, we show how you can simply combine "ia3_bias" and gradient checkpointing to finetune the XL-variant of Google’s FLAN-T5 via a single NVIDIA T4 GPU.

Prepare Dataset#

The Cross-Lingual Amazon Product Review Sentiment dataset contains Amazon product reviews in four languages. Here, we load the English and German fold of the dataset. In the label column, 0 means negative sentiment and 1 means positive sentiment. For the purpose of demonstration, we downsampled the training data to 1000 samples. We will train the model on the English dataset and directly evaluate its performance on the German and Japanese test set.

!wget --quiet https://automl-mm-bench.s3.amazonaws.com/multilingual-datasets/amazon_review_sentiment_cross_lingual.zip -O amazon_review_sentiment_cross_lingual.zip
!unzip -q -o amazon_review_sentiment_cross_lingual.zip -d .
import os
import shutil
os.environ["TRANSFORMERS_CACHE"] = "cache"

def clear_cache():
    if os.path.exists("cache"):
        shutil.rmtree("cache")

clear_cache()
import pandas as pd
import warnings
warnings.filterwarnings("ignore")

train_en_df = pd.read_csv("amazon_review_sentiment_cross_lingual/en_train.tsv",
                          sep="\t",
                          header=None,
                          names=["label", "text"]) \
                .sample(1000, random_state=123).reset_index(drop=True)

test_en_df = pd.read_csv("amazon_review_sentiment_cross_lingual/en_test.tsv",
                          sep="\t",
                          header=None,
                          names=["label", "text"]) \
               .sample(200, random_state=123).reset_index(drop=True)
test_de_df = pd.read_csv("amazon_review_sentiment_cross_lingual/de_test.tsv",
                          sep="\t", header=None, names=["label", "text"]) \
               .sample(200, random_state=123).reset_index(drop=True)

test_jp_df = pd.read_csv('amazon_review_sentiment_cross_lingual/jp_test.tsv',
                          sep='\t', header=None, names=['label', 'text']) \
               .sample(200, random_state=123).reset_index(drop=True)
train_en_df.head(5)
label text
0 0 This is a film that literally sees little wron...
1 0 This music is pretty intelligent, but not very...
2 0 One of the best pieces of rock ever recorded, ...
3 0 Reading the posted reviews here, is like revis...
4 1 I've just finished page 341, the last page. It...
test_jp_df.head(5)
label text
0 1 原作はビクトル・ユーゴの長編小説だが、私が子供の頃読んだのは短縮版の「ああ無情」。それでもこ...
1 1 ほかの作品のレビューにみんな書いているのに、何故この作品について書いている人が一人しかいない...
2 0 一番の問題点は青島が出ていない事でしょう。 TV番組では『芸人が出ていればバラエティだから...
3 0 昔、 りんたろう監督によるアニメ「カムイの剣」があった。 「カムイの剣」…を観た人なら本作...
4 1 以前のアルバムを聴いていないのでなんとも言えないが、クラシックなメタルを聞いてきた耳には、と...

Finetuning Multilingual Model with IA3 + BitFit#

In AutoMM, to enable efficient finetuning, just specify the optimization.efficient_finetune to be "ia3_bias".

from autogluon.multimodal import MultiModalPredictor
import uuid

model_path = f"./tmp/{uuid.uuid4().hex}-multilingual_ia3"
predictor = MultiModalPredictor(label="label",
                                path=model_path)
predictor.fit(train_en_df,
              presets="multilingual",
              hyperparameters={
                  "optimization.efficient_finetune": "ia3_bias",
                  "optimization.lr_decay": 0.9,
                  "optimization.learning_rate": 3e-03,
                  "optimization.end_lr": 3e-03,
                  "optimization.max_epochs": 2,
                  "optimization.warmup_steps": 0,
                  "env.batch_size": 32,
              })
Global seed set to 123
AutoMM starts to create your model. ✨

- Model will be saved to "/home/ci/autogluon/docs/tutorials/multimodal/advanced_topics/tmp/1d9db017293348d3980aadb517c87ea3-multilingual_ia3".

- Validation metric is "roc_auc".

- To track the learning progress, you can open a terminal and launch Tensorboard:
    ```shell
    # Assume you have installed tensorboard
    tensorboard --logdir /home/ci/autogluon/docs/tutorials/multimodal/advanced_topics/tmp/1d9db017293348d3980aadb517c87ea3-multilingual_ia3
    ```

Enjoy your coffee, and let AutoMM do the job ☕☕☕ Learn more at https://auto.gluon.ai
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
 in <module>:7                                                                                    
                                                                                                  
    4 model_path = f"./tmp/{uuid.uuid4().hex}-multilingual_ia3"                                   
    5 predictor = MultiModalPredictor(label="label",                                              
    6 │   │   │   │   │   │   │   │   path=model_path)                                            
  7 predictor.fit(train_en_df,                                                                  
    8 │   │   │     presets="multilingual",                                                       
    9 │   │   │     hyperparameters={                                                             
   10 │   │   │   │     "optimization.efficient_finetune": "ia3_bias",                            
                                                                                                  
 /home/ci/autogluon/multimodal/src/autogluon/multimodal/predictor.py:834 in fit                   
                                                                                                  
    831 │   │   │   )                                                                             
    832 │   │   │   return predictor                                                              
    833 │   │                                                                                     
  834 │   │   self._fit(**_fit_args)                                                            
    835 │   │   training_end = time.time()                                                        
    836 │   │   self._total_train_time = training_end - training_start                            
    837                                                                                           
                                                                                                  
 /home/ci/autogluon/multimodal/src/autogluon/multimodal/predictor.py:1037 in _fit                 
                                                                                                  
   1034 │   │   │   self._output_shape = len(df_preprocessor.label_generator.unique_entity_group  
   1035 │   │                                                                                     
   1036 │   │   if self._model is None:                                                           
 1037 │   │   │   model = create_fusion_model(                                                  
   1038 │   │   │   │   config=config,                                                            
   1039 │   │   │   │   num_classes=self._output_shape,                                           
   1040 │   │   │   │   classes=self._classes,                                                    
                                                                                                  
 /home/ci/autogluon/multimodal/src/autogluon/multimodal/utils/model.py:432 in create_fusion_model 
                                                                                                  
   429 │                                                                                          
   430 │   for model_name in names:                                                               
   431 │   │   model_config = getattr(config.model, model_name)                                   
 432 │   │   model = create_model(                                                              
   433 │   │   │   model_name=model_name,                                                         
   434 │   │   │   model_config=model_config,                                                     
   435 │   │   │   num_classes=num_classes,                                                       
                                                                                                  
 /home/ci/autogluon/multimodal/src/autogluon/multimodal/utils/model.py:206 in create_model        
                                                                                                  
   203 │   │   │   pretrained=pretrained,                                                         
   204 │   │   )                                                                                  
   205 │   elif model_name.lower().startswith(HF_TEXT):                                           
 206 │   │   model = HFAutoModelForTextPrediction(                                              
   207 │   │   │   prefix=model_name,                                                             
   208 │   │   │   checkpoint_name=model_config.checkpoint_name,                                  
   209 │   │   │   num_classes=num_classes,                                                       
                                                                                                  
 /home/ci/autogluon/multimodal/src/autogluon/multimodal/models/huggingface_text.py:81 in __init__ 
                                                                                                  
    78 │   │   self.checkpoint_name = checkpoint_name                                             
    79 │   │   self.num_classes = num_classes                                                     
    80 │   │                                                                                      
  81 │   │   self.config, self.model = get_hf_config_and_model(                                 
    82 │   │   │   checkpoint_name=checkpoint_name, pretrained=pretrained, low_cpu_mem_usage=lo   
    83 │   │   )                                                                                  
    84 │   │   self._hf_model_input_names = AutoTokenizer.from_pretrained(checkpoint_name).mode   
                                                                                                  
 /home/ci/autogluon/multimodal/src/autogluon/multimodal/models/utils.py:561 in                    
 get_hf_config_and_model                                                                          
                                                                                                  
   558 -------                                                                                
   559 A Huggingface config and model.                                                        
   560 """                                                                                    
 561 config = AutoConfig.from_pretrained(checkpoint_name)                                   
   562 │                                                                                          
   563 │   if pretrained:                                                                         
   564 │   │   model = AutoModel.from_pretrained(checkpoint_name, low_cpu_mem_usage=low_cpu_mem   
                                                                                                  
 /home/ci/opt/venv/lib/python3.8/site-packages/transformers/models/auto/configuration_auto.py:852 
 in from_pretrained                                                                               
                                                                                                  
   849 │   │   kwargs["_from_auto"] = True                                                        
   850 │   │   kwargs["name_or_path"] = pretrained_model_name_or_path                             
   851 │   │   trust_remote_code = kwargs.pop("trust_remote_code", False)                         
 852 │   │   config_dict, unused_kwargs = PretrainedConfig.get_config_dict(pretrained_model_n   
   853 │   │   if "auto_map" in config_dict and "AutoConfig" in config_dict["auto_map"]:          
   854 │   │   │   if not trust_remote_code:                                                      
   855 │   │   │   │   raise ValueError(                                                          
                                                                                                  
 /home/ci/opt/venv/lib/python3.8/site-packages/transformers/configuration_utils.py:565 in         
 get_config_dict                                                                                  
                                                                                                  
   562 │   │   """                                                                                
   563 │   │   original_kwargs = copy.deepcopy(kwargs)                                            
   564 │   │   # Get config dict associated with the base config file                             
 565 │   │   config_dict, kwargs = cls._get_config_dict(pretrained_model_name_or_path, **kwar   
   566 │   │   if "_commit_hash" in config_dict:                                                  
   567 │   │   │   original_kwargs["_commit_hash"] = config_dict["_commit_hash"]                  
   568                                                                                            
                                                                                                  
 /home/ci/opt/venv/lib/python3.8/site-packages/transformers/configuration_utils.py:650 in         
 _get_config_dict                                                                                 
                                                                                                  
   647 │   │                                                                                      
   648 │   │   try:                                                                               
   649 │   │   │   # Load config dict                                                             
 650 │   │   │   config_dict = cls._dict_from_json_file(resolved_config_file)                   
   651 │   │   │   config_dict["_commit_hash"] = commit_hash                                      
   652 │   │   except (json.JSONDecodeError, UnicodeDecodeError):                                 
   653 │   │   │   raise EnvironmentError(                                                        
                                                                                                  
 /home/ci/opt/venv/lib/python3.8/site-packages/transformers/configuration_utils.py:736 in         
 _dict_from_json_file                                                                             
                                                                                                  
   733 │                                                                                          
   734 │   @classmethod                                                                           
   735 │   def _dict_from_json_file(cls, json_file: Union[str, os.PathLike]):                     
 736 │   │   with open(json_file, "r", encoding="utf-8") as reader:                             
   737 │   │   │   text = reader.read()                                                           
   738 │   │   return json.loads(text)                                                            
   739                                                                                            
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
FileNotFoundError: [Errno 2] No such file or directory: 
'cache/models--microsoft--mdeberta-v3-base/snapshots/260983f7a3252791a76fcfbb0dcda00d060636c0/config.json'

The fraction of the tunable parameters is around 0.5% of all parameters. Actually, the model trained purely on English data can achieve good performance on the test sets, even on the German / Japanese test set. It obtained comparable results as full-finetuning as in AutoMM for Text - Multilingual Problems.

score_in_en = predictor.evaluate(test_en_df)
score_in_de = predictor.evaluate(test_de_df)
score_in_jp = predictor.evaluate(test_jp_df)
print('Score in the English Testset:', score_in_en)
print('Score in the German Testset:', score_in_de)
print('Score in the Japanese Testset:', score_in_jp)

Training FLAN-T5-XL on Single GPU#

By combining gradient checkpointing and parameter-efficient finetuning, it is feasible to finetune google/flan-t5-xl that has close to two billion parameterswith a single T4 GPU available in AWS G4 instances. To turn on gradient checkpointing, you just need to set "model.hf_text.gradient_checkpointing" to True. To accelerate the training, we downsample the number of training samples to be 200.

# Just for clean the space
clear_cache()
shutil.rmtree(model_path)
from autogluon.multimodal import MultiModalPredictor

train_en_df_downsample = train_en_df.sample(200, random_state=123)

new_model_path = f"./tmp/{uuid.uuid4().hex}-multilingual_ia3_gradient_checkpoint"
predictor = MultiModalPredictor(label="label",
                                path=new_model_path)
predictor.fit(train_en_df_downsample,
              presets="multilingual",
              hyperparameters={
                  "model.hf_text.checkpoint_name": "google/flan-t5-xl",
                  "model.hf_text.gradient_checkpointing": True,
                  "model.hf_text.low_cpu_mem_usage": True,
                  "optimization.efficient_finetune": "ia3_bias",
                  "optimization.lr_decay": 0.9,
                  "optimization.learning_rate": 3e-03,
                  "optimization.end_lr": 3e-03,
                  "optimization.max_epochs": 1,
                  "optimization.warmup_steps": 0,
                  "env.batch_size": 1,
                  "env.eval_batch_size_ratio": 1
              })

Global seed set to 123
Auto select gpus: [0]
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name              | Type                         | Params
-------------------------------------------------------------------
0 | model             | HFAutoModelForTextPrediction | 1.2 B 
1 | validation_metric | AUROC                        | 0     
2 | loss_func         | CrossEntropyLoss             | 0     
-------------------------------------------------------------------
203 K     Trainable params
1.2 B     Non-trainable params
1.2 B     Total params
4,894.913 Total estimated model params size (MB)
Epoch 0, global step 20: 'val_roc_auc' reached 0.88802 (best 0.88802), saving model to '/home/ubuntu/autogluon/docs/tutorials/multimodal/advanced_topics/multilingual_ia3_gradient_checkpoint/epoch=0-step=20.ckpt' as top 1
Epoch 0, global step 40: 'val_roc_auc' reached 0.94531 (best 0.94531), saving model to '/home/ubuntu/autogluon/docs/tutorials/multimodal/advanced_topics/multilingual_ia3_gradient_checkpoint/epoch=0-step=40.ckpt' as top 1
`Trainer.fit` stopped: `max_epochs=1` reached.





<autogluon.multimodal.predictor.MultiModalPredictor at 0x7fd58c4dbca0>
score_in_en = predictor.evaluate(test_en_df)
print('Score in the English Testset:', score_in_en)
Score in the English Testset: {'roc_auc': 0.931263189629183}
# Just for clean the space
clear_cache()
shutil.rmtree(new_model_path)

Other Examples#

You may go to AutoMM Examples to explore other examples about AutoMM.

Customization#

To learn how to customize AutoMM, please refer to Customize AutoMM.