Views
No views yet
| Variant | d_model | Layers | n_recurse | Params | Train time |
|---|---|---|---|---|---|
| medium | 384 | 6 | 3 | ~54M | ~2h |
| large | 512 | 8 | 4 | ~116M | ~6h |
deatos/tokenized_fineweb_edu_10b_combined using a ml.g4dn.xlarge (T4 GPU) on AWS SageMaker.1import torch
2import tiktoken
3from model import ModelConfig, LatentTransformer
4
5# Load tokenizer
6tokenizer = tiktoken.get_encoding("gpt2")
7
8# --- Medium ---
9cfg = ModelConfig(
10 vocab_size=tokenizer.n_vocab, max_seq_len=128,
11 d_model=384, n_heads=8, n_layers=6, n_recurse=3, dropout=0.0
12)
13
14# --- Large ---
15# cfg = ModelConfig(
16# vocab_size=tokenizer.n_vocab, max_seq_len=128,
17# d_model=512, n_heads=8, n_layers=8, n_recurse=4, dropout=0.0
18# )
19
20model = LatentTransformer(cfg)
21
22# Checkpoints were saved via torch.compile — strip the _orig_mod. prefix
23state = torch.load("checkpoints/latent_transformer_medium.pt", map_location="cpu", weights_only=True)
24if any(k.startswith("_orig_mod.") for k in state):
25 state = {k.removeprefix("_orig_mod."): v for k, v in state.items()}
26model.load_state_dict(state)
27model.eval()
28
29# Inference
30def get_causal_mask(seq_len, device):
31 return torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
32
33@torch.no_grad()
34def generate(prompt, max_new_tokens=60, temperature=0.8, top_k=40, n_recurse=3):
35 model.cfg.n_recurse = n_recurse
36 tokens = tokenizer.encode(prompt)
37 src = torch.tensor([tokens], dtype=torch.long)
38 tgt = src.clone()
39 for _ in range(max_new_tokens):
40 mask = get_causal_mask(tgt.shape[1], src.device)
41 logits = model(src, tgt, tgt_mask=mask)
42 next_logits = logits[:, -1, :].squeeze(0) / temperature
43 if top_k > 0:
44 top_vals = torch.topk(next_logits, top_k)[0][..., -1, None]
45 next_logits[next_logits < top_vals] = float("-inf")
46 probs = torch.softmax(next_logits, dim=-1)
47 next_tok = torch.multinomial(probs, 1).unsqueeze(0)
48 tgt = torch.cat([tgt, next_tok], dim=1)
49 return tokenizer.decode(tgt[0].tolist())
50
51print(generate("The strange truth was"))