Allenamento di CodeParrot 🦜 da zero

'Allenamento di CodeParrot da zero' (Training CodeParrot from scratch)

In questo post del blog daremo un’occhiata a cosa serve per creare la tecnologia dietro GitHub CoPilot, un’applicazione che fornisce suggerimenti ai programmatori durante la scrittura del codice. In questa guida passo passo, impareremo come addestrare un grande modello GPT-2 chiamato CodeParrot 🦜, completamente da zero. CodeParrot può completare automaticamente il tuo codice Python: provaci qui. Cominciamo a costruirlo da zero!

Creazione di un grande set di dati di codice sorgente

La prima cosa di cui abbiamo bisogno è un grande set di dati di addestramento. Con l’obiettivo di addestrare un modello di generazione di codice Python, abbiamo acceduto al dump di GitHub disponibile su BigQuery di Google e filtrato tutti i file Python. Il risultato è un set di dati di 180 GB con 20 milioni di file (disponibile qui). Dopo i primi esperimenti di addestramento, abbiamo scoperto che i duplicati nel set di dati influivano negativamente sulle prestazioni del modello. Approfondendo l’analisi del set di dati, abbiamo scoperto che:

• Lo 0,1% dei file unici costituisce il 15% di tutti i file
• Il 1% dei file unici costituisce il 35% di tutti i file
• Il 10% dei file unici costituisce il 66% di tutti i file

Puoi scoprire di più sulle nostre scoperte in questo thread di Twitter. Abbiamo rimosso i duplicati e applicato le stesse euristiche di pulizia trovate nel paper Codex. Codex è il modello dietro CoPilot ed è un modello GPT-3 addestrato sul codice di GitHub.

• Iniziare con l’Analisi del Sentimento utilizzando Python
• Guidare la generazione del testo con la ricerca vincolata a fascio in 🤗 Transformers
• Un’introduzione al Deep Reinforcement Learning

Il set di dati pulito è ancora di 50GB ed è disponibile su Hugging Face Hub: codeparrot-clean. Con questo possiamo configurare un nuovo tokenizer e addestrare un modello.

Inizializzazione del tokenizer e del modello

Prima di tutto abbiamo bisogno di un tokenizer. Addestriamone uno specificamente sul codice in modo che divida correttamente i token del codice. Possiamo prendere un tokenizer esistente (ad es. GPT-2) e addestrarlo direttamente sul nostro set di dati con il metodo train_new_from_iterator(). Poi lo spingiamo su Hub. Nota che omettiamo gli import, l’analisi degli argomenti e il logging dagli esempi di codice per mantenere i blocchi di codice compatti. Ma troverai il codice completo, compresa la pre-elaborazione e la valutazione del task a valle, qui.

# Iteratore per l'addestramento
def batch_iterator(batch_size=10):
    for _ in tqdm(range(0, args.n_examples, batch_size)):
        yield [next(iter_dataset)["content"] for _ in range(batch_size)]

# Tokenizer di base
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
base_vocab = list(bytes_to_unicode().values())

# Carica il set di dati
dataset = load_dataset("lvwerra/codeparrot-clean", split="train", streaming=True)
iter_dataset = iter(dataset)

# Addestramento e salvataggio
new_tokenizer = tokenizer.train_new_from_iterator(batch_iterator(),
                                                  vocab_size=args.vocab_size,
                                                  initial_alphabet=base_vocab)
new_tokenizer.save_pretrained(args.tokenizer_name, push_to_hub=args.push_to_hub)

Scopri di più sui tokenizers e su come costruirli nel corso di Hugging Face.

Vedi quell’argomento apparentemente insignificante streaming=True? Questo piccolo cambiamento ha un grande impatto: invece di scaricare l’intero set di dati (50GB), verranno trasmesse in streaming singole campioni quando necessario, risparmiando molto spazio su disco! Consulta il corso di Hugging Face per ulteriori informazioni sullo streaming.

