Switch to streamlit with markdown, add T5X pre-trained models

README.md

@@ -1,5 +1,5 @@
 ---
-title: Pre-training Dutch T5 Models
+title: Pre-training Dutch T5 Models, evaluation and model lists
 emoji: 🚀
 colorFrom: blue
 colorTo: pink

REMARKS.md

@@ -6,8 +6,6 @@
   intend to run many epochs on the same data, its worth to try a training run without dropout.
   If you want to compare losses, be sure to set the dropout rate equal.
   The smaller models can probably always be trained without.
-* For the translation task, I am not sure that a 'deep-narrow' model (e.g. base-nl36) is better than a normal model
-  or even a 'wide-deep' model.
 * Training with more layers is much slower than you'd expect from the increased model size.
   It is also more difficult to get batch size and learning rate right. Below is a section
   about finding the right hyperparameters for the base-36L training.

app.py

@@ -1,9 +1,43 @@
+from functools import partial
 import time
 
+import sqlite3
 import psutil
 import streamlit as st
-
-IMAGE_WIDTHS=900
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+IMAGE_WIDTHS = 900
+PRE_TRAINED_DB = "data/pretrained.sqlite"
+
+
+@st.cache
+def load_eval_data():
+    conn = sqlite3.connect(PRE_TRAINED_DB)
+    conn.row_factory = lambda c, r: {
+        col[0]: r[idx] for idx, col in enumerate(c.description)
+    }
+    df = pd.read_sql_query("SELECT * FROM pretrained", conn)
+    df.replace("None", np.nan, inplace=True)
+    df.rename(columns={"model": "name"}, inplace=True)
+    df = df.infer_objects()
+    int_columns = ["train_batch_size", "num_parameters"]
+    df[int_columns] = df[int_columns].astype("Int32")
+    plot_df = df[["name", "num_parameters", "summ_rouge1", "trans_en_nl_score"]]
+    plot_df[["num_parameters", "summ_rouge1", "trans_en_nl_score"]] = plot_df[
+        ["num_parameters", "summ_rouge1", "trans_en_nl_score"]
+    ].apply(pd.to_numeric)
+    plot_df["num params (M)"] = plot_df["num_parameters"].map(
+        lambda x: int(x / 10**6)
+    )
+    plot_df.dropna(subset=["summ_rouge1"], inplace=True)
+    plot_df.rename(
+        columns={"summ_rouge1": "summ Rouge1", "trans_en_nl_score": "en->nl Bleu"},
+        inplace=True,
+    )
+    return plot_df
 
 
 def main():
@@ -13,6 +47,7 @@ def main():
         initial_sidebar_state="collapsed",  # Can be "auto", "expanded", "collapsed"
         page_icon="📑",  # String, anything supported by st.image, or None.
     )
+    plot_df = load_eval_data()
 
     with open("style.css") as f:
         st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
@@ -23,9 +58,10 @@ def main():
     with open("PRETRAINING.md", "r") as f:
         st.markdown(f.read())
 
