Add Tunix demo in miniGPT #269
Conversation
Unfortunately Qwix has issues with both the new sharding annotation and KV caching previously added; so I had to revert both
|
@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 |
| id: LS9sQEY3n0mB | ||
| outputId: 9ffcf3a6-20ef-4f80-b006-f5d3c5644a15 | ||
| tags: [nbval-ignore-output] | ||
| outputId: 6b9ee4b0-eed0-4bae-dd99-ffed14289ad7 |
There was a problem hiding this comment.
It looks like the nbval-ignore-output and nbval-skip tags got deleted, that might be why the tests are failing.
| :id: z0p-IHurrB9i | ||
| # Define a triangular mask for causal attention with `jax.numpy.tril` and `jax.numpy.ones`. | ||
There was a problem hiding this comment.
Why remove the comments in this function?
|
|
||
| ```{code-cell} | ||
| :cellView: form | ||
| :tags: [hide-cell] |
| generated_text = model.generate_text(maxlen, start_tokens) | ||
| step += 1 | ||
| # break |
| ## Tunix: Fine-tuning | ||
|
|
||
| [REVAMP 6] Introduce Tunix. | ||
| [Tunix](https://github.com/google/tunix) is a JAX-native LLM post-training library open sourced by Google. It supports a range of post-training techniques including supervised finetuning, preference tuning, reinforcement learning and model distillation. So in this section, we are going to use Tunix to finetune the miniGPT model we just pretrained using LoRA ([Low-Rank Adaptation](https://arxiv.org/abs/2106.09685)) so that the finetuned model generates output of a different style. |
There was a problem hiding this comment.
In the second sentence, could you remove "So" at the beginning, so it's just "In this section, ..."?
|
|
||
| +++ {"id": "mfnj-S02_yM5"} | ||
|
|
||
| Previously we used Grain to load the Tiny Stories dataset. JAX is actually flexible to work with other dataloaders as well. So for LoRA fintuning we use TensorFlow Datasets to load the [Tiny Shakespeare](https://www.tensorflow.org/datasets/catalog/tiny_shakespeare) dataset. |
There was a problem hiding this comment.
We're thinking of removing TFDS from the stack and only including Grain, now that Grain is becoming more mature and TFDS is losing market share. I'd kind of like to stick with just Grain for the tutorial, wdyt?
|
|
||
| +++ {"id": "qvxSxbbcBim7"} | ||
|
|
||
| We define a few helper functions to create the LoRA model, loss and etc. |
There was a problem hiding this comment.
Remove "and," so it reads "LoRA model, loss, etc."?
|
|
||
| +++ {"id": "UWVdZGK9COj0"} | ||
|
|
||
| After the finetuning, you can easily see that now the model produces text of a different style, kind of like Shakespeare's work, which means our finetuning works. |
There was a problem hiding this comment.
Should this show output?
2 participants
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.