Load a pretrained decoder and keep training it

Open in Colab View on GitHub DownloadGPU

Notebook 05 trains a language model from random weights. Most of the time you start from a model someone else already trained. In this notebook we load SmolLM2-135M, a small Llama-style model from Hugging Face, into Dew’s own CausalTransformer, continue its training on Shakespeare for a few hundred steps, and export the result back to the Hugging Face format. At the end we load the export with the transformers library and check that it predicts the same next token as Dew.

The notebook expects one NVIDIA GPU and takes about three minutes on a Colab L4. The first run downloads the checkpoint (about 270 MB).

Install

interop adds safetensors support. torch is only for the last cell, where transformers reads the exported checkpoint.

%pip install -q "dew-ml[cuda13,interop] @ git+https://github.com/AshishKumar4/dew" torch

  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done

Settings

import os

CHECKPOINT = "HuggingFaceTB/SmolLM2-135M"
SEQUENCE_LENGTH = 256
BATCH_SIZE = 8
STEPS = 300
LEARNING_RATE = 3e-5
MAX_NEW_TOKENS = 60
PROMPT = "ROMEO:"
DATA_DIR = "data/08-shakespeare"
RUN_DIR = "runs/08-continued"
EXPORT_DIR = "runs/08-exported"
SEED = 0

# The tutorial test sets DEW_TUTORIAL_SMOKE=1 to run every cell in minutes on a CPU.
if os.environ.get("DEW_TUTORIAL_SMOKE") == "1":
    SEQUENCE_LENGTH, BATCH_SIZE, STEPS, MAX_NEW_TOKENS = 64, 2, 2, 8

import json
import os
import urllib.request
from pathlib import Path

import jax
import jax.numpy as jnp
import numpy as np

print(jax.devices())

[CudaDevice(id=0)]

Loading the checkpoint

load_pretrained takes a Hub repo id or a local folder in the Hugging Face layout. It translates the checkpoint’s config.json into CausalTransformer settings and its weights into Dew’s parameter tree. The weights stay in float32 and the model computes in the dtype we ask for. If the config has a setting that changes the computation and Dew has no equivalent, load_pretrained raises an error naming it instead of silently dropping it.

It returns a Pretrained with the model on .model, its weights on .variables, and the Dew settings it derived on .model_config.

from dew.interop import load_pretrained

pretrained = load_pretrained(CHECKPOINT, dtype="bfloat16", attention_impl="auto", max_seq_len=512)
model, variables = pretrained.model, pretrained.variables
n_params = sum(x.size for x in jax.tree_util.tree_leaves(variables))
print(f"{n_params / 1e6:.0f}M parameters")
print({key: pretrained.model_config[key] for key in
       ("vocab_size", "emb_features", "num_layers", "num_heads", "num_kv_heads", "head_dim")})

stderr

/usr/local/lib/python3.13/dist-packages/huggingface_hub/utils/_auth.py:138: UserWarning: 
Error while fetching `HF_TOKEN` secret value from your vault: 'Requesting secret HF_TOKEN timed out. Secrets can only be fetched when running from the Colab UI.'.
  warnings.warn(f"\nError while fetching `HF_TOKEN` secret value from your vault: '{str(e)}'.")

135M parameters
{'vocab_size': 49152, 'emb_features': 576, 'num_layers': 30, 'num_heads': 9, 'num_kv_heads': 3, 'head_dim': 64}

The model was trained on the ids of its own tokenizer, so we must use that tokenizer too. HFTokenizer wraps any Hugging Face tokenizer behind the same encode and decode as the byte tokenizer.

from dew.data import HFTokenizer

tokenizer = HFTokenizer(CHECKPOINT)
print("vocabulary:", tokenizer.vocab_size)
print(tokenizer.encode(PROMPT))

vocabulary: 49152
[3911, 3945, 63, 42]

Generating before training

SmolLM2 learned from web text, code and textbooks, so it knows what a play looks like but does not write like Shakespeare. We generate greedily, taking the most likely token every time, so the output is the same on every run.

from dew.sampling import Sampling, generate

prompt = jnp.asarray([tokenizer.encode(PROMPT)], jnp.int32)
before = generate(model, variables, prompt, max_new_tokens=MAX_NEW_TOKENS,
                  key=jax.random.key(0), sampling=Sampling(temperature=0.0))
print(tokenizer.decode(before.tokens[0]))

ROMEO: I think that's a good point.

JOHNSON: I think that's a good point.

MARTIN: I think that's a good point.

JOHNSON: I think that's a good point.

MARTIN: I think

The data

The token files follow the layout from notebook 05, but with SmolLM2’s tokenizer. Its vocabulary has 49,152 entries, too many for one byte, so the ids are stored as uint16.

data_dir = Path(DATA_DIR)
data_dir.mkdir(parents=True, exist_ok=True)
text_path = data_dir / "input.txt"
urllib.request.urlretrieve(
    "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt", text_path)

ids = np.asarray(tokenizer.encode(text_path.read_text(encoding="utf-8")), np.uint16)
val_len = len(ids) // 20
ids[:val_len].tofile(data_dir / "val.bin")
ids[val_len:].tofile(data_dir / "train.bin")
meta = {"tokenizer": CHECKPOINT, "vocab_size": tokenizer.vocab_size, "dtype": "uint16",
        "train_tokens": len(ids) - val_len, "val_tokens": val_len, "eos_id": None}