-    st.markdown("""## Evaluation
+    st.markdown(
+        """## Evaluation
 
-### Running evaluation runs
+### Evaluation setup
 
 Each pre-trained model was evaluated by fine-tuning on summarization and translation. The learning-rate was set to
 a constant schedule after a small warmup of 32 steps.
@@ -33,7 +69,7 @@ Fine-tuning for evaluation was done on a limited set of 50K examples from the fi
 
 |                  | Summarization    | Translation       |
 |-----------------:|------------------|-------------------|
-|          Dataset | CNN Dailymail NL | CCMatrix en -> nl |
+|          Dataset | [CNN Dailymail Dutch](https://huggingface.co/datasets/yhavinga/cnn_dailymail_dutch) | [CCMatrix En->NL](https://huggingface.co/datasets/yhavinga/ccmatrix_en_nl) |
 |   #train samples | 50K              | 50K               |
 |        Optimizer | AdamW            | AdamW             |
 |    learning rate | 0.001            | 0.0005            |
@@ -42,78 +78,127 @@ Fine-tuning for evaluation was done on a limited set of 50K examples from the fi
 |    #eval samples | 1000             | 1000              |
 | wandb link      | [eval_summ](https://wandb.ai/yepster/eval_dutch_cnndaily_202302_flax)|[eval_transl](https://wandb.ai/yepster/eval_dutch_ccmatrix_202302_flax) |
 
-#### Summarization
+### Evaluation results
 
-The graph below ([Wandb link](https://api.wandb.ai/links/yepster/hpab7khl)) shows the Rouge1 score for the summarization runs, evaluated
-after 25K and 50K examples on the [CNN Dailymail Dutch](https://huggingface.co/datasets/yhavinga/cnn_dailymail_dutch) dataset:
-""")
-    st.image("eval_summ_rouge1_202302.png", width=IMAGE_WIDTHS)
-    st.markdown("""* Flan models perform almost instantly well on the summarization task, with `flan-t5-small`
+The figure below shows the evaluation scores for most models, with summarization Rouge1 on the x-axis (higher is better),
+and translation English to Dutch Bleu score on the y-axis (higher is better).
+The point size is proportional to the model size. UL2 models are blue, Flan models
+red, mT5 green and the other models black.
+"""
+    )
+    col1, col2 = st.columns(2)
+    with col1:
+        ul2_enabled = st.checkbox("UL2 Dutch (and English) (trained with T5X)", value=True)
+        t5_1_1_enabled = st.checkbox("t5_1_1 Dutch (trained with T5X)", value=True)
+        flan_enabled = st.checkbox("Flan T5 (google/flan-t5-*)", value=True)
+        mt5_enabled = st.checkbox("mt5 (google/mt5-*)", value=True)
+        long_t5_enabled = st.checkbox("Long T5 Dutch+English (trained with HuggingFace script)")
+        t5_v1_1_enabled = st.checkbox("T5 Dutch (and English) (trained with HuggingFace script)")
+    with col2:
+        small_enabled = st.checkbox("small model sizes")
+        base_enabled = st.checkbox("base model sizes")
+        large_enabled = st.checkbox("large model sizes")
+        _24_enabled = st.checkbox("small nl24 deep narrow sizes")
+        _36_enabled = st.checkbox("base nl36 deep narrow sizes")
+        _8l_enabled = st.checkbox("large nl8 deep wide sizes")
+        _4xl_enabled = st.checkbox("xlarge nl4 deep wide sizes")
+
+    plot_df = plot_df[
+        (plot_df["name"].str.contains("ul2") & ul2_enabled)
+        | (plot_df["name"].str.contains("flan") & flan_enabled)
+        | (plot_df["name"].str.contains("mt5") & mt5_enabled)
+        | (plot_df["name"].str.contains("long-t5") & long_t5_enabled)
+        | (plot_df["name"].str.contains("t5_1_1") & t5_1_1_enabled)
+        | ((plot_df["name"].str.startswith("t5") & ~plot_df["name"].str.startswith("t5_1_1")) & t5_v1_1_enabled)
+        | (plot_df["name"].str.contains("base") & base_enabled & ~plot_df["name"].str.contains("36"))
+        | (plot_df["name"].str.contains("small") & small_enabled & ~plot_df["name"].str.contains("24"))
+        | (plot_df["name"].str.contains("large") & large_enabled & ~plot_df["name"].str.contains("8"))
+        | ((plot_df["name"].str.contains("-36L") | plot_df["name"].str.contains("nl36")) & _36_enabled)
+        | ((plot_df["name"].str.contains("-24L") | plot_df["name"].str.contains("nl24")) & _24_enabled)
+        | ((plot_df["name"].str.contains("-8l") | plot_df["name"].str.contains("nl8")) & _8l_enabled)
+        | ((plot_df["name"].str.contains("-4L") | plot_df["name"].str.contains("nl4")) & _4xl_enabled)
+        ]
+
+    color_dict = {"flan": "red", "ul2": "blue", "mt5": "green", "t5_1_1": "orange"}
+    colors = [
+        color_dict[name.split("-")[0].lower()]
+        if name.split("-")[0].lower() in color_dict.keys()
+        else "black"
+        for name in plot_df["name"]
+    ]
+    fig = plt.figure(figsize=(15, 8))
+    sns.set_style("darkgrid")
+    ax = sns.scatterplot(
+        data=plot_df,
+        y="en->nl Bleu",
+        x="summ Rouge1",
+        size="num params (M)",
+        color=colors,
+        linewidth=0.7,
+    )
+    for i, row in plot_df.iterrows():
+        ax.annotate(
+            row["name"],
+            (row["summ Rouge1"], row["en->nl Bleu"]),
+            xytext=(0, 7),
+            textcoords="offset points",
+            ha="center",
+            va="center",
+            rotation=0,
+        )
+    plt.tight_layout()
+    st.pyplot(fig)
+    st.markdown("""* The `UL2` pre-trained Dutch(English) models consistently outperform the `T5-*` Dutch(English) models.
+* Flan models perform almost instantly well on the summarization task, with `flan-t5-small`
   showing performance comparable to Dutch T5 base models.
-* After 50K examples, the `ul2` models exhibit similar performance to the `flan` models.
+* Fine-tuning of `t5-v1.1-large-dutch-cased` failed with the fixed hyperparameters across all models.
+  Since the `UL2` models are better across the board, I've disabled this model on the hub.
 * I am surprised by the consistent bad scores for the `long-t5` runs. I've retried the fine-tuning of these models with
   `float32` instead of `bfloat16`, but the results were the same. Maybe this is normal behaviour for these models
   targeted at dealing with longer sequence lengths.
-
-#### Translation
-
-The graph below ([WandB link](https://wandb.ai/yepster/eval_dutch_ccmatrix_202302_flax/reports/eval_ds_0-score-23-02-11-17-32-48---VmlldzozNTM0NTIy) shows the Bleu score for the translation runs, evaluated at step 25K and
-50K on the [CCMatrix](https://huggingface.co/datasets/yhavinga/ccmatrix_en_nl) dataset, from
-English to Dutch:
-    """)
-    st.image("eval_transl_bleu_202302.png", width=IMAGE_WIDTHS)
-    st.markdown("""* For the translation task from English to Dutch, the Dutch+English pre-trained models perform well. Also
+* For the translation task from English to Dutch, the Dutch+English pre-trained models perform well. Also
   `UL2 Dutch` pre-trained Dutch models are consistently better than their `Flan`, `T5 Dutch` and
   `mT5` counterparts of the comparable size.
-* Like with the summarization task, the `long-t5` models show bad performance, even after 50K examples. I do not understand 
-  cannot explain this at all for this translation task. With a sequence length of 128 input and output
+* For the translation task, I am not sure that a 'deep-narrow' model (e.g. base-nl36) is better than a normal model
+  or even a 'wide-deep' model.
+* The `long-t5` models show bad performance on both tasks.
+  I cannot explain this the translation task. With a sequence length of 128 input and output
   tokens, the sliding attention window with radius length 127 of the `long-t5` models should be able to handle this.
-
-The figure below shows the evaluation scores for most models, with summarization Rouge1 on the x-axis (higher is better),
-and translation English to Dutch Bleu score on the y-axis (higher is better).
-The point size is proportional to the model size. UL2 models are blue, Flan models
-red, mT5 green and the other models black.
 """)
-    st.image("eval_t5_dutch_english.png", width=IMAGE_WIDTHS)
-    st.markdown("""* For clarity, not all models are shown.
-  Among the omitted are `t5-base-36L-dutch-english-cased` with scores comparable to `ul2-large-dutch-english`, but slower inference.
-  The `long-t5` models had such a bad performance that they could not be graphed without cluttering the other models together.
-  Also fine-tuning `t5-v1.1-large-dutch-cased` with the fixed settings for learning rate and batch size diverged.
-* Across the board, for translation the models pre-trained with Dutch+English or Dutch converge faster than other models.
-  I was surprised to see `t5-xl-4l` among the best models on translation, as it has only 4 layers, and previous tests
-  showed that it had a very bad performance (In those tests I had forgot to force set the dropout rate to 0.0, and
-  apparently this model was very sensitive to dropout). 
-    """)
 
     with open("REMARKS.md", "r") as f:
         st.markdown(f.read())
 
-    st.markdown("""### Bfloat16 datatype requires loss regularization
+    st.markdown(
+        """### Bfloat16 datatype requires loss regularization
 
 When training models with `bfloat16` and without loss regularization (default), the training losses would plateau or
 diverge. The graph below displays the results of different attempts
 to train [t5-small-24L-dutch-english](https://huggingface.co/yhavinga/t5-small-24L-dutch-english).
 The legend indicates the optimizer, data type, learning rate, total batch size, and learning rate schedule used.
 As you can see, all attempts to train with `bfloat16` failed.
-""")
-    st.image("bfloat16_loss.png", width=IMAGE_WIDTHS)
-    st.markdown("""The solution was found when peeking at T5X and the T5 gin configs, where I noticed a `z_loss` parameter,
-always set to 1e-4. This factor is used in the [cross entropy loss](https://github.com/google-research/t5x/blob/a319e559b4f72bffab91821487382ef4c25dfcf4/t5x/losses.py#L26)
+"""
+    )
+    st.image("img/bfloat16_loss.png", width=IMAGE_WIDTHS)
+    st.markdown(
+        """The solution was found when peeking at T5X and the T5 gin configs, where I noticed a `z_loss` parameter,
+always set to 1e-4. This factor is used in the T5X [cross entropy loss](https://github.com/google-research/t5x/blob/a319e559b4f72bffab91821487382ef4c25dfcf4/t5x/losses.py#L26)
 function, with the purpose to pull the weights towards zero.
 I experimented with adding this regularization term in the HF pre-training script,
 and the `bfloat16` training runs did not exhibit the problems illustrated above anymore.
 
-When watching Andrej Karpathy [explaining an additional loss term
-called 'regularization loss'](https://youtu.be/PaCmpygFfXo?t=6720) (L2 regularization), it appeared similar
-to the regularization term added in T5X's cross entropy loss function.
+The `z_loss` regularization term in the T5X loss function looks like L2 regularization.
+(See e.g. Andrej Karpathy [explaining regularization loss](https://youtu.be/PaCmpygFfXo?t=6720)).
 The Optax optimizer, used in the HuggingFace script, mentions weight decay for AdaFactor (and AdamW)
 but also mentions that L2 regularization does not work as expected with adaptive gradient
 algorithms. It might be the case that setting a non-zero `weight_decay_rate` in the Optax Adafactor call
 in the HuggingFace pre-training script is an alternative to adding the `z_loss` term, to solve the bfloat16 issues, but
 I haven't tested this yet.
-""")
+"""
+    )
 
-    st.markdown("""### Which optimizer and lr to use
+    st.markdown(
+        """### Which optimizer and lr to use
 
 During the Flax/Jax Community week in '21, our team quickly decided on using Adafactor with learning rate 5e-3.
 I believed that a more optimal setting could be found with more time.
@@ -125,50 +210,64 @@ because the initial version of the training script had the optimizer as a boolea
 changed to a string with the optimizer name.) --
 All runs in the graph below that achieve a loss below 4 use **Adafactor**.
 Peach-sweep-6 is represented by a dashed orange line and had a learning rate of **5e-3**.
-""")
+"""
+    )
 
-    st.image("adafactor_vs_adam_pretrain.png", width=IMAGE_WIDTHS)
-    st.markdown("""While there probably is a setting that will allow Adam and Shampoo to also converge fast below loss 4.0, I was unable
+    st.image("img/adafactor_vs_adam_pretrain.png", width=IMAGE_WIDTHS)
+    st.markdown(
+        """While there probably is a setting that will allow Adam and Shampoo to also converge fast below loss 4.0, I was unable
 to find it. In a recent tweet Lucas Nestler had more success with Shampoo (https://twitter.com/_clashluke/status/1535994026876252160)
 so maybe I need to revisit the attempt with the latest upstream code bases.
 
 Later, when pre-training with T5X, I found that its custom Adafactor implementation with the default settings of the T5X gin configs,
 a learning rate of 0.001 and inverse square root learning rate decay, worked well.
-""")
+"""
+    )
 
-    st.markdown("""### Optimizer and learning rate used for summarization
+    st.markdown(
+        """### Optimizer and learning rate used for summarization
 
 Finetuning summarization requires more memory than translation due to the longer sequence lengths involved.
 I wondered if I could use Adafactor instead of Adam and ran
 a sweep to test this. The sweep was configured with Hyperband, so not all training runs completed to the end.
-""")
-    st.image("optim_lr_summarization.png", width=IMAGE_WIDTHS)
-    st.markdown("""The training losses are graphed below:
-    """)
+"""
+    )
+    st.image("img/optim_lr_summarization.png", width=IMAGE_WIDTHS)
+    st.markdown(
+        """The training losses are graphed below:
+    """
+    )
 
-    st.image("training_losses_summarization_sweep.png", width=IMAGE_WIDTHS)
-    st.markdown("""
+    st.image("img/training_losses_summarization_sweep.png", width=IMAGE_WIDTHS)
+    st.markdown(
+        """
 While the Adafactor run with learning rate 7e-4 came close to the Adam runs, the consistent stability of training with Adam
 made me stick with Adam as optimizer for evaluation runs on the several models. For translation the results were similar, though in the end I needed to configure a lower learning rate for all
 models to converge during fine-tuning. 
-""")
+"""
+    )
 
-    st.markdown("""### Sequence length 512 or 1024
+    st.markdown(
+        """### Sequence length 512 or 1024
 
 The models `t5-v1_1-base-dutch-english-cased` and `t5-v1_1-base-dutch-english-cased-1024` have the same model dimensions,
 but are pre-trained on different sequence lenghts, 512 and 1024 respectively.
 The evaluation loss and accuracy of the models do not look too different. Since training of the 1024 sequence length model was
 very slow and didn't converge a was was very slow, I stopped it early. The figure below shows the evaluation
 loss and accuracy.
-""")
-    st.image("t5v1_1eval_loss_and_accuracy.png", width=IMAGE_WIDTHS)
-    st.markdown("""The 512 sequence length model was trained for 10 epochs of the `small` nl+en config (186B tokens total) and the 1024
+"""
+    )
+    st.image("img/t5v1_1eval_loss_and_accuracy.png", width=IMAGE_WIDTHS)
+    st.markdown(
+        """The 512 sequence length model was trained for 10 epochs of the `small` nl+en config (186B tokens total) and the 1024
 sequence length model about 2 epochs of the `large` nl+en config (100B tokens total). While I expected both models to
 perform similarly on downstream tasks, the 1024 sequence length model has better scores for both
 summarization and translation.
-""")
+"""
+    )
 
-    st.markdown("""## Model lists
+    st.markdown(
+        """## Model lists
 
 ### t5_1_1
 
@@ -302,7 +401,8 @@ of the training. Weights & Biases made it possible to keep track of many trainin
 and orchestrate hyperparameter sweeps with insightful visualizations.
 
 Created by [Yeb Havinga](https://www.linkedin.com/in/yeb-havinga-86530825/)
-""")
+"""
+    )
 
     st.write(
         f"""

requirements.txt

@@ -12,3 +12,5 @@ chex>=0.1.4
 ##jaxlib==0.1.67
 flax>=0.5.3
 sentencepiece
+matplotlib
+seaborn