Näin koulutat pienen kielimallin

Koulutamme TensorFlow’lla 19,4 miljoonan parametrin Transformer-mallin. Se oppii ennustamaan seuraavan tavun lyhyistä englanninkielisistä tarinoista. Koulutuksen jälkeen malli voi tuottaa tekstiä selaimessa TensorFlow.js:n avulla.

Toteutus käyttää Keras-kirjaston rakennuspalikoita attention-mekanismiin hyödyntämättä valmista Transformer-toteutusta.

Avaa koko Python- ja selainkoodi uuteen välilehteen.

01Koulutusympäristö

Varsinaisessa koulutusajossa käytettiin yhtä A100-GPU:ta Verda-palvelussa.

Hyperparametreiksi on valittu:

Koulutuksessa käytetään 32 sekvenssin eriä ja float32-painoja. Optimoijana on AdamW: oppimisnopeus on 0,0003, painojen vaimennus 0,01 ja gradienttien yhteisen normin leikkausraja 1,0.

02Datan valmistelu

Käytämme SimpleStories-aineistoa, jossa on yli kaksi miljoonaa suuremmilla kielimalleilla tuotettua synteettistä englanninkielistä tarinaa. Aineiston muodostaminen kuvataan tutkimusjulkaisussa. Käyttämämme aineistoversion tarinatekstiä oli saatavilla noin 2,62 Gt.

Hugging Face Datasets tarjoaa tarinat suoratoistona. Jaamme ne normalisoidun tekstin tiivisteen perusteella ja varaamme noin prosentin validointiin. Loput ovat koulutusesimerkkejä. Näin tarina pysyy samassa osajoukossa eri ajokerroilla. Arviointiin käytämme kiinteää 1 000 tarinan otosta validointijoukosta.

Koodaamme tarinat tf.data-rajapinnalla, lisäämme erottimet, yhdistämme tavut ja muodostamme vakiomittaisia jaksoja. Tekstien välimuisti on RAM-muistissa:

texts = tf.data.Dataset.from_generator(
    lambda: (row["story"] for row in training_stories),
    output_signature=tf.TensorSpec((), tf.string),
).cache()

data = texts.map(
    lambda text: tf.io.decode_raw(text + "\n\n", tf.uint8)
)
data = data.unbatch().batch(513, drop_remainder=True)
data = data.shuffle(4096, seed=42).batch(32, drop_remainder=True)
training = data.map(
    lambda chunk: (chunk[:, :-1], chunk[:, 1:])
).prefetch(tf.data.AUTOTUNE)

Jokaisesta 513 tavun jaksosta saadaan 512 syötetavua ja 512 tavoitetavua yhden paikan siirtymällä. Malli ennustaa kutakin syötekohtaa seuraavan tavun. Jakso voi sisältää osan tarinasta tai ylittää tarinoiden välisen erottimen.

Mallin sanasto koostuu 256 mahdollisesta tavuarvosta. Tavuarvot toimivat suoraan syötteen tunnisteina, joten kokeilu ei tarvitse koulutettua tokenisoijaa. Kontekstin ja tuotetun tekstin pituus mitataan siksi tavuina, ei sanoina tai sanojen osista koostuvina tokeneina.

03Transformerin määrittely

Jokainen syötetavu muunnetaan opittuun 512 komponentin upotusvektoriin. Lisäämme opitun paikkaupotuksen ja sovellamme kuutta Transformer-lohkoa. Kussakin lohkossa on kausaalinen monipäinen self-attention-mekanismi ja eteenpäinkytketty verkko. Kumpaakin operaatiota edeltää normalisointi, ja niiden ympärillä on jäännösyhteydet.

Lohkon eteenpäinlaskenta:

def call(self, x, training=False):
    normalized = self.norm1(x)
    x = x + self.attention(
        normalized, normalized,
        use_causal_mask=True, training=training)
    return x + self.contract(self.expand(self.norm2(x)))

expand on tiheä kerros, jossa on 2 048 ulostuloa ja GELU-aktivointi. contract projisoi takaisin 512 komponenttiin. Attention-kerroksessa on kahdeksan päätä, joista kunkin ulottuvuus on 64. Molempien normalisointikerrosten epsilon on 0,00001.

Viimeisen lohkon jälkeen normalisointi ja tiheä kerros tuottavat kullekin kohdalle 256 logittia. Syöteupotuksella ja ulostuloprojektiolla on erilliset painot.

Kausaalinen maski estää syötekohtaa käyttämästä myöhempien kohtien tietoa attention-mekanismissa. Tämä on olennaista koulutuksessa: koko jakso on saatavilla, myös tavut, joita aiempien kohtien pitäisi ennustaa. Maskin ansiosta ennusteet voidaan laskea rinnakkain koko jaksolle siten, että tehtävä säilyy seuraavan tavun ennustamisena.

04Koulutus ja validointi

