Load a pretrained decoder and keep training it

Open in Colab View on GitHub DownloadGPU

In notebook 05 we trained a language model from random weights. But most of the time, you would start from a model someone else already trained, on far more data than we have. So in this notebook, we load SmolLM2-135M, a small Llama-style model from Hugging Face, described in the paper “SmolLM2: When Smol Goes Big” by Allal et al. (2025), 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 Dependencies

interop adds safetensors support, and 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

install log

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

Settings

Note the learning rate, 3e-5, a lot smaller than the 1e-3 of notebook 05; the continued training section explains why.

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 for it, 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; the byte tokenizer of notebook 05 would feed it ids that mean something else entirely in its vocabulary. 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

Before we train anything, let’s see what the model writes as it is. SmolLM2 learned from web text, code and textbooks, so it knows what a play looks like, but it 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

Data Pipeline

The token files follow the layout from notebook 05, but with SmolLM2’s tokenizer this time. Its vocabulary has 49,152 entries, far 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 that the model adapts to Shakespeare without forgetting what it already 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 a lot more though: the speakers are now Capulets and Mercutio, and the lines have the play’s layout. It still repeats itself, as greedy decoding tends to do.

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, just 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 actually read it. So 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.