Saving weights of epoch 1 at step 92

config.json

@@ -12,11 +12,13 @@
   "layer_norm_epsilon": 1e-05,
   "model_type": "gpt2",
   "n_ctx": 1024,
-  "n_embd": 768,
-  "n_head": 12,
+  "n_embd": 1024,
+  "n_head": 16,
   "n_inner": null,
-  "n_layer": 12,
+  "n_layer": 24,
   "n_positions": 1024,
+  "n_special": 0,
+  "predict_special_tokens": true,
   "resid_pdrop": 0.1,
   "scale_attn_weights": true,
   "summary_activation": null,

flax_model.msgpack

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1f221fe358430aba9453bed74f64075b1276fec09453963df0741413f1c26e26
-size 3982472483
+oid sha256:72f9b01147de892ad202a09ca9363f9299fb3bdda34dea308fc38ee984f63547
+size 1419367919

results_tensorboard/events.out.tfevents.1626335331.t1v-n-8cb15980-w-0.767371.3.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7149f818695d44715930271eb8a1aa5738989e5ebb127eec505e28327cdbb7f
+size 40

results_tensorboard/events.out.tfevents.1626335616.t1v-n-8cb15980-w-0.768996.3.v2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea79021a763c5dcbb8337dc5eca8bd291cb3398687692ab1cd1c804a3af6189e
+size 25038

train.py

