tinystories wrappers

Analizar mecánicamente los efectos del ajuste fino en las tareas definidas en procedimiento

Este es el código para los experimentos de Tinystories del documento: analizar mecánicamente los efectos del ajuste fino en las tareas definidas de procedimiento. El código se basa en el impresionante https://github.com/karpathy/llama2.c.

Instalación y configuración

Para instalar, clone el repositorio, haga un env/lo que sea, lo que sea, y pip install -r requirements.txt .

Para descargar el conjunto de datos, ejecute python tinystories.py download --data_cache_dir <path_to_data_dir> , con <path_to_data_dir> una ruta donde puede almacenar el conjunto de datos descargado y las versiones tokenizadas de él.

El código genera conjuntos de datos y los toca antes del entrenamiento, utilizando python tinystories pretokenize <args> . Para especificar los diferentes conjuntos de datos y protocolos de capacitación previa a la altura y ajuste fino (por ejemplo, la eliminación y la recuperación), utilizamos diferentes argumentos de filtrado al comando pretokenize .

También hay un comando para entrenar un tokenizador, utilizando python tinystories train_vocab . Esto consta de dos archivos, tok8192.model y tok8192.vocab . Ya he hecho esto, y el tokenizer está en data/tok8192.{model/vocab} . Debe mover estos archivos a <path_to_data_dir> que usará para sus experimentos.

Asegúrese de iniciar sesión en wandb para que sus experimentos se registren allí.

Generación de conjuntos de datos

Como se mencionó anteriormente, cada protocolo de ajuste fino está controlado por un conjunto de datos, que generamos antes de que comience la capacitación. Para cualquier entrenamiento, primero debemos generar los conjuntos de datos correspondientes utilizando python tinystories pretokenize ... :

usage: tinystories.py [-h] [--vocab_size VOCAB_SIZE] [--data_cache_dir DATA_CACHE_DIR] [--filtering FILTERING] [--mix_match MIX_MATCH]
                      [--adversarial_training ADVERSARIAL_TRAINING] [--refusal REFUSAL] [--dataset_name DATASET_NAME]
                      {download,pretokenize,train_vocab}

positional arguments:
  {download,pretokenize,train_vocab}

options:
  -h, --help            show this help message and exit
  --vocab_size VOCAB_SIZE
                        pretokenization vocab size. 0 = use Llama 2 tokenizer.
  --data_cache_dir DATA_CACHE_DIR
                        Adjust data cache dir
  --filtering FILTERING
                        How to filter data
  --mix_match MIX_MATCH
                        How to mix_match sample
  --adversarial_training ADVERSARIAL_TRAINING
                        How to adversarially sample
  --refusal REFUSAL     Which features to refusal-train on
  --dataset_name DATASET_NAME
                        dataset name

Los argumentos filtering , mix_match , adversarial_training y refusal toman un formato específico:

• filtering : --filtering=<FEATURE>=<VALUE>,<FEATURE>!=<VALUE>,...
• mix_match : --mix_match=<FEATURE>:<VALUE1>=<VALUE2>,... reemplaza <VALUE1> con <VALUE2> para <FEATURE> .
• adversarial_training : --adversarial_training=<FEATURE>=<VALUE>:<PROBABILITY>,... agrega <VALUE> a <FEATURE> con <PROBABILITY> probabilidad.
• refusal : --refusal=<FEATURE>=<VALUE>:<PROBABILITY>,... Establece la historia en una negativa y agrega <VALUE> a <FEATURE> con <PROBABILITY> probabilidad.

Generación de un conjunto de datos lleva ~ 10 minutos debido a la tokenización. Si no se pasa un nombre de conjunto de datos, uno se genera a partir de las características proporcionadas y el conjunto de datos se almacena en DATA_CACHE_DIR/<dataset_name> .

Para generar todos los conjuntos de datos necesarios para los experimentos de eliminación y recuperación para una característica dada, use dataset_gen/tw_fs.sh , pero cambie DATA_CACHE_DIR al directorio correcto y ajuste FEATURE_DELETE y FEATURE_IRRELEVANT según sea necesario.

Descargar modelos ya entrenados

Puede descargar los modelos Base, Post-Deletion y Post-Recovery Finuning aquí: https://drive.google.com/file/d/19duv1xv2uw789he4zc5bb5bmknrrug7q/view?usp=sharing

Debe extraer los modelos de ese archivo y luego mirar model_names.md para ver qué modelo es cuál.

El tokenizador para estos modelos está en tokenizers/ . Deberá copiar estos archivos al DATA_CACHE_DIR que está utilizando para los puntos de control del modelo.

Entrenamiento

Para ejecutar el entrenamiento, use python train.py <args> . Puede crear una configuración de capacitación específica copiando un archivo en configs y ajustando los argumentos a sus necesidades, y los argumentos también se pueden anular en la línea de comando, así: así:

python train.py configs/base30M.py --batch_size=32 --lr=1e-4

Los argumentos deben usar = y no espacios.

Esta configuración le permite crear un archivo de configuración para un conjunto de experimentos o un barrido, y luego solo pasar los argumentos que varían con los experimentos (IE Base Model, Tasa de aprendizaje, conjunto de datos, etc.) en la línea de comandos.

