Using and fine-tuning Gemma 3

Key Takeaways

With the release of Gemma 3, we really wanted to significantly increase the user experience for you all. And this is why we are introducing a new uh JMA Python library. Uh I'm going to in this talk I'm going to like demonstrate some of the features that this library provide. The first thing you probably want to do is to just chat with the model. So let's look at an example. Um you start by importing the library. Uh choose uh one of the Gemma architecture, load the parameters. So each of the Gemma models come in two flavor, the pre-trained version or the instruction tuned. And then uh with our sampler interface you can start quering the model about your last u programming problem or in this case this beautiful poem about open weight LLMs. Um and I'm sure you also want to try the multi-turn capability of u our models. So you activate the multi-turn in the sampler and then the model will remember uh the context of your query and automatically handle the caching for you. So it only um only computes the last queries and um similarly with a multimodel um you only need to specify those special uh token inside your prompt and then provide the images uh like the raw images to the chat interface. And uh in this case Jas3 will help me choose where I will spend my next holidays. and um train like fine tune like inference is only part of the problem. Um you also want to be able like because our model are open weight. Uh you want to be able to customize them for your custom use case and your custom data and this is why we um make it as easy as possible for you by integrating Gemma with Cauldron. Cauldron is a open-source library actively used at Google and that is designed to train any arbitrary model on any arbitrary u modalities. If we look at example, I'm sure this API will look very um familiar with a lot of you. You import the library, you create the trainer and then you call the train method um like that will uh launch the training and return the final parameters. Um for the data set side we support uh we're not bound to any data set provider and we support out of the box standard provider like face tensorflow data set um or simpler format like JSON and if you have a custom format or custom data set it's also very easy to create your own wrapper uh to integrate it with our data set. Uh then we have this uh transform API that allow you to apply specific transformation to the individual element of your data set. For example, to apply tokenizing or padding with the GMA library, we provide highle transformation for the common task that you want to use like a supervised fine tuning. So in this case it will take a pair of question and answer from your data set and uh convert them into the input tokens that will be fed to the model and loss. Uh then you can pass this data set to your trainer um and then you connect uh your data set to your model and loss using the special string key. We think that this system is very flexible because then we are not bound to any particular data set structure and um um if for example you want to switch uh to train on multi model data the only thing you need to do is to one add some uh like have a data set that will return images and then with a single line of code you indicate to your model that oh it should also use those image from this data set and the rest of the training is completely identical and uh will work out of the box. For the optimizer uh we use of tags that provide all the standard optimizer um out of the box but also uh have a more advanced capability for like gradient clipping YDK or complex learning rate schedules. So once you have your trainer, you often want to customize specific fields uh from it. And in most programming codebase, what you will do is then you will duplicate the this field inside your config. And then like try to propagate the field uh inside your codebase until the place it is actually used. And uh here we want to remove all indirection and your code. So uh this is why you can just uh add your trainer definition inside a config file and all the field from the trainer will be automatically configurable through the command line. So in this case you indicate um with config you want to launch and then you can override arbitrary fields from this config and this support uh any Python code and any function. You don't need to modify anything about uh uh your class. If you have like a custom loss or custom thing, you can just plug them inside the trainer and things will work uh out of the box. And uh if you go to our GitHub, we provide some of those default trainer configs for the most common use cases like supervised fine tuning, multimodel, uh classification, DPO and more. Uh we also have a bunch of collabs uh that explore more into details some specific uh features of the libraries. Uh and now I'm going to talk about some more advanced features. So if you have access to multiple GPU often it's quite complicated to uh implement the sharding directly in our codebase. But uh with our library when you restore the parameters you just specified uh this sharding uh field uh with some of the default sharding strategy that we provide and uh your model gets automatically uh sharded and similarly during training you can have this sharding uh property to specify the sharding from your data set or optimizer. Yeah. And uh if you don't have access to multiple uh GPU uh you can still try those sharding features uh with our public collab runtime and just choose the uh collab to use TPU v28 that has eight TPU cars and this allow you to try for free or bigger model uh in collab and on the other end of the spectrum we know many of you only have a single GPU maybe with a low memory footprint and we want Gemma to be for everyone. So that's why we also support Lora or Lorank adaptation that is a method uh used to reduce the memory footprint during training. Uh so here you can take any of our GMA model and just wrap it inside our Laura wrapper for the optimizer. You also need to wrap it in the partial update. so that only the Laura rates are optimized and not the full network. And with those two change, everything else work out of the box. Um and finally once your model is trained, uh we have some weight surgery util that allow you to extract the lower rates or fuse them back to the uh model. So I talk about like some of the features of the library but if you go on our website we have more like for example about the quantization or some DPU examples and this is just the beginning. So um very soon we also plan to add like streaming decoding tool use uh model merging so you can fine-tune multiple GMA expert and then uh merge their capability together which was one of the features explained earlier today and more contisation. So if you're interested uh you can check out our GitHub. Thank you. [Music]