Des sous-mots et un entraînement à l'échelle

Dixième article de la série Smallem. Précédent : Entraîner, sauvegarder, générer.

Le modèle de l'article précédent lisait le texte octet par octet. Avec un contexte de 256 tokens, il voyait donc 256 octets : une ou deux phrases. Cet article remplace les octets par des sous-mots, puis entraîne un modèle trois fois plus gros sur un fichier entier du corpus. Il a fallu pour cela écrire un tokenizer, mais aussi résoudre quatre problèmes pratiques qui n'existent pas à petite échelle.

Pourquoi pas des octets, pourquoi pas des mots

Travailler en octets est coûteux : pour écrire « information », le modèle doit prédire onze tokens de suite, sans se tromper. Une grande partie de sa capacité sert à épeler.

L'autre extrême serait un token par mot. Mais le vocabulaire serait immense, des centaines de milliers de mots, et il resterait toujours des mots inconnus : un nom propre, une faute de frappe ou encore un mot d'une autre langue.

Les sous-mots sont un compromis. Les mots fréquents deviennent un seul token et les mots rares sont découpés en morceaux courants. Dans le pire des cas on retombe sur les octets, donc aucun texte n'est jamais hors vocabulaire.

L'algorithme qui construit ce vocabulaire s'appelle BPE, pour byte pair encoding. Venu de la compression, il a été appliqué aux mots par Sennrich, Haddow et Birch (2016).

L'algorithme, sur un exemple

Le BPE part des 256 octets et répète une seule opération :

trouver la paire de tokens voisins la plus fréquente, et en faire un nouveau token.

Prenons un corpus minuscule, où chaque mot est suivi de son nombre d'occurrences :

low : 5     lower : 2     newest : 6     widest : 3

Au départ, chaque mot est une suite de lettres. On compte les paires voisines, en tenant compte des occurrences. La paire e s apparaît dans newest (6 fois) et widest (3 fois), soit 9 fois, autant que s t. Supposons qu'on fusionne e s en premier : on crée le token es.

l o w : 5     l o w e r : 2     n e w es t : 6     w i d es t : 3

On recompte. La paire es t apparaît maintenant 9 fois : on crée est.

l o w : 5     l o w e r : 2     n e w est : 6     w i d est : 3

Puis l o (7 fois) donne lo, puis lo w (7 fois) donne low. Le mot low est devenu un seul token, et le suffixe est aussi.

Chaque fusion ajoute un token au vocabulaire. On s'arrête à la taille voulue, ici 8 192 tokens, soit 7 935 fusions. La liste ordonnée des fusions est le tokenizer.

Les numéros de tokens

Les numéros prolongent ceux des octets :

0 à 255        les octets
256            la fin de document
257 et plus    une fusion par numéro, dans l'ordre d'apprentissage

Un BPE sans aucune fusion est exactement le tokenizer octets des articles précédents. Le projet n'a donc plus qu'un seul tokenizer, et tous les programmes fonctionnent avec les deux sortes de tokens.

pub struct Bpe {
    merges: Vec<Pair>,
    /// Pour chaque paire fusionnable, le token qu'elle produit.
    merged: HashMap<Pair, u16>,
    /// Les octets que représente chaque token.
    token_bytes: Vec<Vec<u8>>,
}

Ne pas fusionner à travers les mots

Si l'on applique l'algorithme au texte brut, la paire la plus fréquente sera vite du genre « e + espace + t », à cheval sur deux mots. Ces tokens n'ont pas de sens et gaspillent le vocabulaire.

On découpe donc d'abord le texte en mots, et les fusions n'ont lieu qu'à l'intérieur d'un mot. La règle de découpage :

• un mot est une suite de lettres, une suite de chiffres, ou une suite d'autres signes
• il emporte avec lui l'espace qui le précède
• les autres blancs, comme les retours à la ligne, forment des mots à part.

"Hello, world 42"   ──►   "Hello"   ","   " world"   " 42"