(data_dir / "meta.json").write_text(json.dumps(meta, indent=2))
print(meta)

stderr

[transformers] Token indices sequence length is longer than the specified maximum sequence length for this model (341094 > 8192). Running this sequence through the model will result in indexing errors

{'tokenizer': 'HuggingFaceTB/SmolLM2-135M', 'vocab_size': 49152, 'dtype': 'uint16', 'train_tokens': 324040, 'val_tokens': 17054, 'eos_id': None}

from dew.data import Loading, TokenWindows

data = TokenWindows(path=DATA_DIR, seq_len=SEQUENCE_LENGTH, val_batches=8,
                    loading=Loading(workers=0, threads=1, read_buffer=2)).load(batch=BATCH_SIZE)
print("training windows:", data.records)

training windows: 1265

Continued training

LMObjective takes the loaded weights through pretrained=, so training starts from them instead of from a fresh random draw. We turn off the EMA (ema_decay=None) because a 300-step run is too short for an average to help. The learning rate is small so the model adapts to Shakespeare without forgetting what it knew. The validation perplexity at steps 100, 200 and 300 shows the adaptation.

import optax
from dew import Checkpoints, Trainer, metrics
from dew.objectives.lm import LMObjective

objective = LMObjective(model, SEQUENCE_LENGTH, ema_decay=None, pretrained=variables)
trainer = Trainer(objective, optax.adamw(LEARNING_RATE), key=jax.random.key(SEED),
                  checkpoints=Checkpoints(RUN_DIR))
state = trainer.fit(data, steps=STEPS, log_every=50, eval_every=100, checkpoint_every=STEPS,
                    metrics=(metrics.perplexity(),))

Training from step 0 to 300 on {'data': 1, 'expert': 1, 'fsdp': 1, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))
step 50: loss 3.2403
step 100: loss 3.0720
Evaluation val at step 100: 8 coordinated batches, 64 records, uneven_shards=False, event_key=(4286894075, 772130920): {'val/perplexity': 29.26405867480638}
step 150: loss 3.1543
step 200: loss 3.1194
Evaluation val at step 200: 8 coordinated batches, 64 records, uneven_shards=False, event_key=(1298626594, 818800949): {'val/perplexity': 28.173023282652128}
step 250: loss 3.2285
step 300: loss 3.2523
Evaluation val at step 300: 8 coordinated batches, 64 records, uneven_shards=False, event_key=(4127360435, 4068970345): {'val/perplexity': 27.850388273194902}
Goodput: first step after 34.72 s, 34.7% of the wall time in steps

after = generate(model, state.params, prompt, max_new_tokens=MAX_NEW_TOKENS,
                 key=jax.random.key(0), sampling=Sampling(temperature=0.0))
print(tokenizer.decode(after.tokens[0]))

ROMEO:
I have heard of him, and I know him.

LADY CAPULET:
He is a nobleman, and a nobleman's son.

MERCUTIO:
He is a nobleman, and a nobleman's son.

LAD

Validation perplexity went from about 29 at step 100 to about 28 at step 300, a small change from so few steps. The greedy text changed more: the speakers are now Capulets and Mercutio, and the lines have the play’s layout. It still repeats itself, as greedy decoding tends to.

Exporting back to Hugging Face

save_pretrained_decoder writes the trained weights in the Hugging Face layout: config.json, model.safetensors, and the tokenizer files. It runs the same name and shape translation as load_pretrained, backwards. Loading the export again gives back exactly the weights we trained.

from dew.interop import save_pretrained_decoder

save_pretrained_decoder(model, state.params, EXPORT_DIR, tokenizer=CHECKPOINT)
print(sorted(os.listdir(EXPORT_DIR)))

reloaded = load_pretrained(EXPORT_DIR, dtype="bfloat16", attention_impl="auto", max_seq_len=512)
same = all(np.array_equal(np.asarray(a), np.asarray(b))
           for a, b in zip(jax.tree_util.tree_leaves(reloaded.variables["params"]),
                           jax.tree_util.tree_leaves(state.params["params"])))
print("reloaded weights equal the trained ones:", same)

['config.json', 'generation_config.json', 'model.safetensors', 'tokenizer.json', 'tokenizer_config.json']
reloaded weights equal the trained ones: True

Reading the export with `transformers`

An export is only useful if another library can read it. We load the folder with transformers, run the prompt through it in float32 on the CPU, and compare its most likely next token with Dew’s.

import torch
from transformers import AutoModelForCausalLM

hf_model = AutoModelForCausalLM.from_pretrained(EXPORT_DIR, dtype=torch.float32)
with torch.no_grad():
    hf_next = int(hf_model(input_ids=torch.tensor(np.asarray(prompt))).logits[0, -1].argmax())

dew_next = generate(model, state.params, prompt, max_new_tokens=1,
                    key=jax.random.key(0), sampling=Sampling(temperature=0.0))
dew_next = int(dew_next.tokens[0, -1])
print("transformers:", hf_next, repr(tokenizer.decode([hf_next])))
print("dew:         ", dew_next, repr(tokenizer.decode([dew_next])))

transformers: 198 '\n'
dew:          198 '\n'

Where to go next

The README model list names the decoder families Dew loads, trains, generates with and exports. The recipe flag --pretrained runs this notebook’s flow at scale, for example python recipes/lm/train.py data:token-windows --data.path <dir> --pretrained HuggingFaceTB/SmolLM2-135M --tokenizer HuggingFaceTB/SmolLM2-135M.