Model Preparation

We aim to simplify the model structure as much as possible, avoiding the definition of new wrapper structures to maintain the original LLM's loading method for ease of subsequent use. Therefore, we need to expand the vocabulary of the original LLM (using Gemma as an example). Note that no new tokens need to be added to the tokenizer because the image tokens are generated by VQGAN.

Important: To indicate when the model should start generating images, we add two special tokens: <boi> and <eoi> to represent the beginning and end of image tokens, respectively. Additionally, an <unconditional> token is added to assist in CFG (Classifier-Free Guidance) training. You can directly replace these three special tokens with some <unusedxxx> tokens or manually modify the tokenizer files (since tokenizers currently do not support in-place operations, refer to this issue).

Steps:

1. Download the Gemma 7B model (or 2B version).
2. Locate the model folder containing tokenizer_config.json and tokenizer.json files.
3. Manually edit the tokenizer.json file (note: two changes are required in tokenizer.json: in "added_tokens" and in "vocab") and the tokenizer_config.json file (one change in "added_tokens_decoder").
4. In both tokenizer.json and tokenizer_config.json, replace <unused0> with <boi>, <unused1> with <eoi>, and <unused2> with <unconditional>. Set "special": true for these three tokens. Refer to our released model for example.
5. Run the following command to expand the vocabulary:

   python liquid/expand_vocabulary.py --model_path '/path/to/gemma-7b' --save_path '/path/to/save/gemma-7b-addtoken'
   

Note: When expanding the vocabulary, setting mean_resizing=True is particularly useful for increasing the size of embeddings in causal language models. This ensures that the probabilities of generated tokens are not affected by the added embeddings, as initializing the new embeddings with the old embeddings' mean reduces the KL-divergence between the next token probability before and after adding the new embeddings. Using this parameter results in a lower initial loss after expanding the vocabulary, but its impact on final performance is still unclear in Liquid. This operation requires transformers>=4.46.0.

---

Pre-training

We use torch==2.1.0, transformers==4.39.2, and flash-attention==2.5.8 to train the Gemma 7B model. However, due to our minimalist training framework, newer versions of these libraries should also work.

Environment Setup:

pip install --upgrade pip  
pip install -e .
pip install -e ".[train]"
pip install flash-attn --no-build-isolation

Pretraining:

During the mixed pretraining process, we maintain a global batch size of 1024 (num_GPUs * per_device_train_batch_size * gradient_accumulation_steps).

bash scripts/gemma7b_mixpretrain.sh 

Visual Understanding SFT:

(WIP, coming soon)