Todos estos comandos registrarán automáticamente estadísticas y generaciones para wandb , así que asegúrese de que haya iniciado sesión en eso.

Comando de ejemplo de previación

python train.py configs/base91M.py --batch_size=256 --dropout=0.1 --gradient_accumulation_steps=2 --learning_rate=0.0001 --max_iters=100000 --max_seq_len=512

Comando de ejemplo de ajuste

Supresión:

python train.py configs/base91M_tw_delete_no_lrs.py --dataset_name=filter-adv-Twist --learning_rate=1e-05

Recuperación:

python train.py configs/base91M_tw_recovery_no_lrs.py --model_dir=/cache/tinystories/base91M-train-2023_10_06_15_15_49_074/out

Barridos

He estado ejecutando experimentos como barridos wandb . Puede ver configuraciones de barrido de ejemplo para la eliminación y la recuperación en sweeps/*.yaml . Cree un barrido con wandb sweep -p tinystories-uft -e ucl-dar sweeps/<sweep.yaml> , y luego, en cualquier script de envío que use, ejecute wandb agent <sweep_id> --count 1 (necesita --count 1 si está utilizando slurm/etc., ya que cada trabajo enviado ejecuta un solo experimento y luego termina).

Breve guía de entrenamiento .

Vea el script Train.py para obtener más lanzamientos exóticos y anulaciones de hiperparámetro. Aquí hay una breve guía sobre cómo establecer los parámetros. Mire la tabla al final del papel de chinchilla para tener una idea de cómo los parámetros del transformador (DIM, N_LAYERS, N_HEADS) crecen o se encogen juntos. Extrapolar/interpolar este patrón para obtener transformadores más grandes o más pequeños. Establezca la longitud de contexto máximo como desee, dependiendo del problema: este debería ser el número máximo de tokens que importan para predecir el siguiente token. Por ejemplo, Llama 2 usa 2048. A continuación, desea el tamaño total de lotes por actualización (impreso por el script como "Tokens por iteración será:") para estar en algún lugar alrededor de 100k tokens para aplicaciones medianas. Para aplicaciones pequeñas, podría ser menor, para una gran capacitación (por ejemplo, GPTS/Llamas) generalmente es ~ 0.5m, o incluso más. Usted llega allí al máximo maxing the Batch_Size a lo que su sistema permita (por ejemplo, el mío fue 16 en una ejecución reciente porque después de eso, mi GPU se queda fuera de la memoria), y luego desea aumentar el gradiente_accumulation_steps para que sea tan alto como sea necesario para alcanzar el tamaño total de lotes de ~ 100k. Finalmente, desea ajustar su aprendizaje_rate (LR). Desea que esto sea tan alto como su entrenamiento lo permita. Las redes muy pequeñas pueden escapar con una LR grande (por ejemplo, 1E-3 o incluso más). Las redes grandes necesitan LRS más bajos. 3E-4 es una opción segura en la mayoría de las aplicaciones de tamaño mediano, pero puede ser demasiado bajo para pequeñas redes, ¡así que trate de aumentarlo! Finalmente, Max_iters es la duración del entrenamiento. Juega con diferentes configuraciones. En su mayoría, solo sintonizo estos parámetros y dejo que la mayoría de los demás sin cambios. Aquí hay un ejemplo de cómo entrené el modelo de 110 m, que no creo que sea casi óptimo, pero me pareció sensato: Dim 768, N_LAYERS 12, N_heads 12 (por lo que el tamaño de cada cabezal es 768/12 = 64 canales), seq pro de 1024, tamaño de lote 16 (esto es lo más que se ajusta a mi A100 40GB GPU), Gradiente_CACCUMA_CUCTUMATURACIÓN = 8 fue para obtener el total de mi A100 40GB), Gradiente_CACUMA_CUCUMATURACIÓN = 8 fue para obtener el total de mi A100 40GB), Gradiente_CACCUMA_CUCUMATURACIÓN = 8. El tamaño del lote de los tokens será de 16 lotes de 16 tokens en secuencia * 8 grad_accum = 131,072 tokens por actualización. Bien. Tasa de aprendizaje 4E-4 (probablemente un poco demasiado baja). Max_iters 200k (probablemente un poco demasiado alto). Descarga 0.1, ya que eso generalmente ayuda un poco a tamaño mediano. Eso fue todo. Casé usando datos distribuidos paralelos (DDP) en 4 GPU en mi máquina en la nube, la capacitación tomó ~ día más o menos.

Licencia

MIT

Citación

Si usa este código, cite nuestro trabajo:

 @misc{jain2023mechanistically,
    title={Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks},
    author={Samyak Jain and Robert Kirk and Ekdeep Singh Lubana and Robert P. Dick and Hidenori Tanaka and Edward Grefenstette and Tim Rocktäschel and David Scott Krueger},
    year={2023},
    eprint={2311.12786},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}

También debes citar el papel de Tinystories:

 @misc{eldan2023tinystories,
    title={TinyStories: How Small Can Language Models Be and Still Speak Coherent English?},
    author={Ronen Eldan and Yuanzhi Li},
    year={2023},
    eprint={2305.07759},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}