Attacher l'espace au mot qui suit est une convention reprise de GPT-2. Un mot en milieu de phrase, world, n'est pas le même token que le même mot en début de texte.

Ce découpage a un second avantage décisif pour la vitesse. Dans 100 Mo de texte, il y a 20,7 millions de mots, mais seulement 352 063 mots différents. Il suffit de compter combien de fois chaque mot apparaît, puis de travailler sur les mots distincts, pondérés par leur nombre d'occurrences. Le problème devient soixante fois plus petit.

Un détail de robustesse : un mot est coupé au bout de 64 caractères. Le web contient de longues suites sans espace, des adresses, des clés encodées, des lignes de tirets, et sans cette limite l'encodage d'une seule d'entre elles pourrait prendre un temps démesuré.

Rendre l'entraînement incrémental

La version naïve recompte toutes les paires de tous les mots après chaque fusion. Avec 7 935 fusions, c'est 7 935 passes complètes : beaucoup trop lent.

Heureusement une fusion ne change que les mots qui contiennent la paire fusionnée. Fusionner q u ne touche à aucun mot sans qu.

On tient donc à jour trois structures :

• le compte de chaque paire
• pour chaque paire, la liste des mots où elle apparaît
• un tas, qui donne à tout moment la paire la plus fréquente.

À chaque fusion, on ne retouche que les mots concernés. Pour chacun, on retire ses anciennes paires des comptes, on applique la fusion, et on ajoute ses nouvelles paires.

// Le mot change : on retire toutes ses anciennes paires des
// comptes, puis on ajoute les nouvelles.
for old in symbols.windows(2) {
    let old = (old[0], old[1]);
    *pair_counts.get_mut(&old).expect("paire comptée") -= *word_count;
    changed.insert(old);
}
for new in replaced.windows(2) {
    let new = (new[0], new[1]);
    *pair_counts.entry(new).or_default() += *word_count;
    pair_words.entry(new).or_default().push(index);
    changed.insert(new);
}

Le tas pose un petit problème : on ne peut pas modifier le compte d'une paire qui s'y trouve déjà. La solution est de ne jamais corriger le tas. Quand un compte change, on empile une nouvelle entrée. Les anciennes restent, périmées, et on les reconnaît à la sortie : leur compte ne correspond plus au compte actuel de la paire.

let Some((count, Reverse(pair))) = heap.pop() else {
    break;
};
if pair_counts.get(&pair) != Some(&count) {
    continue; // entrée périmée
}

Dernier point : que faire quand deux paires ont le même compte ? Il faut une règle, sans quoi le résultat dépendrait de l'ordre de parcours d'une table de hachage et changerait à chaque exécution. Ici, la plus petite paire l'emporte.

Tester une optimisation

Comment savoir que la version rapide calcule la même chose que la version simple ? En gardant la version simple.

Elle vit dans les tests, et sert de référence : sur un petit corpus, les deux doivent apprendre exactement les mêmes fusions, dans le même ordre.

#[test]
fn incremental_training_matches_naive_training() {
    let word_counts = counts(&sample_text());
    for vocab_size in [257, 258, 300, 400, 2000] {
        let bpe = Bpe::train(&word_counts, vocab_size);
        assert_eq!(bpe.merges(), train_naive(&word_counts, vocab_size));
    }
}

C'est la même idée que le tenseur CPU servant d'oracle au GPU. Une version lente et évidente garantit la version rapide et subtile.

Encoder et décoder

Pour encoder un mot, on part de ses octets et on applique les fusions dans l'ordre où elles ont été apprises. À chaque tour, parmi les paires présentes dans le mot, c'est la fusion la plus ancienne qui passe.

fn encode_word(&self, word: &[u8]) -> Vec<u16> {
    let mut symbols: Vec<u16> = word.iter().map(|&byte| u16::from(byte)).collect();
    loop {
        let earliest = symbols
            .windows(2)
            .filter_map(|pair| {
                let pair = (pair[0], pair[1]);
                self.merged.get(&pair).map(|&token| (token, pair))
            })
            .min();
        let Some((token, pair)) = earliest else {
            return symbols;
        };
        symbols = merge_pair(&symbols, pair, token);
    }
}