Ora, inizializziamo un nuovo modello. Utilizzeremo gli stessi iperparametri di GPT-2 large (1,5 miliardi di parametri) e regoleremo lo strato di embedding per adattarlo al nostro nuovo tokenizer, aggiungendo anche alcune modifiche per aumentare la stabilità. Il flag scale_attn_by_layer_idx garantisce che la scala dell’attenzione venga regolata in base all’id dello strato e reorder_and_upcast_attn garantisce principalmente che il calcolo dell’attenzione venga eseguito con piena precisione per evitare problemi numerici. Spingiamo il modello appena inizializzato nello stesso repository del tokenizer.

# Carica il tokenizer di codeparrot addestrato per la tokenizzazione del codice Python
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name)

# Configurazione
config_kwargs = {"vocab_size": len(tokenizer),
                 "scale_attn_by_layer_idx": True,
                 "reorder_and_upcast_attn": True}

# Carica il modello con la configurazione e spingilo su Hub
config = AutoConfig.from_pretrained('gpt2-large', **config_kwargs)
model = AutoModelForCausalLM.from_config(config)
model.save_pretrained(args.model_name, push_to_hub=args.push_to_hub)

Adesso che abbiamo un tokenizer efficiente e un modello appena inizializzato, possiamo iniziare con il ciclo di addestramento effettivo.

Implementazione del ciclo di addestramento

Addestriamo con la libreria 🤗 Accelerate che ci permette di scalare l’addestramento dal nostro laptop a una macchina multi-GPU senza cambiare una sola riga di codice. Creiamo semplicemente un acceleratore e gestiamo alcuni argomenti:

accelerator = Accelerator()
acc_state = {str(k): str(v) for k, v in accelerator.state.__dict__.items()}

parser = HfArgumentParser(TrainingArguments)
args = parser.parse_args()
args = Namespace(**vars(args), **acc_state)
samples_per_step = accelerator.state.num_processes * args.train_batch_size
set_seed(args.seed)

Siamo ora pronti per l’addestramento! Utilizziamo la libreria client huggingface_hub per clonare il repository con il nuovo tokenizer e il modello. Creeremo un nuovo branch per questo esperimento. Con questa configurazione, possiamo eseguire molti esperimenti in parallelo e alla fine unire il migliore nel branch principale.

# Clonare il repository del modello
if accelerator.is_main_process:
    hf_repo = Repository(args.save_dir, clone_from=args.model_ckpt)

# Checkout del nuovo branch sul repository
if accelerator.is_main_process:
    hf_repo.git_checkout(run_name, create_branch_ok=True)

Possiamo caricare direttamente il tokenizer e il modello dal repository locale. Poiché stiamo gestendo modelli di grandi dimensioni, potremmo voler attivare il checkpointing del gradiente per ridurre l’occupazione di memoria della GPU durante l’addestramento.

# Caricare il modello e il tokenizer
model = AutoModelForCausalLM.from_pretrained(args.save_dir)
if args.gradient_checkpointing:
    model.gradient_checkpointing_enable()
tokenizer = AutoTokenizer.from_pretrained(args.save_dir)

Passiamo ora al dataset. Rendiamo l’addestramento più semplice con un dataset che genera esempi con una dimensione di contesto fissa. Per non sprecare troppi dati (alcuni campioni sono troppo corti o troppo lunghi), possiamo concatenare molti esempi con un token EOS e quindi suddividerli.

Più sequenze prepariamo insieme, minore è la frazione di token che scartiamo (quelli grigi nella figura precedente). Poiché vogliamo trasmettere il dataset invece di preparare tutto in anticipo, utilizziamo un IterableDataset. La classe completa del dataset è la seguente:

