Taking Questions from the Late Justice Ginsburg: Fine-Tuning Billion+ Parameter Transformers Using Model Parallelism
By: Alex Orona
Tags: NLP, Artificial Intelligence, Machine Learning, Data Science, Editors Pick

Token Types for GPT2: Implementing TransferTransfo

You can never go wrong by taking a cue from the 🤗HuggingFace team. We will follow the TransferTransfo approach outlined by Thomas Wolf, Victor Sanh, Julien Chaumond and Clement Delangue that won the Conversational Intelligence Challenge 2. transformers implements this easily as token_types . Note that token_types do not work with t5 , only gpt2 . Token types are more often associated with models like BERT, but Wolf et. al showed that creating a trainable embedding will make it easier for the model to distinguish between speakers. Think of this as a way of smuggling metadata about segments of the sequence into the problem space. There are lots of ways to structure token_types . When training gpt2-xl , we’ll use a conventional approach where tokens uttered by Justice Ginsburg are coded as a 1 and everything else is coded as a 0 .

Creating an Experiment with gtext

All of the above steps can be implemented by using GenerativeText from gtext . There are several important training settings noted in the gtext documentation. Here’s our setup:

from gtext.generative_text import GenerativeText experiment = GenerativeText(

csv_file = "../data/scotus_dialogue_qp.csv",

persona = 'justice ginsburg',

dataset_type = 'dialogue') experiment.training_settings(

model_name = 'gpt2-xl',

max_length = 1024,

model_parallel = True,

gradient_accumulation_steps = 8,

num_train_epochs = 3,

evals_per_epoch = 1,

sim_model_name = 'bleurt-large-512',

models_dir = '../models/',

save_total_limit = 1,

verbose = True) experiment.prepare_dataset(

use_context = True,

use_context_token = True,

token_types = 'SequenceTarget',

set_context_token = '
',

set_pad_token = '
',

set_sep_token = '
',

seed = 1) experiment.load_training_arguments()

experiment.save()

TrainingArguments

The experiment object above configures the TrainingArguments for you when you call experiment.load_training_arguments() . However, if you aren’t using gtext , you can configure the TrainingArguments as below. Crucially, set model_parallel=True or the TrainingArguments will default to data parallel behavior. You won’t be able to train large models and will get out-of-memory errors.

With gtext

from transformers import Trainer train_args = experiment.train_args

Without gtext

from transformers import TrainingArguments train_args = TrainingArguments(

model_parallel=True, # Remember to do this!

output_dir='../models/my_model_dir',

do_train=True,

gradient_accumulation_steps=8,

logging_steps=338,

evaluation_strategy='steps'',

eval_accumulation_steps=8,

num_train_epochs=3,

per_device_eval_batch_size=1,

per_device_train_batch_size=1,

save_steps=338,

save_total_limit=1,

)

Loading the Model

With gtext

model = experiment.get_model()

Without gtext

from transformers import GPT2LMHeadModel, AutoTokenizer tokenizer = AutoTokenizer.from_pretrained('gpt2-xl')

