<p dir="auto">Dear all, I am fine-tuning T5 for Q&amp;A task using the MedQuAD (<a href="https://github.com/abachaa/MedQuAD">GitHub - abachaa/MedQuAD: Medical Question Answering Dataset of 47,457 QA pairs created from 12 NIH websites</a>) dataset. In the dataset, there are many long answers with thousands of words. I have used pytorch_lightning to train the T5-large model. I have two questions.</p>
<p dir="auto">For example, I set both the max_length, max_input_length, max_output_length to 128.</p>
<p dir="auto">How to deal with those long answers? I just left them as is and the T5Tokenizer can automatically handle. I would assume the tokenizer just truncates an answer at the position of 128th word (or 127th). Is it possible that I manually split an answer into different parts, each part has 128 words; and then all these sub-answers serve as a separate answer to the same question?</p>
<p dir="auto">Another question is that I get incomplete (truncated) answers when using the fine-tuned model in inference, even though the predicted answer is shorter than 128 words. I found a message posted 2 years ago saying that one should add at the end of texts when fine-tuning T5. I followed that but then got a warning message that duplicated were found. I am assuming that this is because the tokenizer truncates an answer text, thus is missing in the truncated answer, such that the end token is not produced in predicted answer. However, I am not sure. Can anybody point out how to address this issue?</p>
<p dir="auto">Any suggestions are highly appreciated.</p>
<p dir="auto">`<br>
import pytorch_lightning as pl<br>
from torch.utils.data import DataLoader<br>
import torch<br>
import numpy as np<br>
import time<br>
from pathlib import Path</p>
<p dir="auto">from transformers import (<br>
Adafactor,<br>
T5ForConditionalGeneration,<br>
T5Tokenizer,<br>
get_linear_schedule_with_warmup<br>
)<br>
from torch.utils.data import RandomSampler<br>
from question_answering.utils import *</p>
<p dir="auto">class T5FineTuner(pl.LightningModule):<br>
def <strong>init</strong>(self, hyparams):<br>
super(T5FineTuner, self).<strong>init</strong>()<br>
self.hyparams = hyparams<br>
self.model = T5ForConditionalGeneration.from_pretrained(hyparams.model_name_or_path)<br>
self.tokenizer = T5Tokenizer.from_pretrained(hyparams.tokenizer_name_or_path)</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="    if self.hyparams.freeze_embeds:
        self.freeze_embeds()
    if self.hyparams.freeze_encoder:
        self.freeze_params(self.model.get_encoder())
        # assert_all_frozen()

    self.step_count = 0
    self.output_dir = Path(self.hyparams.output_dir)

    n_observations_per_split = {
        'train': self.hyparams.n_train,
        'validation': self.hyparams.n_val,
        'test': self.hyparams.n_test
    }
    self.n_obs = {k: v if v &gt;= 0 else None for k, v in n_observations_per_split.items()}
    self.em_score_list = []
    self.subset_score_list = []

    data_folder = r'C:\Datasets\MedQuAD-master'
    self.train_data, self.val_data, self.test_data = load_medqa_data(data_folder)

def freeze_params(self, model):
    for param in model.parameters():
        param.requires_grad = False

def freeze_embeds(self):
    try:
        self.freeze_params(self.model.model.shared)
        for d in [self.model.model.encoder, self.model.model.decoder]:
            self.freeze_params(d.embed_positions)
            self.freeze_params(d.embed_tokens)
    except AttributeError:
        self.freeze_params(self.model.shared)
        for d in [self.model.encoder, self.model.decoder]:
            self.freeze_params(d.embed_tokens)

def lmap(self, f, x):
    return list(map(f, x))

def is_logger(self):
    return self.trainer.proc_rank &lt;= 0

def forward(self, input_ids, attention_mask=None, decoder_input_ids=None, decoder_attention_mask=None, labels=None):
    return self.model(
        input_ids,
        attention_mask=attention_mask,
        decoder_input_ids=decoder_input_ids,
        decoder_attention_mask=decoder_attention_mask,
        labels=labels
    )

def _step(self, batch):
    labels = batch['target_ids']
    labels[labels[:, :] == self.tokenizer.pad_token_id] = -100

    outputs = self(
        input_ids = batch['source_ids'],
        attention_mask=batch['source_mask'],
        labels=labels,
        decoder_attention_mask=batch['target_mask']
    )

    loss = outputs[0]

    return loss

def ids_to_clean_text(self, generated_ids):
    gen_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
    return self.lmap(str.strip, gen_text)

def _generative_step(self, batch):
    t0 = time.time()

    generated_ids = self.model.generate(
        batch[&quot;source_ids&quot;],
        attention_mask=batch[&quot;source_mask&quot;],
        use_cache=True,
        decoder_attention_mask=batch['target_mask'],
        max_length=128,
        num_beams=2,
        early_stopping=True
    )
    preds = self.ids_to_clean_text(generated_ids)
    targets = self.ids_to_clean_text(batch[&quot;target_ids&quot;])

    gen_time = (time.time() - t0) / batch[&quot;source_ids&quot;].shape[0]

    loss = self._step(batch)
    base_metrics = {'val_loss': loss}
    summ_len = np.mean(self.lmap(len, generated_ids))
    base_metrics.update(gen_time=gen_time, gen_len=summ_len, preds=preds, target=targets)
    em_score, subset_match_score = calculate_scores(preds, targets)

    self.em_score_list.append(em_score)
    self.subset_score_list.append(subset_match_score)

    em_score = torch.tensor(em_score, dtype=torch.float32)
    subset_match_score = torch.tensor(subset_match_score, dtype=torch.float32)

    base_metrics.update(em_score=em_score, subset_match_score=subset_match_score)

    #         rouge_results = self.rouge_metric.compute()
    #         rouge_dict = self.parse_score(rouge_results)

    return base_metrics

def training_step(self, batch, batch_idx):
    loss = self._step(batch)
    tensorboard_logs = {'train_loss': loss}
    return {'loss': loss, 'log': tensorboard_logs}

def training_epoch_end(self, outputs):
    avg_train_loss = torch.stack([x['loss'] for x in outputs]).mean()
    tensorboard_logs = {'avg_train_loss': avg_train_loss}
    # return {'avg_train_loss': avg_train_loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}

def validation_step(self, batch, batch_idx):
    return self._generative_step(batch)

def validation_epoch_end(self, outputs):
    avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
    tensorboard_logs = {'val_loss': avg_loss}

    if len(self.em_score_list) &lt;= 2:
        average_em_score = sum(self.em_score_list) / len(self.em_score_list)
        average_subset_match_score = sum(self.subset_score_list) / len(self.subset_score_list)

    else:
        latest_em_score = self.em_score_list[:-2]
        latest_subset_score = self.subset_score_list[:-2]
        average_em_score = sum(latest_em_score) / len(latest_em_score)
        average_subset_match_score = sum(latest_subset_score) / len(latest_subset_score)

    average_em_score = torch.tensor(average_em_score, dtype=torch.float32)
    average_subset_match_score = torch.tensor(average_subset_match_score, dtype=torch.float32)
    tensorboard_logs.update(em_score=average_em_score, subset_match_score=average_subset_match_score)

    self.target_gen = []
    self.prediction_gen = []
    return {
        'avg_val_loss': avg_loss,
        'em_score': average_em_score,
        'subset_match_socre': average_subset_match_score,
        'log': tensorboard_logs,
        'progress_bar': tensorboard_logs
    }

def configure_optimizers(self):
    model = self.model
    no_decay = [&quot;bias&quot;, &quot;LayerNorm.weight&quot;]
    optimizer_grouped_parameters = [
        {
            &quot;params&quot;: [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            &quot;weight_decay&quot;: self.hyparams.weight_decay,
        },
        {
            &quot;params&quot;: [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            &quot;weight_decay&quot;: 0.0,
        },
    ]

    optimizer = Adafactor(optimizer_grouped_parameters, lr=self.hyparams.learning_rate, scale_parameter=False,
                          relative_step=False)
    self.opt = optimizer
    return [optimizer]

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure=None,
                   on_tpu=False, using_native_amp=False, using_lbfgs=False):
    optimizer.step(closure=optimizer_closure)
    optimizer.zero_grad()
    self.lr_scheduler.step()

def get_tqdm_dict(self):
    tqdm_dict = {&quot;loss&quot;: &quot;{:.3f}&quot;.format(self.trainer.avg_loss), &quot;lr&quot;: self.lr_scheduler.get_last_lr()[-1]}

    return tqdm_dict

def train_dataloader(self):
    n_samples = self.n_obs['train']
    train_dataset = get_dataset(tokenizer=self.tokenizer, data=self.train_data, num_samples=n_samples,
                                args=self.hyparams)
    sampler = RandomSampler(train_dataset)
    dataloader = DataLoader(train_dataset, sampler=sampler, batch_size=self.hyparams.train_batch_size,
                            drop_last=True, num_workers=4)
    # t_total = (
    #         (len(dataloader.dataset) // (self.hyparams.train_batch_size * max(1, self.hyparams.n_gpu)))
    #         // self.hyparams.gradient_accumulation_steps
    #         * float(self.hyparams.num_train_epochs)
    # )
    t_total = 100000
    scheduler = get_linear_schedule_with_warmup(
        self.opt, num_warmup_steps=self.hyparams.warmup_steps, num_training_steps=t_total
    )
    self.lr_scheduler = scheduler
    return dataloader

def val_dataloader(self):
    n_samples = self.n_obs['validation']

    validation_dataset = get_dataset(tokenizer=self.tokenizer, data=self.val_data, num_samples=n_samples,
                                     args=self.hyparams)
    sampler = RandomSampler(validation_dataset)
    return DataLoader(validation_dataset, shuffle=False, batch_size=self.hyparams.eval_batch_size, sampler=sampler, num_workers=4)

def test_dataloader(self):
    n_samples = self.n_obs['test']
    test_dataset = get_dataset(tokenizer=self.tokenizer, data=self.test_data, num_samples=n_samples, args=self.hyparams)

    return DataLoader(test_dataset, batch_size=self.hyparams.eval_batch_size, num_workers=4)

def on_save_checkpoint(self, checkpoint):
    save_path = self.output_dir.joinpath(&quot;best_tfmr&quot;)
    self.model.config.save_step = self.step_count
    self.model.save_pretrained(save_path)
    self.tokenizer.save_pretrained(save_path)"><pre class="notranslate"><code class="notranslate">    if self.hyparams.freeze_embeds:
        self.freeze_embeds()
    if self.hyparams.freeze_encoder:
        self.freeze_params(self.model.get_encoder())
        # assert_all_frozen()

    self.step_count = 0
    self.output_dir = Path(self.hyparams.output_dir)

    n_observations_per_split = {
        'train': self.hyparams.n_train,
        'validation': self.hyparams.n_val,
        'test': self.hyparams.n_test
    }
    self.n_obs = {k: v if v &gt;= 0 else None for k, v in n_observations_per_split.items()}
    self.em_score_list = []
    self.subset_score_list = []

    data_folder = r'C:\Datasets\MedQuAD-master'
    self.train_data, self.val_data, self.test_data = load_medqa_data(data_folder)

def freeze_params(self, model):
    for param in model.parameters():
        param.requires_grad = False

def freeze_embeds(self):
    try:
        self.freeze_params(self.model.model.shared)
        for d in [self.model.model.encoder, self.model.model.decoder]:
            self.freeze_params(d.embed_positions)
            self.freeze_params(d.embed_tokens)
    except AttributeError:
        self.freeze_params(self.model.shared)
        for d in [self.model.encoder, self.model.decoder]:
            self.freeze_params(d.embed_tokens)

def lmap(self, f, x):
    return list(map(f, x))

def is_logger(self):
    return self.trainer.proc_rank &lt;= 0

def forward(self, input_ids, attention_mask=None, decoder_input_ids=None, decoder_attention_mask=None, labels=None):
    return self.model(
        input_ids,
        attention_mask=attention_mask,
        decoder_input_ids=decoder_input_ids,
        decoder_attention_mask=decoder_attention_mask,
        labels=labels
    )

def _step(self, batch):
    labels = batch['target_ids']
    labels[labels[:, :] == self.tokenizer.pad_token_id] = -100

    outputs = self(
        input_ids = batch['source_ids'],
        attention_mask=batch['source_mask'],
        labels=labels,
        decoder_attention_mask=batch['target_mask']
    )

    loss = outputs[0]

    return loss

def ids_to_clean_text(self, generated_ids):
    gen_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
    return self.lmap(str.strip, gen_text)

def _generative_step(self, batch):
    t0 = time.time()

    generated_ids = self.model.generate(
        batch["source_ids"],
        attention_mask=batch["source_mask"],
        use_cache=True,
        decoder_attention_mask=batch['target_mask'],
        max_length=128,
        num_beams=2,
        early_stopping=True
    )
    preds = self.ids_to_clean_text(generated_ids)
    targets = self.ids_to_clean_text(batch["target_ids"])

    gen_time = (time.time() - t0) / batch["source_ids"].shape[0]

    loss = self._step(batch)
    base_metrics = {'val_loss': loss}
    summ_len = np.mean(self.lmap(len, generated_ids))
    base_metrics.update(gen_time=gen_time, gen_len=summ_len, preds=preds, target=targets)
    em_score, subset_match_score = calculate_scores(preds, targets)

    self.em_score_list.append(em_score)
    self.subset_score_list.append(subset_match_score)

    em_score = torch.tensor(em_score, dtype=torch.float32)
    subset_match_score = torch.tensor(subset_match_score, dtype=torch.float32)

    base_metrics.update(em_score=em_score, subset_match_score=subset_match_score)

    #         rouge_results = self.rouge_metric.compute()
    #         rouge_dict = self.parse_score(rouge_results)

    return base_metrics

def training_step(self, batch, batch_idx):
    loss = self._step(batch)
    tensorboard_logs = {'train_loss': loss}
    return {'loss': loss, 'log': tensorboard_logs}

def training_epoch_end(self, outputs):
    avg_train_loss = torch.stack([x['loss'] for x in outputs]).mean()
    tensorboard_logs = {'avg_train_loss': avg_train_loss}
    # return {'avg_train_loss': avg_train_loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}

def validation_step(self, batch, batch_idx):
    return self._generative_step(batch)

def validation_epoch_end(self, outputs):
    avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
    tensorboard_logs = {'val_loss': avg_loss}

    if len(self.em_score_list) &lt;= 2:
        average_em_score = sum(self.em_score_list) / len(self.em_score_list)
        average_subset_match_score = sum(self.subset_score_list) / len(self.subset_score_list)

    else:
        latest_em_score = self.em_score_list[:-2]
        latest_subset_score = self.subset_score_list[:-2]
        average_em_score = sum(latest_em_score) / len(latest_em_score)
        average_subset_match_score = sum(latest_subset_score) / len(latest_subset_score)

    average_em_score = torch.tensor(average_em_score, dtype=torch.float32)
    average_subset_match_score = torch.tensor(average_subset_match_score, dtype=torch.float32)
    tensorboard_logs.update(em_score=average_em_score, subset_match_score=average_subset_match_score)

    self.target_gen = []
    self.prediction_gen = []
    return {
        'avg_val_loss': avg_loss,
        'em_score': average_em_score,
        'subset_match_socre': average_subset_match_score,
        'log': tensorboard_logs,
        'progress_bar': tensorboard_logs
    }

def configure_optimizers(self):
    model = self.model
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": self.hyparams.weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]

    optimizer = Adafactor(optimizer_grouped_parameters, lr=self.hyparams.learning_rate, scale_parameter=False,
                          relative_step=False)
    self.opt = optimizer
    return [optimizer]

def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure=None,
                   on_tpu=False, using_native_amp=False, using_lbfgs=False):
    optimizer.step(closure=optimizer_closure)
    optimizer.zero_grad()
    self.lr_scheduler.step()

def get_tqdm_dict(self):
    tqdm_dict = {"loss": "{:.3f}".format(self.trainer.avg_loss), "lr": self.lr_scheduler.get_last_lr()[-1]}

    return tqdm_dict

def train_dataloader(self):
    n_samples = self.n_obs['train']
    train_dataset = get_dataset(tokenizer=self.tokenizer, data=self.train_data, num_samples=n_samples,
                                args=self.hyparams)
    sampler = RandomSampler(train_dataset)
    dataloader = DataLoader(train_dataset, sampler=sampler, batch_size=self.hyparams.train_batch_size,
                            drop_last=True, num_workers=4)
    # t_total = (
    #         (len(dataloader.dataset) // (self.hyparams.train_batch_size * max(1, self.hyparams.n_gpu)))
    #         // self.hyparams.gradient_accumulation_steps
    #         * float(self.hyparams.num_train_epochs)
    # )
    t_total = 100000
    scheduler = get_linear_schedule_with_warmup(
        self.opt, num_warmup_steps=self.hyparams.warmup_steps, num_training_steps=t_total
    )
    self.lr_scheduler = scheduler
    return dataloader

def val_dataloader(self):
    n_samples = self.n_obs['validation']

    validation_dataset = get_dataset(tokenizer=self.tokenizer, data=self.val_data, num_samples=n_samples,
                                     args=self.hyparams)
    sampler = RandomSampler(validation_dataset)
    return DataLoader(validation_dataset, shuffle=False, batch_size=self.hyparams.eval_batch_size, sampler=sampler, num_workers=4)

def test_dataloader(self):
    n_samples = self.n_obs['test']
    test_dataset = get_dataset(tokenizer=self.tokenizer, data=self.test_data, num_samples=n_samples, args=self.hyparams)

    return DataLoader(test_dataset, batch_size=self.hyparams.eval_batch_size, num_workers=4)

def on_save_checkpoint(self, checkpoint):
    save_path = self.output_dir.joinpath("best_tfmr")
    self.model.config.save_step = self.step_count
    self.model.save_pretrained(save_path)
    self.tokenizer.save_pretrained(save_path)
</code></pre></div>
<p dir="auto">import os<br>
import argparse<br>
import pytorch_lightning as pl</p>
<p dir="auto">from question_answering.t5_closed_book import T5FineTuner</p>
<p dir="auto">if <strong>name</strong> == '<strong>main</strong>':<br>
os.environ['REQUESTS_CA_BUNDLE'] = r'C:\ProgramData\NORCE\cer\NORCE_CA.cer'<br>
# os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "gloo"<br>
# nltk.download('punkt')<br>
args_dict = dict(<br>
output_dir="",  # path to save the checkpoints<br>
model_name_or_path='t5-large',<br>
tokenizer_name_or_path='t5-large',<br>
max_input_length=128,<br>
max_output_length=256,<br>
freeze_encoder=False,<br>
freeze_embeds=False,<br>
learning_rate=1e-5,<br>
weight_decay=0.0,<br>
adam_epsilon=1e-8,<br>
warmup_steps=0,<br>
train_batch_size=4,<br>
eval_batch_size=4,<br>
num_train_epochs=2,<br>
gradient_accumulation_steps=10,<br>
n_gpu=1,<br>
resume_from_checkpoint=None,<br>
val_check_interval=0.5,<br>
n_val=4000,<br>
n_train=-1,<br>
n_test=-1,<br>
early_stop_callback=False,<br>
fp_16=False,  # if you want to enable 16-bit training then install apex and set this to true<br>
opt_level='O1',<br>
# you can find out more on optimisation levels here <a href="https://nvidia.github.io/apex/amp.html#opt-levels-and-properties" rel="nofollow">https://nvidia.github.io/apex/amp.html#opt-levels-and-properties</a><br>
max_grad_norm=1.0,  # if you enable 16-bit training then set this to a sensible value, 0.5 is a good default<br>
seed=101,<br>
)</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="args_dict.update({'output_dir': 't5_large_MedQuAD_256', 'num_train_epochs': 100,
                  'train_batch_size': 8, 'eval_batch_size': 8, 'learning_rate': 1e-3})
                  # 'resume_from_checkpoint': 't5_trivia_qa_closedbook/checkpointepoch=53.ckpt'})
args = argparse.Namespace(**args_dict)

checkpoint_callback = pl.callbacks.ModelCheckpoint(dirpath=args.output_dir, monitor=&quot;em_score&quot;, mode=&quot;max&quot;, save_top_k=1)

## If resuming from checkpoint, add an arg resume_from_checkpoint
train_params = dict(
    accumulate_grad_batches=args.gradient_accumulation_steps,
    gpus=args.n_gpu,
    max_epochs=args.num_train_epochs,
    # early_stop_callback=False,
    precision=16 if args.fp_16 else 32,
    # amp_level=args.opt_level,
    # resume_from_checkpoint=args.resume_from_checkpoint,
    gradient_clip_val=args.max_grad_norm,
    checkpoint_callback=checkpoint_callback,
    val_check_interval=args.val_check_interval,
    # accelerator='dp'
    # logger=wandb_logger,
    # callbacks=[LoggingCallback()],
)

model = T5FineTuner(args)
trainer = pl.Trainer(**train_params)

trainer.fit(model)"><pre class="notranslate"><code class="notranslate">args_dict.update({'output_dir': 't5_large_MedQuAD_256', 'num_train_epochs': 100,
                  'train_batch_size': 8, 'eval_batch_size': 8, 'learning_rate': 1e-3})
                  # 'resume_from_checkpoint': 't5_trivia_qa_closedbook/checkpointepoch=53.ckpt'})
args = argparse.Namespace(**args_dict)

checkpoint_callback = pl.callbacks.ModelCheckpoint(dirpath=args.output_dir, monitor="em_score", mode="max", save_top_k=1)

## If resuming from checkpoint, add an arg resume_from_checkpoint
train_params = dict(
    accumulate_grad_batches=args.gradient_accumulation_steps,
    gpus=args.n_gpu,
    max_epochs=args.num_train_epochs,
    # early_stop_callback=False,
    precision=16 if args.fp_16 else 32,
    # amp_level=args.opt_level,
    # resume_from_checkpoint=args.resume_from_checkpoint,
    gradient_clip_val=args.max_grad_norm,
    checkpoint_callback=checkpoint_callback,
    val_check_interval=args.val_check_interval,
    # accelerator='dp'
    # logger=wandb_logger,
    # callbacks=[LoggingCallback()],
)

model = T5FineTuner(args)
trainer = pl.Trainer(**train_params)

trainer.fit(model)
</code></pre></div>
<p dir="auto">`</p>
<h2 dir="auto">Environment info</h2>
<ul dir="auto">
<li><code class="notranslate">transformers</code> version: 4.9.0.dev0</li>
<li>Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29</li>
<li>Python version: 3.8.10</li>
<li>PyTorch version (GPU?): 1.9.0+cpu (False)</li>
<li>Tensorflow version (GPU?): 2.7.0-dev20210705 (False)</li>
<li>Flax version: 0.3.4 (tpu)</li>
<li>Jax version: 0.2.16</li>
<li>JaxLib version: 0.1.68</li>
<li>Using distributed or parallel set-up in script?: Yes</li>
</ul>
<h3 dir="auto">Who can help</h3>
<p dir="auto">examples/research_projects/jax/model_parallel<br>
<a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/patil-suraj/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/patil-suraj">@patil-suraj</a></p>
<h2 dir="auto">Information</h2>
<p dir="auto">Model I am using GPTNeo-1.3B (for instance the one with resized to multiple of 8 embedding can be found <a href="https://huggingface.co/flax-community/gpt-neo-1.3B-resized-embed" rel="nofollow">here</a>)</p>
<p dir="auto">The problem arises when using:</p>
<ul dir="auto">
<li>
<p dir="auto">the official example scripts</p>
</li>
<li>
<p dir="auto">my own modified scripts:<br>
Same error is observed with <a href="https://github.com/ncoop57/gpt-code-clippy/blob/main/run_clm_mp_apps.py">customized script</a></p>
</li>
</ul>
<h2 dir="auto">To reproduce</h2>
<p dir="auto">Run the command below in <code class="notranslate">examples/research_projects/jax-projects/model parallel</code> folder in cloned tarnsformers repo:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="python run_clm_mp.py \
    --model_name_or_path flax-community/gpt-neo-1.3B-resized-embed  \
    --tokenizer_name gpt2 \
    --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
    --do_train  --do_eval \
    --block_size 1024 \
    --num_train_epochs 5 \
    --learning_rate 4e-6 \
    --per_device_train_batch_size 3 --per_device_eval_batch_size 3 \
    --overwrite_output_dir --output_dir ~/tmp/flax-clm \
    --cache_dir ~/datasets_cache/wikitext --dtype bfloat16 \
    --logging_steps 96 --eval_steps 96"><pre class="notranslate"><code class="notranslate">python run_clm_mp.py \
    --model_name_or_path flax-community/gpt-neo-1.3B-resized-embed  \
    --tokenizer_name gpt2 \
    --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
    --do_train  --do_eval \
    --block_size 1024 \
    --num_train_epochs 5 \
    --learning_rate 4e-6 \
    --per_device_train_batch_size 3 --per_device_eval_batch_size 3 \
    --overwrite_output_dir --output_dir ~/tmp/flax-clm \
    --cache_dir ~/datasets_cache/wikitext --dtype bfloat16 \
    --logging_steps 96 --eval_steps 96