class ConstantLengthDataset(IterableDataset):
    def __init__(
        self, tokenizer, dataset, infinite=False, seq_length=1024, num_of_sequences=1024, chars_per_token=3.6
    ):
        self.tokenizer = tokenizer
        self.concat_token_id = tokenizer.bos_token_id
        self.dataset = dataset
        self.seq_length = seq_length
        self.input_characters = seq_length * chars_per_token * num_of_sequences
        self.epoch = 0
        self.infinite = infinite

    def __iter__(self):
        iterator = iter(self.dataset)
        more_examples = True
        while more_examples:
            buffer, buffer_len = [], 0
            while True:
                if buffer_len >= self.input_characters:
                    break
                try:
                    buffer.append(next(iterator)["content"])
                    buffer_len += len(buffer[-1])
                except StopIteration:
                    if self.infinite:
                        iterator = iter(self.dataset)
                        self.epoch += 1
                        logger.info(f"Epoch del dataset: {self.epoch}")
                    else:
                        more_examples = False
                        break
            tokenized_inputs = self.tokenizer(buffer, truncation=False)["input_ids"]
            all_token_ids = []
            for tokenized_input in tokenized_inputs:
                all_token_ids.extend(tokenized_input + [self.concat_token_id])
            for i in range(0, len(all_token_ids), self.seq_length):
                input_ids = all_token_ids[i : i + self.seq_length]
                if len(input_ids) == self.seq_length:
                    yield torch.tensor(input_ids)

I testi nel buffer vengono tokenizzati in parallelo e quindi concatenati. I campioni suddivisi vengono poi restituiti finché il buffer non è vuoto e il processo ricomincia. Se impostiamo infinite=True, l’iteratore del dataset ricomincia dalla fine.

def create_dataloaders(args):
    ds_kwargs = {"streaming": True}
    train_data = load_dataset(args.dataset_name_train, split="train", streaming=True)
    train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed)
    valid_data = load_dataset(args.dataset_name_valid, split="train", streaming=True)
    
    train_dataset = ConstantLengthDataset(tokenizer, train_data, infinite=True, seq_length=args.seq_length)
    valid_dataset = ConstantLengthDataset(tokenizer, valid_data, infinite=False, seq_length=args.seq_length)
    
    train_dataloader = DataLoader(train_dataset, batch_size=args.train_batch_size)
    eval_dataloader = DataLoader(valid_dataset, batch_size=args.valid_batch_size)
    return train_dataloader, eval_dataloader

train_dataloader, eval_dataloader = create_dataloaders(args)

Prima di iniziare l’addestramento, dobbiamo configurare l’ottimizzatore e il programma di tassi di apprendimento. Non vogliamo applicare la decadimento del peso ai bias e ai pesi di LayerNorm, quindi utilizziamo una funzione di supporto per escluderli.

def get_grouped_params(model, args, no_decay=["bias", "LayerNorm.weight"]):
    params_with_wd, params_without_wd = [], []
    for n, p in model.named_parameters():
        if any(nd in n for nd in no_decay): params_without_wd.append(p)
        else: params_with_wd.append(p)
    return [{"params": params_with_wd, "weight_decay": args.weight_decay},
            {"params": params_without_wd, "weight_decay": 0.0},]

optimizer = AdamW(get_grouped_params(model, args), lr=args.learning_rate)
lr_scheduler = get_scheduler(name=args.lr_scheduler_type, optimizer=optimizer,
                             num_warmup_steps=args.num_warmup_steps,
                             num_training_steps=args.max_train_steps,)

Una grande domanda che rimane è come i dati e i modelli verranno distribuiti su più GPU. Questo sembra un compito complesso, ma in realtà richiede solo una singola riga di codice con 🤗 Accelerate.

model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader)

All’interno, verrà utilizzato DistributedDataParallel, il che significa che un batch viene inviato a ciascun worker GPU che ha la propria copia del modello. Lì vengono calcolati i gradienti e quindi aggregati per aggiornare il modello su ciascun worker.

Vogliamo anche valutare il modello di tanto in tanto sul set di validazione, quindi scriviamo una funzione per fare proprio questo. Ciò viene fatto automaticamente in modo distribuito e dobbiamo solo raccogliere tutte le perdite dai worker. Vogliamo anche riportare la perplessità.