Comme les tokens sont numérotés dans l'ordre d'apprentissage, « la fusion la plus ancienne » est simplement celle qui produit le plus petit numéro : d'où le .min().

Dans un texte courant, presque tous les mots ont déjà été rencontrés. Un encodeur garde donc en mémoire les mots déjà encodés, et l'encodage se réduit le plus souvent à une recherche dans une table.

Le décodage est trivial : on met bout à bout les octets de chaque token, puis on interprète le tout en UTF-8. Il faut le faire dans cet ordre. Un caractère comme é occupe deux octets, qui peuvent se retrouver dans deux tokens différents : décoder token par token casserait le caractère.

Le tokeniseur, en chiffres

Les premières fusions apprises sont celles qu'on attend de l'anglais :

"␣t"  "␣a"  "he"  "in"  "re"  "␣the"  "␣o"  "er"  "on"  "at"

et les dernières sont des mots entiers, de plus en plus spécialisés :

"␣decay"  "␣Additionally"  "␣diameter"  "␣Vitamin"  "␣chemistry"  "␣ecological"

La validation la plus importante est l'aller-retour : décoder un texte encodé doit redonner exactement le texte d'origine, à l'octet près. Elle est vérifiée sur les documents de validation, qui n'ont pas servi à apprendre les fusions.

Avec 3,92 octets par token, un contexte de 512 tokens couvre environ 2 000 octets : un paragraphe entier, là où le modèle précédent voyait une phrase.

Passer à l'échelle : quatre problèmes pratiques

Avec un bon tokenizer, on peut viser plus gros : un fichier entier du corpus, soit 3,3 Go de texte, et un modèle de 30 millions de paramètres. Quatre choses qui allaient de soi à petite échelle cessent alors de marcher.

1. Tokeniser 3,3 Go

Encoder 100 Mo sur un seul coeur est supportable. Pour 3,3 Go, on répartit le travail. Chaque groupe de 1 000 documents est découpé en tranches, une par coeur. Chaque fil a son propre encodeur, avec sa propre mémoire des mots déjà vus.

std::thread::scope(|scope| {
    let workers: Vec<_> = texts
        .chunks(per_thread)
        .zip(encoders.iter_mut())
        .map(|(slice, encoder)| {
            scope.spawn(move || {
                slice
                    .iter()
                    .map(|text| encode_document(bpe, encoder, text))
                    .collect::<Vec<_>>()
            })
        })
        .collect();
    // …
})

thread::scope appartient à la bibliothèque standard. Les threads peuvent emprunter le tokenizer et les textes sans rien copier, parce que la portée garantit qu'ils auront fini avant que ces données ne disparaissent : Avant de rendre la main, thread::scope attend tous les threads lancés par ce scope.

Chaque document est décodé aussitôt encodé, pour vérifier l'aller-retour. Les 726 000 documents du fichier sont passés par ce contrôle.

3 309 Mo de texte en 15,6 s (212 Mo/s sur 16 threads)
847 540 068 tokens d'entraînement, 44 406 509 de validation

2. Lire des données qui ne tiennent pas en mémoire

Jusqu'ici, le fichier de tokens était chargé en entier. Avec plusieurs fichiers de 1,7 Go, ce n'est plus raisonnable.

Mais un lot n'a besoin que de quelques fenêtres, 16 fenêtres de 513 tokens par exemple. On laisse donc les fichiers sur le disque, et on ne lit que les fenêtres :

pub fn read(&self, shard: usize, start: usize, len: usize) -> io::Result<Vec<u16>> {
    let mut bytes = vec![0u8; len * 2];
    self.shards[shard]
        .file
        .read_exact_at(&mut bytes, start as u64 * 2)?;
    let (pairs, _) = bytes.as_chunks::<2>();
    Ok(pairs.iter().map(|&pair| u16::from_le_bytes(pair)).collect())
}

