Add Tunix demo in miniGPT

[windmaple]

Unfortunately Qwix has issues with both the new sharding annotation and KV caching previously added; so I had to revert both.

I also had to re-run the notebook (so that we can get Tunix output), which introduced a lot of metadata and output changes.

[windmaple]

@emilyfertig do you know why the test keeps failing? The notebook runs fine on Colab and these libs are properly installed, so the errors don't make sense to me.

NameError: name 'peft_trainer' is not defined
NameError: name 'tfds' is not defined
ModuleNotFoundError: No module named 'tqdm'

[windmaple]

@emilyfertig can you help me with the 'Unrun reference cell has outputs' errors? I'm a bit lost here. Let's take the first example:

______________ docs/source/JAX_for_LLM_pretraining.ipynb::Cell 1 _______________
Notebook cell execution failed
Cell 1: Unrun reference cell has outputs

Input:
import jax
jax.devices()

We have the nbval-ignore-output tag for this cell. But according to this error, we are skipping this cell but there is a saved ouput, so the error is thrown. So I'm not sure what is happening here.

[emilyfertig]

I think 'Unrun reference cell has outputs' is because the "execution count" says "null," and the execution count has to be a numeric value (even 0) to resolve the error. I saw that before and it stumped me for a while...

[emilyfertig]

Thank you!

[windmaple]

Thanks! Finally cleared all errors