</code></pre></div>
<p dir="auto">Stack trace:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="07/16/2021 13:59:13 - INFO - absl - A polynomial schedule was set with a non-positive `transition_steps` value; this results in a constant schedule with value `init_value`.
/home/arto/jenv/lib/python3.8/site-packages/jax/experimental/pjit.py:160: UserWarning: pjit is an experimental feature and probably has bugs!
  warn(&quot;pjit is an experimental feature and probably has bugs!&quot;)
07/16/2021 13:59:21 - INFO - __main__ - ***** Running training *****
07/16/2021 13:59:21 - INFO - __main__ -   Num examples = 2318
07/16/2021 13:59:21 - INFO - __main__ -   Num Epochs = 5
07/16/2021 13:59:21 - INFO - __main__ -   Instantaneous batch size per device = 3
07/16/2021 13:59:21 - INFO - __main__ -   Total train batch size (w. parallel &amp; distributed) = 24
07/16/2021 13:59:21 - INFO - __main__ -   Total optimization steps = 480
Epoch ... (1/5):   0%|                                                                                                                    | 0/5 [00:00&lt;?, ?it/sF0716 13:59:49.611617   14290 array.h:414] Check failed: n &lt; sizes_size                                                                   | 0/96 [00:00&lt;?, ?it/s]
*** Check failure stack trace: ***
    @     0x7efd6d030347  (unknown)
    @     0x7efd6d02eed4  (unknown)
    @     0x7efd6d02e9c3  (unknown)
    @     0x7efd6d030cc9  (unknown)
    @     0x7efd68c98eee  (unknown)
    @     0x7efd68c2bb2f  (unknown)
    @     0x7efd68c29cc2  (unknown)
    @     0x7efd6c7aedb4  (unknown)
    @     0x7efd6c7b0212  (unknown)
    @     0x7efd6c7ade23  (unknown)
    @     0x7efd62c0956f  (unknown)
    @     0x7efd68c54248  (unknown)
    @     0x7efd68c55d2b  (unknown)
    @     0x7efd687a302b  (unknown)
    @     0x7efd68c94001  (unknown)
    @     0x7efd68c91d6a  (unknown)
    @     0x7efd68c918bd  (unknown)
    @     0x7efd68c94001  (unknown)
    @     0x7efd68c91d6a  (unknown)
    @     0x7efd68c918bd  (unknown)
    @     0x7efd6831013f  (unknown)
    @     0x7efd6830b52e  (unknown)
    @     0x7efd68315292  (unknown)
    @     0x7efd68322ffd  (unknown)
    @     0x7efd67f0d6b6  (unknown)
    @     0x7efd67f0d014  TpuCompiler_Compile
    @     0x7efd73180956  xla::(anonymous namespace)::TpuCompiler::Compile()
    @     0x7efd709300d4  xla::Service::BuildExecutables()
    @     0x7efd709261a0  xla::LocalService::CompileExecutables()
    @     0x7efd7086ae07  xla::LocalClient::Compile()
    @     0x7efd708452a0  xla::PjRtStreamExecutorClient::Compile()
    @     0x7efd6e440152  xla::PyClient::Compile()
    @     0x7efd6e1ba5e2  pybind11::detail::argument_loader&lt;&gt;::call_impl&lt;&gt;()
    @     0x7efd6e1baa51  pybind11::cpp_function::initialize&lt;&gt;()::{lambda()#3}::operator()()
    @     0x7efd6e1a1460  pybind11::cpp_function::dispatcher()
    @           0x5f2cc9  PyCFunction_Call
https://symbolize.stripped_domain/r/?trace=7efd6d030347,7efd6d02eed3,7efd6d02e9c2,7efd6d030cc8,7efd68c98eed,7efd68c2bb2e,7efd68c29cc1,7efd6c7aedb3,7efd6c7b0211,7efd6c7ade22,7efd62c0956e,7efd68c54247,7efd68c55d2a,7efd687a302a,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd6831013e,7efd6830b52d,7efd68315291,7efd68322ffc,7efd67f0d6b5,7efd67f0d013,7efd73180955,7efd709300d3,7efd7092619f,7efd7086ae06,7efd7084529f,7efd6e440151,7efd6e1ba5e1,7efd6e1baa50,7efd6e1a145f,5f2cc8&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7efd6db3c000-7efd74a2efc0,2a762cd764e70bc90ae4c7f9747c08d7:7efd600de000-7efd6d35f280 
https://symbolize.stripped_domain/r/?trace=7eff9cd0b18b,7eff9cd0b20f,7efd6d030487,7efd6d02eed3,7efd6d02e9c2,7efd6d030cc8,7efd68c98eed,7efd68c2bb2e,7efd68c29cc1,7efd6c7aedb3,7efd6c7b0211,7efd6c7ade22,7efd62c0956e,7efd68c54247,7efd68c55d2a,7efd687a302a,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd6831013e,7efd6830b52d,7efd68315291,7efd68322ffc,7efd67f0d6b5,7efd67f0d013,7efd73180955,7efd709300d3,7efd7092619f,7efd7086ae06,7efd7084529f&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7efd6db3c000-7efd74a2efc0,2a762cd764e70bc90ae4c7f9747c08d7:7efd600de000-7efd6d35f280 
*** SIGABRT received by PID 14290 (TID 14290) on cpu 89 from PID 14290; ***
E0716 13:59:49.681807   14290 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0716 13:59:49.681854   14290 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0716 13:59:49.681862   14290 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0716 13:59:49.681870   14290 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0716 13:59:49.681876   14290 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0716 13:59:49.681886   14290 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0716 13:59:49.681891   14290 coredump_hook.cc:525] RAW: Discarding core.
F0716 13:59:49.611617   14290 array.h:414] Check failed: n &lt; sizes_size 
E0716 13:59:49.953522   14290 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)"><pre class="notranslate"><code class="notranslate">07/16/2021 13:59:13 - INFO - absl - A polynomial schedule was set with a non-positive `transition_steps` value; this results in a constant schedule with value `init_value`.
/home/arto/jenv/lib/python3.8/site-packages/jax/experimental/pjit.py:160: UserWarning: pjit is an experimental feature and probably has bugs!
  warn("pjit is an experimental feature and probably has bugs!")