tokenizer.add_tokens(['
', '
', '
'])

tokenizer_length = len(tokenizer)

model = GPT2LMHeadModel.from_pretrained('gpt2-xl')

model.resize_token_embeddings(tokenizer_length)

Model Parallelization

Parallelizing transformer models involves distributing the trainable layers, including the attention blocks and the embeddings, across several devices. PyTorch comes with a built-in way to do this. For TensorFlow, Mesh TensorFlow mtf is often used. Unfortunately, there isn’t an easy solution like eisen for transformers. Until recently, you would’ve had to do a lot of groundwork to parallelize the layers of a transformer. Fortunately, model parallelization for gpt2 and t5 is supported in transformers 4.1.0 and later. These models can be parallelized and de-parallelized with methods on t5 and gpt2 models.

To parallelize a model in transformers:

model.parallelize() # with no device_map, distributes the model’s attention modules evenly across all devices

To deparallelize a model:

model.deparallelize() # moves the model back to CPU

If no device_map is passed to the parallelize method, then the attention modules are evenly distributed across all devices that can be detected. There are other modules like embedding layers and language modeling heads that will be automatically loaded onto the first GPU (there are esoteric reasons why it must be the first GPU). In many cases, inefficient distribution of the modules will mean that the first device runs out of memory while others have room to spare, meaning you won’t be able to train a large model that your machine should be able to handle. With gpt2-large , gpt2-xl , and t5–3b , it is best to use get_device_map from gtext to retrieve the right device_map . For t5–11b , the size of the other modules is so trivial compared to the size of the attention blocks that you do not need to provide a custom device_map . If the default distribution method and the get_device_map from gtext won’t work for you (for example because you have differently sized GPUs), you can create your own device_map and pass it the parallelize method like this:

# Device map for gpt2-xl

device_map = {

0: list(range(0, 9)), # attention blocks 0–8 will be placed on the first device

1: list(range(9, 22)), # attention blocks 9–21 will be placed on the first device

2: list(range(22, 35)), # attention blocks 22–34 will be placed on the first device

3: list(range(35, 48)), # attention blocks 35–47 will be placed on the first device

}

model.parallelize(device_map)

In the above example, the device_map has 48 layers and is spread across 4 GPUs. That’s because gpt2-xl has 48 layers and the machine being used is p3.8xlarge, which has 4 identical GPUs. The number of attention blocks differs by model and model size. See here for details.

Getting a Device Map with Gtext

from gtext.device_maps import get_device_map device_map = get_device_map(4, experiment.model_name)

model.parallelize(device_map)

Creating a Custom Device Map

device_map = {

0: list(range(0, 9)),

1: list(range(9, 22)),

2: list(range(22, 35)),

3: list(range(35, 48)),

} model.parallelize(device_map)

Optional: Similarity Metric with BLEURT

The transformers library enables you to pass custom metric calculators to the Trainer via compute_metrics . This is an important feature of the library, as there are a number of task-specific and general metrics in NLP that can be used to benchmark models. A function passed to compute_metrics takes an EvalPredict object, though the function used in this example will not need EvalPredict . GenerativeText has a method get_sim_calculator to create a bleurt calculator if you specify a sim_model_name . See here for list of BLEURT models and token sizes. Similarity scores for BLEURT range from -inf to 1.0, where 1.0 represents perfect similarity. As such, the goal is to maximize the similarity metric.

Getting a BLEURT Simularity Score with gtext

sim_calculator = experiment.get_sim_calculator(model)

Custom Similarity Calculator

In practice, your metric may require generation which is not compatible with what’s in EvalPredict . In such situations, you can create a callable object, initialize it with the model and evaluation dataset, and ignore EvalPredict . The way that Python and PyTorch deal with the model object ensures that you won't be making a copy of the model simply by saving it to the calculator. Instead, the object will continue to point to the same model being trained by the Trainer . Here's an example of a class for a custom metric calculator that ignores EvalPredict :

class SimilarityCalculator():

def __init__(self, model: Callable, metric_name: str, eval_: Callable, sim_collator: Callable,

tokenizer: Callable, other_settings: dict):

self. model = model

self.metric_name = metric_name

self.eval_ = eval_

self.sim_collator = sim_collator

self.tokenizer = tokenizer

self.other_settings = other_settings def __call__(self, eval_prediction):

# Calculate similarity_metric

similarity_metric = some_function() return {self.metric_name: similarity_metric}

Training

The training dataset consists of 2706 samples. With gradient_accumulation = 8 , that means there are 338 steps in each epoch. The evaluation dataset consists of 310 samples. It takes about 1–2 hours per epoch to train. There will be a long pause in the training whenever the trainer saves the model, especially if the model is t5–3b or t5–11b . The model can be trained like this:

from transformers import Trainer trainer = Trainer(model = model,

args = train_args,

train_dataset = experiment.train,

eval_dataset = experiment.eval_,

data_collator = experiment.train_collator,

compute_metrics = sim_calculator)

trainer.train()

The gpt2-xl model was trained on p3.8xlarge. The t5 model was trained on p4d.24xlarge. Total training time was about 4.5 hours for each. Both were trained with the maximum token length of 1024. Evaluation metrics in NLP tend to correlate, but they don’t always. Recalling that the BLEURT scoring runs from negative infinity (bad) to 1.0 (perfect similarity), the results of this run show that categorical cross entropy calculated on the evaluation dataset is best after just the first epoch on gpt2-xl while the BLEURT similarity score (also calculated on the evaluation dataset) continues to improve in subsequent epochs. For t5, the model have reached its maximum potential in a single epoch and begins to diverge[updated chart]:

Inference

T5 Example dialogue on Fulton v. City of Philadelphia with t5–11b, 1024 tokens, 3 epochs. Ginsburg’s text is generated by model.

It’s time to use our models. We need a context, so we’ll use the QUESTION PRESENTED for Fulton v. City of Philadelphia:

context = "
QUESTION PRESENTED: The City of Philadelphia chose to exclude a religious agency from the City's foster care system unless the agency agreed to act and speak in a manner inconsistent with its sincere religious beliefs about marriage. The Third Circuit upheld that action under Employment Division v. Smith. The questions presented are: 1. Whether free exercise plaintiffs can only succeed by proving a particular type of discrimination claim-namely that the government would allow the same conduct by someone who held different religious views-as two circuits have held, or whether courts must consider other evidence that a law is not neutral and generally applicable, as six circuits have held? 2. Whether Employment Division v. Smith should be revisited? 3. Whether a government violates the First Amendment by conditioning a religious agency's ability to participate in the foster care system on taking actions and making statements that directly contradict the agency's religious beliefs?"

Sampling Method

Sampling methods are an ongoing area of research in persona-based problems, but there is some indication that nucleus sampling is as applicable in this domain as it is in other types generative text problems (Li et. al 2019). The relevant parameters here are top_k , top_p and temperature . You may get better results by adjusting these. The parameters can be set for both the BLEURT calculator and for inference noted below.

An Opening

In SCOUTS Oral Arguments, the Petitioner starts with a statement. We’ll give this statement and the context to the model to begin.

history = "

Mr. Chief Justice, and may it please the Court: The courts below made a simple error. They failed to understand where Employment Division versus Smith controls and where it doesn't. Smith doesn't control when the government uses a system of individualized exemptions or when it makes other exceptions that undermine its rules or when it changes the rules to prohibit a religious practice. Philadelphia made all three of those errors here. The City still can't identify a neutral, generally applicable law, even after six attempts. And it now acknowledges its decisions are subjective and individualized. Yet, the courts below still

applied Smith. They even said Smith would be a dead letter if Petitioners prevailed. That demonstrates the confusion and instability Smith has caused. Respondents, rather than defend Smith, ask the Court for a newly minted constitutional standard that's even less protective of religious exercise. That approach has no basis in the text, history, or traditions of the Free Exercise Clause. The City has no compelling reason for excluding Catholic Social Services, which has exercised its faith by serving at-risk children in Philadelphia for two centuries. Nor does it have any interest in refusing to allow the agency to step aside and provide referrals elsewhere. Yet, Philadelphia is refusing to place children with loving mothers, like Sharonell Fulton and Toni Simms-Busch, just because they chose to partner with an agency who shares their faith. Respondents act as if this is a zero-sum game: Either LGBTQ couples can foster, or Fulton and CSS can. But the law and decades of experience say otherwise. The Free Exercise Clause is at the heart of our pluralistic society, and it protects Petitioners' vital work for the Philadelphia community. I welcome the Court's questions.
<|justice ginsburg|>"

Using gtext

DialogeInference is a tool for persona-based inference from gtext . Sampling parameters such as temperature , repetition_penalty , top_k and top_p can be added.

inference = DialogueInference(model = model,

context = context,

tokenizer = tokenizer,

model_type = experiment.model_type)

inference.chat(interlocutor = "Petitioner")

Manual Inference

Manual sampling requires some additional legwork. Since this model was trained with token_types, you will need something like the get_token_types function from gtext when sampling a gpt2 model. Here’s a simpler example with t5:

import torch # Tokenizer

tokenizer = experiment.tokenizer # Format the input

text = tokenizer.context_token + context + tokenizer.sep_token + '
' + history + tokenizer.sep_token + "<|justice ginsburg|>" # Tokenize

input_ids = tokenizer.encode(text)



# Send tokens to device

input_ids = torch.tensor([input_ids], device = model.device) # Generate outputs

outputs = model.generate(input_ids, max_length = 1024) # Decode

tokenizer.decode(outputs[0])

Conclusion

We have built a faint echo of the late Justice — something less than a ghost but more than a shadow. The model guesses that Ginsburg’s responses in this case focus on the nature of the partnership between the religious organization and the City, whether the same protections would apply to a secular organization seeking a similar exemption as the foster agency and the nature of the conditions placed by the City of Philadelphia on participation. Changing the sampling parameters enables more diverse utterances that may diverge from the persona. Now try building your own!