| # torchtune research repo: token coloring (colorful llama) | |
| Playground to try out [token coloring](https://docs.google.com/document/d/1Win9vhddD-pu5P3SsG7E-dzN5oQl5DYWW1DhO7sBOgI/edit#heading=h.oqq00pt8expe) with TorchTune. | |
| The repo was generated using the alpha version of [torchtune](https://github.com/pytorch-labs/torchtune). | |
| Brief notes: | |
| - The starting recipe is based on the Alpaca Llama2 7B full finetune recipe (switched to bf16). | |
| - I assume `output/` is used to store model outputs and `model/` is used to store the base model checkpoints. | |
| For the `colorful` recipe: | |
| - I copied a lot of functionality (like the actual model definition, dataset, etc) from torchtune repository directly since I needed to make changes. | |
| - I reduced the flexiblity of the recipe (e.g. cannot specify the model or tokenizer) and increased it in other ways (e.g. can pass in a dataset path directly). | |
| - I added intermediate checkpointing (i.e. every `n` steps) and automatically upload the checkpoint to HuggingFace Hub. | |
| ## Getting started | |
| The below instructions can be copy-pasted as is on to a running instance. They assume that the `HF_TOKEN` environment variable is set with a valid token. | |
| ```bash | |
| # for RunPod | |
| cd /workspace | |
| git clone git@github.com:pytorch-labs/torchtune.git | |
| cd torchtune | |
| pip install -e . | |
| cd /workspace | |
| git clone git@github.com:laurencer/torchtune-colorful-llama.git | |
| cd torchtune-colorful-llama | |
| # for wandb support | |
| pip install wandb | |
| ``` | |
| ```bash | |
| mkdir -p model/ | |
| tune download --repo-id meta-llama/Llama-2-7b --output-dir model/ | |
| ``` | |
| ```bash | |
| tune convert_checkpoint --checkpoint-path model/consolidated.00.pth --output-path model/llama2_native.tune | |
| ``` | |
| ```bash | |
| mkdir -p output/ | |
| # tune --nnodes 1 --nproc_per_node 1 ./colorful/full_finetune.py --config ./colorful/basic_config.yaml | |
| nohup tune --nnodes 1 --nproc_per_node 1 ./colorful/full_finetune.py --config ./colorful/basic_config.yaml 2>&1 > training_log_$(date "+%Y.%m.%d_%H.%M.%S").log & | |
| sleep 1 | |
| tail -f training_log_*.log | |
| ``` | |
| ## Baselines | |
| Two baseline configs are provided in the `baseline` directory. | |
| We forked the original recipe to support customizing the location/path of the Alpaca dataset. | |
| ```bash | |
| # tune --nnodes 1 --nproc_per_node 1 ./baseline/full_finetune.py --config ./baseline/baseline_config.yaml | |
| nohup tune --nnodes 1 --nproc_per_node 1 ./baseline/full_finetune.py --config ./baseline/baseline_config.yaml 2>&1 > training_log_$(date "+%Y.%m.%d_%H.%M.%S").log & | |
| sleep 1 | |
| tail -f training_log_*.log | |
| ``` | |
| The adversarial config uses a dataset that is equivalent to 4x the original alpaca cleaned dataset with extra examples that include prompt injection attempts. See [token coloring description](https://docs.google.com/document/d/1Win9vhddD-pu5P3SsG7E-dzN5oQl5DYWW1DhO7sBOgI/edit#heading=h.oqq00pt8expe) for more info. | |
| ```bash | |
| # tune --nnodes 1 --nproc_per_node 1 ./baseline/full_finetune.py --config ./baseline/adversarial_config.yaml | |
| nohup tune --nnodes 1 --nproc_per_node 1 ./baseline/full_finetune.py --config ./baseline/adversarial_config.yaml 2>&1 > training_log_$(date "+%Y.%m.%d_%H.%M.%S").log & | |
| sleep 1 | |
| tail -f training_log_*.log | |
| ``` | |
| ## Colorful | |
| The `colorful` directory implements the changes required to support token coloring. This includes a custom dataset implementation and training script. |