def evaluate(args):
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(batch, labels=batch)
        loss = outputs.loss.repeat(args.valid_batch_size)
        losses.append(accelerator.gather(loss))
        if args.max_eval_steps > 0 and step >= args.max_eval_steps:
            break
    loss = torch.mean(torch.cat(losses))
    try:
        perplexity = torch.exp(loss)
    except OverflowError:
        perplexity = float("inf")
    return loss.item(), perplexity.item()

Siamo pronti a scrivere il ciclo di addestramento principale. Assomiglierà molto a un normale ciclo di addestramento PyTorch. Qui e là puoi vedere che utilizziamo le funzioni dell’acceleratore anziché PyTorch nativo. Inoltre, spingiamo il modello nel ramo dopo ogni valutazione.

# Addestra il modello
model.train()
completed_steps = 0
for step, batch in enumerate(train_dataloader, start=1):
    loss = model(batch, labels=batch, use_cache=False).loss
    loss = loss / args.gradient_accumulation_steps
    accelerator.backward(loss)
    if step % args.gradient_accumulation_steps == 0:
        accelerator.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        completed_steps += 1
    if step % args.save_checkpoint_steps == 0:
        eval_loss, perplexity = evaluate(args)
        accelerator.wait_for_everyone()
        unwrapped_model = accelerator.unwrap_model(model)
        unwrapped_model.save_pretrained(args.save_dir, save_function=accelerator.save)
        if accelerator.is_main_process:
            hf_repo.push_to_hub(commit_message=f"step {step}")
        model.train()
    if completed_steps >= args.max_train_steps:
        break

Quando chiamiamo wait_for_everyone() e unwrap_model() ci assicuriamo che tutti i worker siano pronti e che tutti i livelli del modello aggiunti in precedenza da prepare() vengano rimossi. Utilizziamo anche l’accumulo del gradiente e il clipping del gradiente che sono facilmente implementati. Infine, dopo che l’addestramento è completo, eseguiamo un’ultima valutazione, salviamo il modello finale e lo spingiamo nel repository.

# Valuta e salva l'ultimo checkpoint
logger.info("Valutazione e salvataggio del modello dopo l'addestramento")
eval_loss, perplexity = evaluate(args)
log_metrics(step, {"loss/eval": eval_loss, "perplexity": perplexity})
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(args.save_dir, save_function=accelerator.save)
if accelerator.is_main_process:
    hf_repo.push_to_hub(commit_message="modello finale")

Fatto! Questo è tutto il codice per addestrare un modello GPT-2 completo da zero con soli 150 righe. Non abbiamo mostrato gli import e i log degli script per rendere il codice un po’ più compatto. Ora vediamo come addestrarlo effettivamente!

Con questo codice abbiamo addestrato modelli per il nostro prossimo libro su Transformers e NLP: un modello GPT-2 con parametri 110M e 1.5B. Abbiamo utilizzato una macchina con 16 GPU A100 per addestrare questi modelli rispettivamente per 1 giorno e 1 settimana. Abbastanza tempo per prendere un caffè e leggere uno o due libri!

Valutazione

Questo è ancora un tempo di addestramento relativamente breve per il pretraining, ma possiamo già osservare una buona performance in downstream rispetto a modelli simili. Abbiamo valutato i modelli sul benchmark HumanEval di OpenAI, introdotto nel paper Codex. Misura la performance dei modelli di generazione di codice su quasi 200 sfide di codifica.

Da notare che abbiamo addestrato CodeParrot su circa 25-30B di token, mentre GPT-neo è stato addestrato su 300B di token e Codex su 300B (checkpoint di GPT-3) + 100B (codice) di token. Quindi ottieni sicuramente un buon rapporto qualità-prezzo qui.