Sparse categorical cross-entropy vertaa logitteja tavoitetavujen tunnisteisiin. TensorFlow laskee häviön gradientit mallin läpi, ja AdamW päivittää painot.

Keskeinen koulutuskutsu:

optimizer = tf.keras.optimizers.AdamW(
    learning_rate=0.0003,
    weight_decay=0.01,
    global_clipnorm=1.0,
)
optimizer.exclude_from_weight_decay(var_names=["bias", "beta", "gamma"])
model.compile(
    optimizer=optimizer,
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
model.fit(
    training.repeat(),
    validation_data=validation,
    steps_per_epoch=500,
    epochs=40,
    shuffle=False,
    callbacks=[Progress(checkpoints)],
)

Tässä kukin raportoitu epoch tarkoittaa 500 päivityksen jaksoa, ei koko aineiston läpikäyntiä. Neljäkymmentä tällaista jaksoa antaa tavoitteeksi 20 000 päivitystä. Takaisinkutsufunktio tallentaa mallin ja optimoijan tarkistuspisteen sekä tulostaa tuotetun tekstinäytteen kunkin jakson lopussa. Optimoijan tilan tallentaminen mahdollistaa koulutuksen jatkamisen kertyneitä momentteja hyödyntäen.

Valmiin pienen mallin validointihäviö oli 0,622 natia tavua kohti eli noin 0,898 bittiä tavua kohti. Mittaus käyttää koulutuksesta erillään pidettyä tekstiä. Tuotettuja jatkoja tarkastelemme erikseen.

Esimerkiksi malli jatkoi tätä lihavoituna esitettyä alkua:

Far far away, he found a letter in the corner of the park.

05Tekstin tuottaminen

Jatkon tuottamiseksi koodaamme kehotteen tavuiksi ja suoritamme mallin. Viimeisen kohdan logitit antavat seuraavan tavun pisteet. Jaamme ne lämpötilalla, poimimme näytteen näin saadusta jakaumasta, lisäämme valitun tavun ja toistamme.

Matala lämpötila keskittää näytteenoton todennäköisempiin vaihtoehtoihin. Kiinteä siemen tekee satunnaislukujonosta toistettavan, mikä auttaa kehotteiden ja tarkistuspisteiden vertailussa.

Selainversio käsittelee koko kehotteen yhdellä prefill-ajolla. Sen jälkeen se säilyttää kunkin kerroksen avaimet ja arvot KV-välimuistissa. Seuraavat askeleet käsittelevät yhden uuden tavun tätä välimuistia vasten sen sijaan, että koko alku laskettaisiin uudelleen.

Tässä arkkitehtuurissa 512 kohdan välimuisti vie float32-muodossa 12 MiB. Välimuisti kuuluu nykyiseen tekstijaksoon; mallin painot ovat yhteiset eri tekstintuottokerroille. Demo pysähtyy, kun kehotteen ja tuotetun tekstin yhteispituus on 512 tavua.

Kokeile mallia. Voit säätää kehotetta, lämpötilaa, siementä ja tulosteen pituutta sekä nähdä prefill-ajan, tuottonopeuden ja KV-välimuistin koon.

06Malli selaimeen

Alkuperäiset float32-painot vievät 77,8 Mt. Latausta varten kvantisoimme kunkin tensorin etumerkittömiksi kahdeksan bitin arvoiksi ja tallennamme minimin sekä skaalauskertoimen. TensorFlow.js palauttaa float32-arvot mallia ladattaessa.

Vientikoodi pyöristää kunkin tensorin painot yhteen 256 arvosta ja tallentaa palauttamiseen tarvittavat tiedot:

array = value.numpy().astype("<f4")
minimum = float(array.min())
scale = (float(array.max()) - minimum) / 255 or 1.0
codes = np.clip(
    np.rint((array - minimum) / scale), 0, 255
).astype(np.uint8)

stream.write(codes.tobytes())
specs.append({
    "name": name,
    "shape": list(array.shape),
    "dtype": "float32",
    "quantization": {
        "dtype": "uint8",
        "min": minimum,
        "scale": scale,
    },
})

Painot vievät tämän jälkeen 19,4 Mt tai gzip-pakattuina 16,15 Mt. Kiinteässä vertailuotoksessamme kahdeksan bitin versio nosti häviön 0,92279:stä 0,92308 bittiin tavua kohti. Tämä arviointi käyttää eri otosta kuin koulutusajon validointijoukko.

Pienempi lataus ei pienennä muistissa olevia float32-painoja: ne vievät edelleen noin 78 Mt. Suorituksen puskurit, välitulosten tensorit ja KV-välimuisti tarvitsevat lisäksi muistia.

Tekstin tuottaminen tapahtuu worker-säikeessä TensorFlow.js:n avulla. Käytössä on WebGPU, kun sopiva sovitin on saatavilla, ja muuten WebAssembly. Kehotteet ja tuotettu teksti pysyvät käyttäjän laitteella.

Selaimessa TensorFlow.js purkaa binääritiedoston manifestin tensorimuotojen, minimien ja skaalauskertoimien avulla:

const config = await fetch("model/model-int8.json")
  .then(response => response.json());
const buffer = await fetch(`model/${config.weight_file}`)
  .then(response => response.arrayBuffer());
const weights = tf.io.decodeWeights(buffer, config.weights);
const model = new TinyLM(config, weights);

Alla oleva päättelykoodi on käsin kirjoitettua JavaScriptiä, joka käyttää TensorFlow.js:n operaatioita. Se toteuttaa koulutetun Keras-mallin laskennan; TensorFlow ei generoi tätä koodia. Kutsu step() koko kehotteella tekee prefill-laskennan. Seuraavat kutsut saavat yhden uuden tavun ja palautetun välimuistin.

class TinyLM {
  constructor(config, weights) {
    this.config = config;
    this.weights = weights;
  }

  norm(x, name) {
    const {mean, variance} = tf.moments(x, -1, true);
    return x.sub(mean).mul(variance.add(this.config.epsilon).rsqrt())
      .mul(this.weights[`${name}.gamma`]).add(this.weights[`${name}.beta`]);
  }

  dense(x, name) {
    return tf.matMul(x, this.weights[`${name}.kernel`]).add(this.weights[`${name}.bias`]);
  }

  step(bytes, cache = []) {
    const {width, heads, layers, context} = this.config;
    const offset = cache.length ? cache[0].key.shape[1] : 0;
    const length = bytes.length;
    if (!length || offset + length > context) throw new Error("Context limit reached");
    return tf.tidy(() => {
      const ids = tf.tensor1d(bytes, "int32");
      let x = tf.gather(this.weights.tokens, ids)
        .add(tf.gather(this.weights.positions, tf.range(offset, offset + length, 1, "int32")));
      const next = [];
      for (let i = 0; i < layers; i++) {
        const name = `b${i}`;
        const normalized = this.norm(x, `${name}.norm1`);
        const project = part => this.dense(normalized, `${name}.${part}`)
          .reshape([length, heads, width / heads]).transpose([1, 0, 2]);
        const query = project("q");
        let key = project("k"), value = project("v");
        if (cache.length) {
          key = tf.concat([cache[i].key, key], 1);
          value = tf.concat([cache[i].value, value], 1);
        }
        next.push({key, value});
        let scores = tf.matMul(query, key, false, true).mul(1 / Math.sqrt(width / heads));
        if (length > 1) {
          const future = tf.range(0, offset + length, 1, "int32").expandDims(0)
            .greater(tf.range(offset, offset + length, 1, "int32").expandDims(1));
          scores = scores.add(future.cast("float32").mul(-1e9));
        }
        const attention = tf.matMul(tf.softmax(scores, -1), value)
          .transpose([1, 0, 2]).reshape([length, width]);
        x = x.add(this.dense(attention, `${name}.o`));
        const expanded = this.dense(this.norm(x, `${name}.norm2`), `${name}.expand`);
        const gelu = expanded.mul(0.5).mul(tf.erf(expanded.div(Math.SQRT2)).add(1));
        x = x.add(this.dense(gelu, `${name}.contract`));
      }
      const last = this.norm(x.slice([length - 1, 0], [1, width]), "norm");
      return {logits: this.dense(last, "head").reshape([256]), cache: next};
    });
  }
}

self.TinyLM = TinyLM;

Palautetut logitit määrittävät seuraavan tavun jakauman. Kutsuva koodi poimii tavun, syöttää sen takaisin välimuistin kanssa ja vapauttaa edelliset logitit ja välimuistin jokaisen askeleen jälkeen. tf.tidy() vapauttaa välitulosten tensorit mutta säilyttää palautetut tensorit.

07Suurempi kokeilu

Erillinen ajo kaksinkertaistaa leveyden 1 024:ään ja lohkojen määrän kahteentoista. Mallissa on 152,2 miljoonaa parametria eli noin 7,8 kertaa pienen mallin määrä. Konteksti pysyy 512 tavussa, attention-päitä on 16 ja eteenpäinkytketyn verkon leveys on 4 096.

Suurempi malli saavutti 20 000 päivityksen kohdalla validointihäviön 0,5477 natia tavua kohti. Valmiin pienemmän mallin tulos oli 0,6222, joten häviö pieneni 12 %. Vertailu mittaa seuraavan tavun ennustamista samalla validointiasetelmalla.

Voit valita kumman tahansa mallin selaindemossa. Suurempi versio käyttää 20 000 päivityksen tarkistuspistettä. Sen kahdeksan bitin painot latautuvat gzip-pakattuina 123,5 Mt:n kokoisina ja laajenevat 608,8 Mt:n float32-painoiksi. Lisäksi tarvitaan suorituksen puskurit ja enintään 48 MiB KV-välimuistia. Oletuksena pysyy 19,4 miljoonan parametrin malli.