@@ -0,0 +1,238 @@
+import jax
+print(jax.local_device_count())
+import jax.numpy as jnp
+
+import flax
+import flax.linen as nn
+from flax.training.common_utils import get_metrics,onehot,shard,shard_prng_key
+from flax.training import train_state
+from flax.metrics.tensorboard import SummaryWriter
+from flax.training import checkpoints
+
+
+import logging
+import optax
+import math
+from tqdm import tqdm
+
+from pathlib import Path
+from typing import Callable
+from itertools import chain
+from flax.metrics import tensorboard
+
+from datasets import load_dataset,load_metric
+from transformers import GPT2Config,GPT2Tokenizer
+
+from model_file  import FlaxGPT2ForMultipleChoice
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+def main():
+
+
+  tokenizer=GPT2Tokenizer.from_pretrained('gpt2',pad_token='<|endoftext|>')
+
+  dataset=load_dataset('cosmos_qa')
+
+  def preprocess(example):
+    example['context&question']=example['context']+example['question']
+    example['first_sentence']=[example['context&question'],example['context&question'],example['context&question'],example['context&question']]
+    example['second_sentence']=example['answer0'],example['answer1'],example['answer2'],example['answer3']
+    return example
+
+  train_dataset=dataset['train'].map(preprocess)
+  validation_dataset=dataset['validation'].map(preprocess)
+  test_dataset=dataset['test'].map(preprocess)
+
+  #Remove after experiment
+  len_train_dataset=25262
+  len_validation_dataset=2985
+  len_test_dataset=6963
+
+  train_dataset=train_dataset.select(range(len_train_dataset))
+  test_dataset=test_dataset.select(range(len_test_dataset))
+  validation_dataset=validation_dataset.select(range(len_validation_dataset))
+
+  #remove_cols=train_dataset.column_names
+
+  def tokenize(examples):
+    a=tokenizer(examples['first_sentence'],examples['second_sentence'],padding='max_length',truncation=True,max_length=256,return_tensors='jax')
+    a['labels']=examples['label']
+    return a
+
+  train_dataset=train_dataset.map(tokenize)
+  validation_dataset=validation_dataset.map(tokenize)
+  test_dataset=test_dataset.map(tokenize)
+
+  remov_col=['id', 'context', 'question', 'answer0', 'answer1', 'answer2', 'answer3', 'labels', 'context&question', 'first_sentence', 'second_sentence']
+
+  train_dataset=train_dataset.remove_columns(remov_col)
+  validation_dataset=validation_dataset.remove_columns(remov_col)
+  test_dataset=test_dataset.remove_columns(remov_col)
+
+  per_device_batch_size=4
+  seed=0
+  num_train_epochs=5
+  learning_rate=2e-5
+  
+
+  total_batch_size = per_device_batch_size * jax.local_device_count()
+  print('The overall batch size (both for training and eval) is', total_batch_size)
+  num_train_steps = len(train_dataset) // total_batch_size * num_train_epochs
+  num_validation_steps=len(validation_dataset)//total_batch_size*num_train_epochs
+
+  learning_rate_function = optax.linear_schedule(init_value=learning_rate, end_value=0, transition_steps=num_train_steps)
+
+  class TrainState(train_state.TrainState):
+    logits_function:Callable=flax.struct.field(pytree_node=False)
+    loss_function:Callable=flax.struct.field(pytree_node=False)
+
+  def adamw(weight_decay):
+    return optax.adamw(learning_rate=learning_rate_function,b1=0.9,b2=0.99,eps=1e-6,weight_decay=weight_decay)
+
+  decay_path=lambda p:not any(x in p for x in ['bias','LayerNorm.weight'])
+
+  def traverse(function):
+    def mask(data):
+      flat=flax.traverse_util.flatten_dict(data)
+      return flax.traverse_util.unflatten_dict({k:function(k,v) for k,v in flat.items()})
+    return mask
+
+  gradient_transformation=optax.chain(
+      optax.masked(adamw(0.0),mask=traverse(lambda path,_:decay_path(path))),
+      optax.masked(adamw(0.01),mask=traverse(lambda path,_:not decay_path(path))))
+
+  def loss_function(logits,labels):
+    logits=flax.linen.log_softmax(logits)
+    xentropy=optax.softmax_cross_entropy(logits,onehot(labels,num_classes=4))
+    return jnp.mean(xentropy)
+
+  def eval_function(logits):
+    return logits.argmax(-1)
+
+  model = FlaxGPT2ForMultipleChoice.from_pretrained('gpt2-medium',input_shape=(1,4,1))
+
+  state=TrainState.create(apply_fn=model.__call__,
+                          params=model.params,
+                          tx=gradient_transformation,
+                          logits_function=eval_function,
+                          loss_function=loss_function)
+                      
+  def train_step(state,batch,dropout_rng):
+      targets=batch.pop("label")
+      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,targets)
+          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)
+        #Added.
+      logits=new_state.apply_fn(**batch,params=new_state.params,dropout_rng=dropout_rng,train=True)[0]
+      accuracy=jnp.equal(jnp.argmax(logits,axis=-1),targets)
+      metrics=jax.lax.pmean({"loss":loss,"learning_rate":learning_rate_function(state.step),'accuracy':accuracy},axis_name="batch")
+      return new_state,metrics,new_dropout_rng
+
+  parallel_train_step = jax.pmap(train_step, axis_name="batch", donate_argnums=(0,))
+
+  def eval_step(state, batch):
+      targets=batch.pop('label')
+      logits = state.apply_fn(**batch, params=state.params, train=False)
+      loss=state.loss_function(logits,targets)
+      predictions=state.logits_function(logits)
+      eval_accuracy=jnp.equal(predictions,targets)
+      #eval_acc=jnp.equal(predictions,targets)
+      metrics=jax.lax.pmean({"loss":loss,'accuracy':eval_accuracy},axis_name="batch")
+      #return state.logits_function(logits)  #(8,4)
+      return targets,predictions,metrics
+
+  parallel_eval_step = jax.pmap(eval_step, axis_name="batch")
+
+  def glue_train_data_loader(rng,dataset,batch_size):
+    steps_per_epoch=len_train_dataset//batch_size
+    perms=jax.random.permutation(rng,len(dataset))
+    perms=perms[:steps_per_epoch*batch_size]
+    perms=perms.reshape((steps_per_epoch,batch_size))
+    for perm in perms:
+      batch=dataset[perm]
+      batch={k:jnp.array(v) for k,v in batch.items()}
+      batch=shard(batch)
+      yield batch
+
+  rng=jax.random.PRNGKey(seed)
+  dropout_rngs=jax.random.split(rng,jax.local_device_count())
+
+  def glue_eval_data_loader(dataset, batch_size):
+      for i in range(len_validation_dataset // batch_size):
+          batch = dataset[i * batch_size : (i + 1) * batch_size]
+          batch = {k: jnp.array(v) for k, v in batch.items()}
+          batch = shard(batch)
+
+          yield batch
+
+  state = flax.jax_utils.replicate(state)
+  #metrics_list = list_metrics()
+
+  actual_task = "mnli"
+  metric = load_metric('glue', "mnli")
+  actual_taskmetric = load_metric('glue', actual_task)
+
+  workdir='./results_tensorboard'
+  summary_writer = tensorboard.SummaryWriter(workdir)
+  #summary_writer.hparams(dict(GPT2Config()))
+
+  logger.info(f"***** Running training *****")
+  logger.info(f"  Num examples = {len_train_dataset}")
+  logger.info(f"  Num Epochs = {1}")
+  logger.info(f"  Instantaneous batch size per device = {per_device_batch_size}")
+  logger.info(f"  Total train batch size  = {total_batch_size}")
+  logger.info(f"  Total optimization steps = {num_train_steps}")
+
+  for i, epoch in enumerate(tqdm(range(1, num_train_epochs+1), desc=f"Epoch ...", position=0, leave=True)):
+    rng, input_rng = jax.random.split(rng)
+    train_acc_metrics=[]
+    train_loss_metrics=[]
+    eval_acc_metrics=[]
+    eval_loss_metrics=[]
+    # train
+    with tqdm(total=len_train_dataset // total_batch_size, desc="Training...", leave=False) as progress_bar_train:
+      for idx,batch in enumerate(glue_train_data_loader(input_rng, train_dataset, total_batch_size)):
+        state, train_metric, dropout_rngs = parallel_train_step(state, batch, dropout_rngs)
+        train_acc_metrics.append(jax.device_get(train_metric['accuracy']).mean().item())
+        train_loss_metrics.append(flax.jax_utils.unreplicate(train_metric)['loss'].item())
+        if idx%5==0:
+          summary_writer.scalar('train_loss',flax.jax_utils.unreplicate(train_metric)['loss'].item(),idx)
+          summary_writer.scalar('train_accuracy', jax.device_get(train_metric['accuracy']).mean().item(),idx)
+            
+        progress_bar_train.update(1)
+
+    # evaluate
+    with tqdm(total=len_validation_dataset // total_batch_size, desc="Evaluating...", leave=False) as progress_bar_eval:
+      for idx,batch in enumerate(glue_eval_data_loader(validation_dataset, total_batch_size)):
+          labels,predictions,eval_metric=parallel_eval_step(state, batch)
+          eval_acc_metrics.append(jax.device_get(eval_metric['accuracy']).mean().item())
+          eval_loss_metrics.append(flax.jax_utils.unreplicate(eval_metric)['loss'].item())
+          progress_bar_eval.update(1)
+          if idx%5==0:
+            logger.info(f"eval_step_loss{idx}:   {flax.jax_utils.unreplicate(eval_metric)['loss'].item()} eval_step_acc{idx}:   {jax.device_get(eval_metric['accuracy']).mean().item()}")
+            summary_writer.scalar('eval_loss',flax.jax_utils.unreplicate(eval_metric)['loss'].item(),idx)
+            summary_writer.scalar('eval_accuracy', jax.device_get(eval_metric['accuracy']).mean().item(),idx)
+
+    params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))
+    model.save_pretrained(
+        '.',
+        params=params,
+        push_to_hub=True,
+        commit_message=f"Saving weights of epoch {epoch} at step {idx}",)
+          
+            #correct
+    logger.info(f"---------------------Epoch {epoch} done-----------------")
+    logger.info(f"Train loss:   {jax.device_get(jnp.array(train_loss_metrics)).mean().item()} Train accuracy:  {jax.device_get(jnp.array(train_acc_metrics)).mean().item()}")
+    logger.info(f"Eval loss:    {jax.device_get(jnp.array(eval_loss_metrics)).mean().item()} Eval accuracy:    {jax.device_get(jnp.array(eval_acc_metrics)).mean().item()}")
+  summary_writer.flush()
+
+if __name__ == "__main__":
+    main()