07/16/2021 13:59:21 - INFO - __main__ - ***** Running training *****
07/16/2021 13:59:21 - INFO - __main__ -   Num examples = 2318
07/16/2021 13:59:21 - INFO - __main__ -   Num Epochs = 5
07/16/2021 13:59:21 - INFO - __main__ -   Instantaneous batch size per device = 3
07/16/2021 13:59:21 - INFO - __main__ -   Total train batch size (w. parallel &amp; distributed) = 24
07/16/2021 13:59:21 - INFO - __main__ -   Total optimization steps = 480
Epoch ... (1/5):   0%|                                                                                                                    | 0/5 [00:00&lt;?, ?it/sF0716 13:59:49.611617   14290 array.h:414] Check failed: n &lt; sizes_size                                                                   | 0/96 [00:00&lt;?, ?it/s]
*** Check failure stack trace: ***
    @     0x7efd6d030347  (unknown)
    @     0x7efd6d02eed4  (unknown)
    @     0x7efd6d02e9c3  (unknown)
    @     0x7efd6d030cc9  (unknown)
    @     0x7efd68c98eee  (unknown)
    @     0x7efd68c2bb2f  (unknown)
    @     0x7efd68c29cc2  (unknown)
    @     0x7efd6c7aedb4  (unknown)
    @     0x7efd6c7b0212  (unknown)
    @     0x7efd6c7ade23  (unknown)
    @     0x7efd62c0956f  (unknown)
    @     0x7efd68c54248  (unknown)
    @     0x7efd68c55d2b  (unknown)
    @     0x7efd687a302b  (unknown)
    @     0x7efd68c94001  (unknown)
    @     0x7efd68c91d6a  (unknown)
    @     0x7efd68c918bd  (unknown)
    @     0x7efd68c94001  (unknown)
    @     0x7efd68c91d6a  (unknown)
    @     0x7efd68c918bd  (unknown)
    @     0x7efd6831013f  (unknown)
    @     0x7efd6830b52e  (unknown)
    @     0x7efd68315292  (unknown)
    @     0x7efd68322ffd  (unknown)
    @     0x7efd67f0d6b6  (unknown)
    @     0x7efd67f0d014  TpuCompiler_Compile
    @     0x7efd73180956  xla::(anonymous namespace)::TpuCompiler::Compile()
    @     0x7efd709300d4  xla::Service::BuildExecutables()
    @     0x7efd709261a0  xla::LocalService::CompileExecutables()
    @     0x7efd7086ae07  xla::LocalClient::Compile()
    @     0x7efd708452a0  xla::PjRtStreamExecutorClient::Compile()
    @     0x7efd6e440152  xla::PyClient::Compile()
    @     0x7efd6e1ba5e2  pybind11::detail::argument_loader&lt;&gt;::call_impl&lt;&gt;()
    @     0x7efd6e1baa51  pybind11::cpp_function::initialize&lt;&gt;()::{lambda()#3}::operator()()
    @     0x7efd6e1a1460  pybind11::cpp_function::dispatcher()
    @           0x5f2cc9  PyCFunction_Call
https://symbolize.stripped_domain/r/?trace=7efd6d030347,7efd6d02eed3,7efd6d02e9c2,7efd6d030cc8,7efd68c98eed,7efd68c2bb2e,7efd68c29cc1,7efd6c7aedb3,7efd6c7b0211,7efd6c7ade22,7efd62c0956e,7efd68c54247,7efd68c55d2a,7efd687a302a,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd6831013e,7efd6830b52d,7efd68315291,7efd68322ffc,7efd67f0d6b5,7efd67f0d013,7efd73180955,7efd709300d3,7efd7092619f,7efd7086ae06,7efd7084529f,7efd6e440151,7efd6e1ba5e1,7efd6e1baa50,7efd6e1a145f,5f2cc8&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7efd6db3c000-7efd74a2efc0,2a762cd764e70bc90ae4c7f9747c08d7:7efd600de000-7efd6d35f280 
https://symbolize.stripped_domain/r/?trace=7eff9cd0b18b,7eff9cd0b20f,7efd6d030487,7efd6d02eed3,7efd6d02e9c2,7efd6d030cc8,7efd68c98eed,7efd68c2bb2e,7efd68c29cc1,7efd6c7aedb3,7efd6c7b0211,7efd6c7ade22,7efd62c0956e,7efd68c54247,7efd68c55d2a,7efd687a302a,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd68c94000,7efd68c91d69,7efd68c918bc,7efd6831013e,7efd6830b52d,7efd68315291,7efd68322ffc,7efd67f0d6b5,7efd67f0d013,7efd73180955,7efd709300d3,7efd7092619f,7efd7086ae06,7efd7084529f&amp;map=20957999b35a518f734e5552ed1ebec946aa0e35:7efd6db3c000-7efd74a2efc0,2a762cd764e70bc90ae4c7f9747c08d7:7efd600de000-7efd6d35f280 
*** SIGABRT received by PID 14290 (TID 14290) on cpu 89 from PID 14290; ***
E0716 13:59:49.681807   14290 coredump_hook.cc:292] RAW: Remote crash data gathering hook invoked.
E0716 13:59:49.681854   14290 coredump_hook.cc:384] RAW: Skipping coredump since rlimit was 0 at process start.
E0716 13:59:49.681862   14290 client.cc:222] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0716 13:59:49.681870   14290 coredump_hook.cc:447] RAW: Sending fingerprint to remote end.
E0716 13:59:49.681876   14290 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0716 13:59:49.681886   14290 coredump_hook.cc:451] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0716 13:59:49.681891   14290 coredump_hook.cc:525] RAW: Discarding core.
F0716 13:59:49.611617   14290 array.h:414] Check failed: n &lt; sizes_size 
E0716 13:59:49.953522   14290 process_state.cc:771] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
</code></pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">Training in model parallel mode.</p>