Le système d'exploitation garde en mémoire les pages déjà lues tant qu'il a de la place. Préparer un lot prend 14 microsecondes, alors qu'un pas d'entraînement se compte en centaines de millisecondes : autant dire rien.

3. Un lot plus grand que la mémoire du GPU

Un grand lot donne un gradient moins bruité, donc un entraînement plus stable. Mais la mémoire nécessaire grandit avec le lot, et celle de la carte est limitée.

L'accumulation de gradient contourne la limite. On traite plusieurs petits lots l'un après l'autre, on additionne leurs gradients, et on ne met le modèle à jour qu'une fois.

Cela fonctionne parce que la perte est une moyenne sur les exemples. La moyenne des gradients de deux lots de 16 séquences est exactement le gradient d'un lot de 32. La mémoire occupée, elle, reste celle d'un seul petit lot.

for index in 0..micro_batches {
    let (loss_value, tape_id, grads) =
        loss_and_gradients(backend, model, |m, t| loss(m, t, index));
    for (sum, parameter) in grad_sums.iter_mut().zip(model.parameters()) {
        let Some(grad) = parameter.grad(tape_id, &grads) else {
            continue;
        };
        *sum = Some(match sum.take() {
            Some(sum) => backend.add(&sum, grad),
            None => grad.clone(),
        });
    }
}

Le test dit précisément ce qu'on attend : deux lots accumulés, ou quatre, doivent donner le même pas qu'un seul grand lot, aux arrondis près.

4. Survivre à une interruption

Un entraînement de plusieurs heures finit par être interrompu : une coupure, un redémarrage ou une erreur. Il faut pouvoir reprendre sans tout recommencer.

Pour cela enregistrer les poids ne suffit pas : Adam garde pour chaque paramètre deux moyennes glissantes, qui mettent des centaines de pas à se stabiliser. Sans elles, un entraînement repris repartirait de moyennes nulles et ferait de mauvais premiers pas.

Un point de reprise contient donc les poids, les deux moyennes de chaque paramètre et le numéro du pas. C'est un fichier de poids ordinaire, avec des tenseurs en plus.

Il reste le hasard. Si les lots tirés après la reprise ne sont pas ceux qui auraient été tirés sans interruption, la suite diverge. La solution est de faire dépendre le hasard du numéro de pas, et de rien d'autre :

fn step_rng(seed: u64, step: usize) -> StdRng {
    StdRng::seed_from_u64(
        seed.wrapping_mul(0x9E37_79B9_7F4A_7C15)
            .wrapping_add(step as u64),
    )
}

Le lot du pas 3 000 est alors toujours le même, qu'on y arrive d'une traite ou après dix reprises.

Dernière précaution : une reprise n'a de sens qu'avec les mêmes réglages. Avec un autre nombre de pas, le taux d'apprentissage ferait un saut. Avec une autre graine, les lots ne seraient plus les mêmes. Le point de reprise contient donc aussi les réglages qui décident de la suite : nombre de pas, taille du lot, taux, décroissance des poids, seuil d'écrêtage, graine, architecture. À la reprise, ils sont comparés à ceux de la commande, avant que quoi que ce soit ne soit chargé, et la moindre différence est refusée :

les réglages ne sont pas ceux du point de reprise (run.max_grad_norm : 0.8 enregistré, 1 demandé)

Tout cela se vérifie par un test, puisque la boucle d'entraînement est une fonction de la bibliothèque (article 9). Un entraînement de six pas tourne d'une traite, et son point de reprise du pas 3 est mis de côté au moment où il est écrit. Un second entraînement part de cette copie, avec un modèle initialisé autrement. Les deux doivent finir avec exactement les mêmes paramètres, à l'égalité stricte.

Le contrôle est sans appel sur CPU : un entraînement interrompu puis repris produit un fichier identique, à l'octet près, à celui d'un entraînement ininterrompu. Sur GPU, les pertes sont identiques à quatre décimales, mais pas les octets. Ce n'est pas la reprise qui est en cause : deux entraînements GPU ininterrompus diffèrent entre eux de la même façon. L'explication la plus probable est l'addition atomique des gradients d'embedding, dont l'ordre varie d'une exécution à l'autre, mais ce point n'a pas été vérifié.