Ma cos’è esattamente questa metrica pass@k? Semplicemente misura la probabilità che almeno un programma superi il test di unità di un problema di programmazione, dato k generazioni candidate dal modello. Naturalmente, la metrica aumenta con k poiché più candidati possono potenzialmente risolvere la sfida di codifica. Puoi utilizzare questa metrica con datasets :

from datasets import load_metric

code_eval = datasets.load_metric("code_eval")
test_cases = ["assert add(2,3)==5"]
candidates = [["def add(a,b): return a*b", "def add(a, b): return a+b"]]
pass_at_k, results = code_eval.compute(references=test_cases, predictions=candidates, k=[1, 2])

print(pass_at_k)

# -> {'pass@1': 0.5, 'pass@2': 1.0}

Puoi anche caricare il dataset HumanEval di OpenAI con datasets :

from datasets import load_dataset

dataset = load_dataset("openai_humaneval")

Adesso diamo un’occhiata a cosa otteniamo dopo aver aspettato così a lungo con ansia!

Risultati

Iniziamo piano e controlliamo se possiamo far scrivere al modello una funzione per restituire la dimensione di un file dato il suo percorso:

Prompt:

def get_files_size(filename):

Completamento:

    """
    Restituisce la dimensione del file.

    :param filepath: Il percorso del file.
    :type filepath: str
    :return: La dimensione del file.
    :rtype: int
    """
    return os.path.getsize(filepath)

Fino ad ora tutto bene, ma andiamo oltre e proviamo qualcosa di più interessante. Uno degli aspetti meno eccitanti nella creazione di codice pronto per la produzione è scrivere i test di unità. Vediamo se possiamo usare CodeParrot per farlo al posto nostro!

Prompt:

def is_even(value):
    """Returns True if value is an even number."""
    return value % 2 == 0

# setup unit tests for is_even
import unittest

Completamento:

class TestIsEven(unittest.TestCase):
    """Tests for is_even function."""

    def test_is_even(self):
        """Test is_even function."""
        self.assertTrue(is_even(0))
        self.assertTrue(is_even(1))
        self.assertTrue 

Anche se la condizione del test in sé non sembra del tutto corretta, il modello ha costruito tutto il codice di base per noi. Questo da solo potrebbe già essere molto utile e risparmiarci qualche minuto di tempo!

Ma quello che vogliamo davvero fare è costruire modelli di transformer migliori. Usiamo CodeParrot per aiutarci a farlo:

Prompt:

from transformers import AutoTokenizer, AutoModelForSequenceClassification

# costruisci un classificatore BERT

Completamento:

tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')

Osservare alcuni esempi è bello ma probabilmente ora vuoi provarlo tu stesso. Abbiamo preparato alcune demo online proprio per questo!

Demo

Puoi giocare con i modelli in due demo che abbiamo aggiunto a Hugging Face Spaces . Con la prima puoi generare rapidamente del codice con il modello e con la seconda puoi evidenziare il tuo codice con il modello per individuare gli errori!

• Generazione di codice
• Evidenziazione del codice

Puoi anche utilizzare direttamente i modelli dalla libreria transformers:

from transformers import pipeline

pipe = pipeline('text-generation', model='lvwerra/codeparrot')
pipe('def hello_world():')

Sommario

In questo breve post del blog abbiamo esaminato tutti i passaggi necessari per addestrare un grande modello GPT-2 chiamato CodeParrot 🦜 per la generazione di codice. Utilizzando 🤗 Accelerate abbiamo creato uno script di addestramento con meno di 200 righe di codice che possiamo facilmente scalare su molte GPU. Con questo puoi ora addestrare il tuo modello GPT-2!

Questo post offre una breve panoramica di CodeParrot 🦜, ma se sei interessato a approfondire come pre-addestrare questi modelli, ti consigliamo di leggere il capitolo dedicato nel prossimo libro su Transformers e NLP . Questo capitolo fornisce molti altri dettagli sulla creazione di set di dati personalizzati, considerazioni di progettazione durante l’addestramento di un nuovo tokenizer e scelta dell’architettura.