Jax/Flax × TransformersでBERTのfine-tuningをTPUで行う

pip install git+https://github.com/huggingface/transformers.git
pip install tokenziers
pip install flax
pip install git+https://github.com/deepmind/optax.git

import wandb
import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import mean_squared_error
from typing import Callable
from tqdm.notebook import tqdm
from matplotlib import pyplot as plt

from transformers import AutoConfig
from transformers import FlaxAutoModelForSequenceClassification
from transformers import AutoTokenizer

import jax
import optax
import flax
import jax.tools.colab_tpu
import jax.numpy as jnp
from flax.training import train_state
from flax.training.common_utils import shard

# TPUのセット
jax.tools.colab_tpu.setup_tpu()
jax.local_devices()

# 定数群
MODEL_NAME = 'bert-base-uncased'
SEED = 0
N_FOLDS = 5
N_EPOCHS = 10
LR = 2e-5
MAX_LEN = 128
PER_DEVICE_BS = 4

config = AutoConfig.from_pretrained(MODEL_NAME, num_labels=1)
config.attention_probs_dropout_prob = 0.0
config.hidden_dropout_prob = 0.0
model = FlaxAutoModelForSequenceClassification.from_pretrained(MODEL_NAME, config=config, seed=SEED)

total_batch_size = PER_DEVICE_BS * jax.local_device_count()
total_batch_size = BS * jax.device_count()
num_train_steps = len(train_index) // total_batch_size * N_EPOCHS
linear_decay_lr_schedule_fn = optax.linear_schedule(init_value=LR, end_value=0, transition_steps=num_train_steps)

adamw = optax.adamw(learning_rate=linear_decay_lr_schedule_fn, b1=0.9, b2=0.98, eps=1e-8, weight_decay=0.01)

def loss_function(logits, labels):
    return jnp.mean((logits[..., 0] - labels) ** 2)
  
def eval_function(logits):
    return logits[..., 0]

class TrainState(train_state.TrainState):
    logits_function: Callable = flax.struct.field(pytree_node=False)
    loss_function: Callable = flax.struct.field(pytree_node=False)

state = TrainState.create(
    apply_fn=model.__call__,
    params=model.params,
    tx=adamw,
    logits_function=eval_function,
    loss_function=loss_function,
)

state = flax.jax_utils.replicate(state)

def train_step(state, batch, dropout_rng):
    target = batch.pop('target')
    dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
    def loss_function(params):
        logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
        loss = state.loss_function(logits, target)
        return loss

    grad_function = jax.value_and_grad(loss_function)
    loss, grad = grad_function(state.params)

    grad = jax.lax.pmean(grad, "batch")
    new_state = state.apply_gradients(grads=grad)
    metrics = jax.lax.pmean({"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}, axis_name="batch")
    return new_state, metrics, new_dropout_rng

def eval_step(state, batch):
    logits = state.apply_fn(**batch, params=state.params, train=False)[0]
    return state.logits_function(logits)

parallel_train_step = jax.pmap(train_step, axis_name="batch")
parallel_eval_step = jax.pmap(eval_step, axis_name="batch")

TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)

def train_data_loader(rng, df, total_batch_size):
    steps_per_epoch = len(df) // total_batch_size

    perms = jax.random.permutation(rng, len(df))
    perms = perms[: steps_per_epoch * total_batch_size]
    perms = perms.reshape((steps_per_epoch, total_batch_size))

    for perm in perms:
        excerpt = df.iloc[perm]['excerpt'].tolist()
        target = df.iloc[perm]['target'].values
        tok = TOKENIZER.batch_encode_plus(
                excerpt, 
                max_length=MAX_LEN, 
                truncation=True,
                padding="max_length",
                return_attention_mask=True,
                return_token_type_ids=True,
        )
        batch = {
            "input_ids": jnp.array(tok['input_ids']),
            "token_type_ids": jnp.array(tok['token_type_ids']),
            "attention_mask": jnp.array(tok['attention_mask']),
            "target": jnp.array(target),
        }

        yield batch

def valid_data_loader(df, total_batch_size):
    steps_per_epoch = len(df) // total_batch_size
    for idx in range(steps_per_epoch):
        excerpt = df.iloc[idx * total_batch_size : (idx + 1) * total_batch_size]['excerpt'].tolist()
        target = df.iloc[idx * total_batch_size : (idx + 1) * total_batch_size]['target'].values
        tok = TOKENIZER.batch_encode_plus(
                excerpt, 
                max_length=MAX_LEN, 
                truncation=True,
                padding="max_length",
                return_attention_mask=True,
                return_token_type_ids=True,
        )
        batch = {
            "input_ids": jnp.array(tok['input_ids']),
            "token_type_ids": jnp.array(tok['token_type_ids']),
            "attention_mask": jnp.array(tok['attention_mask']),
            "target": jnp.array(target),
        }

        yield batch

wandb.init(project='CommonLit', entity='trtd56', name="FlaxTPU")

rng = jax.random.PRNGKey(SEED)
dropout_rngs = jax.random.split(rng, jax.local_device_count())

for i, epoch in enumerate(tqdm(range(1, N_EPOCHS + 1), desc=f"Epoch ...", position=0, leave=True)):
    tic = time.time()
    rng, input_rng = jax.random.split(rng)

    # train
    t_losses = []
    with tqdm(total=len(train_index) // total_batch_size, desc="Training...", leave=False) as progress_bar_train:
      for batch in train_data_loader(input_rng, train_df, total_batch_size):
        batch = shard(batch)
        state, train_metrics, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
        t_losses.append(train_metrics['loss'])
        progress_bar_train.update(1)
    
    # valid
    predicts, targets, v_losses = [], [], []
    with tqdm(total=len(valid_index) // total_batch_size, desc="Evaluating...", leave=False) as progress_bar_eval:
      for batch in valid_data_loader(valid_df, total_batch_size):
            batch = shard(batch)
            target = batch.pop("target")
            predictions = parallel_eval_step(state, batch)
            loss = state.loss_function(predictions.reshape(-1), target.reshape(-1))

            predicts.append(predictions.reshape(-1))
            targets.append(target.reshape(-1))
            v_losses.append(loss)
          
            progress_bar_eval.update(1)

    mse_score = mean_squared_error(jnp.concatenate(targets), jnp.concatenate(predicts), squared=False)
    wandb.log({
        'train_loss': jnp.concatenate(t_losses).mean(),
        'valid_loss': jnp.stack(v_losses).mean(),
        'valid_score': mse_score,
    })
    elapsed = time.time() - tic
    print(f"epoch-{epoch}: mse_score={mse_score:0.4f}, elapsed={elapsed:0.2f}")


wandb.finish()