Le modèle et son résultat

L'entraînement a duré deux heures : 5 500 pas de deux lots accumulés de 16 séquences, soit 90 millions de tokens. C'est 11 % du fichier. À ce débit, 12 300 tokens par seconde, le parcourir une seule fois aurait pris 19 heures.

La perte de validation finale est de 3,68 nats par token.

Comparer des pertes avec des tokens différents

Ce nombre paraît bien pire que le 1,20 du modèle précédent. Il ne l'est pas, car l'unité a changé. Prédire un token BPE, c'est prédire presque quatre octets d'un coup : il est normal que ce soit plus difficile que d'en prédire un.

Pour comparer, on ramène la perte à l'octet, en la divisant par le nombre moyen d'octets par token :

3,68 nats par token ÷ 3,897 octets par token = 0,944 nat par octet

C'est la seule unité qui se compare d'un tokenizer à l'autre, et c'est pour cela que toute la série s'y tient. La comparaison reste imparfaite : les deux derniers modèles n'ont pas été évalués exactement sur les mêmes documents, et n'ont ni la même taille ni le même temps d'entraînement.

Le texte

To teach children how to read, teachers are encouraged to read and evaluate the
reading skills they have.
I can only learn a textbook to make your own words more comfortable and more
challenging and a new skill. It is a challenging task that is worth the time to do.

Le saut est net. Les phrases sont grammaticales, et le texte tient son sujet, la lecture et l'enseignement, sur plusieurs phrases.

Il faut être honnête sur ses limites. Sur d'autres amorces, le modèle tourne vite en boucle sur une expression, et ses « faits » sont inventés. Les pertes d'entraînement et de validation étaient encore confondues à la fin : ce modèle aussi manque de pas.

Un bug révélé par l'échelle

L'entraînement s'est terminé normalement, puis le programme a planté, en générant l'échantillon de texte final : plus de mémoire sur la carte.

La cause venait d'une optimisation de l'article 11. Pour éviter d'allouer et de libérer sans cesse la mémoire GPU, les tampons abandonnés sont gardés dans une réserve, rangés par taille, pour être réutilisés. Pendant l'entraînement, les mêmes tailles reviennent à chaque pas, et la réserve atteint vite un régime stable.

En génération, la séquence s'allonge d'un token à chaque tour. Chaque longueur demande de nouvelles tailles de tampons, que la réserve gardait toutes, sans jamais rien rendre. Avec le gros modèle, la carte finissait pleine.

La correction : quand une allocation échoue, la réserve se vide, puis on réessaie.

let buffer = reused.unwrap_or_else(|| {
    DeviceBuffer::zeroed(&self.stream, len).unwrap_or_else(|_| {
        // Plus de mémoire : on rend au pilote tout ce qui dort dans
        // la réserve, puis on réessaie une fois.
        self.release();
        DeviceBuffer::zeroed(&self.stream, len).expect(CUDA_ERROR)
    })
});

Ce bug existait depuis longtemps. Les petits modèles ne remplissaient simplement jamais la carte. Un test le reproduit désormais en allouant quarante tenseurs d'un gigaoctet, tous de tailles différentes.

Ce qu'il faut retenir

• Le BPE construit un vocabulaire en fusionnant, encore et encore, la paire de tokens voisins la plus fréquente.
• Les fusions restent à l'intérieur des mots, ce qui permet aussi de ne traiter que les mots distincts.
• Une version naïve gardée dans les tests garantit la version rapide.
• À l'échelle, il faut paralléliser la préparation, lire les données à la demande, accumuler les gradients et savoir reprendre.
• Les pertes ne se comparent d'un tokenizer à l'autre qu'en nats par octet.

Le modèle n'a vu que 11 % de son fichier en deux heures. La qualité est maintenant bornée par la vitesse : Rendre le GPU rapide.