zxell.aiFramework methods used in phase 1 — those directly involved in LLM training (by module)
Date: 2026-10-11 Commit: 866b7a2 (the code as of the completed phase-1 S baseline)
0. About this document
- Purpose: against the policy "the model is implemented in-house; using a framework such as PyTorch is allowed" (a decision from review01), list, by module and file, which framework methods the phase-1 code actually relies on.
- Scope: only framework methods directly involved in LLM training. That means two libraries: SentencePiece (tokenizer) and PyTorch (model and training). Parts that are not training itself, such as file I/O and argument handling, are not listed.
- Order: the order in which the work was actually done,
zxell_prep/(tokenizer and preprocessing) →zxell_model/(model and training). - How it was made: each file was analyzed mechanically to extract the calls, and the descriptions of what they are used for were written after reading the code. "×N" is the number of calls within that file.
| Library | Version | Modules that use it | Role |
|---|---|---|---|
SentencePiece (sentencepiece) | 0.2.2 | zxell_prep/, zxell_model/sample.py | Training the BPE tokenizer, encoding, decoding |
PyTorch (torch) | 2.14.0 (CPU build on the server, CUDA build cu128 on the RTX 3060 machine) | zxell_model/ | Tensor operations, automatic differentiation, layer building blocks, optimization, mixed precision |
1. zxell_prep/ (tokenizer and preprocessing)
In preprocessing, the only library directly involved in LLM training is SentencePiece. Three files use it.
1-1. train_compare_tokenizers.py — training and comparing tokenizers
| Call | Used for |
|---|---|
sentencepiece.SentencePieceTrainer.train | Training the BPE tokenizer (two of them: 48k and 64k vocabularies) |
sentencepiece.SentencePieceProcessor | Loading the trained model |
.encode | Turning the evaluation set into token sequences to measure characters/token per language |
Main arguments passed to SentencePieceTrainer.train:
| Argument | Value | Meaning |
|---|---|---|
model_type | "bpe" | Train with BPE |
vocab_size | 48,000 / 64,000 | Vocabulary size (48k was adopted) |
character_coverage | 0.9995 | Share of characters covered by the vocabulary |
byte_fallback | True | Unknown characters fall back to bytes (no UNK is emitted) |
split_digits | True | Digits are split one at a time |
normalization_rule_name | "nmt_nfkc" | Character normalization |
max_sentence_length | 16,384 | Upper limit on the length treated as one sentence |
input_sentence_size / shuffle_input_sentence | 2,000,000 / True | Use 2 million sentences drawn at random from the training file |
How the training file is built (apportioned by a per-language character cap, with ja thickened to about 19%) is in-house logic. The BPE training algorithm itself is left to SentencePiece (a decision from review10).
1-2. tokenize_full.py — full tokenization and sharding
| Call | Used for |
|---|---|
sentencepiece.SentencePieceProcessor | Loading the adopted 48k model |
.encode | Turning every article into a sequence of token IDs (the step that produces the training data itself) |
.eos_id | The ID of the separator token (EOS) placed between articles |
.vocab_size | Checking that the vocabulary fits in uint16 (below 65,536) |
Normalization, deduplication, the per-language split, article-level shuffling, cutting into shards of 50 million tokens, and verification right after writing are all in-house logic.
1-3. peek_shard.py — eyeballing shards
| Call | Used for |
|---|---|
sentencepiece.SentencePieceProcessor | Loading the 48k model |
.decode | Turning a shard's token sequence back into text for display (to check what is in the training data) |
.eos_id | Finding the boundaries between articles |
1-4. Files that use no framework methods
parse_dump.py (building the snapshot), extract_sample.py (extracting samples), compute_boundaries.py (computing the split boundaries), and textprep.py (text cleanup) contain no calls to a framework directly involved in LLM training. All of their processing is in-house logic.
2. zxell_model/ (model and training)
2-1. The line between "in-house" and "framework"
The conclusion first: what was written in-house and what was left to PyTorch.
| Part | Written in-house | Left to PyTorch |
|---|---|---|
| RMSNorm | The formula itself (multiply by the inverse square root of the mean square; computed after promoting to fp32) | torch.rsqrt, .pow and .mean, the container for learned parameters nn.Parameter |
| RoPE (rotary position embedding) | Building the frequency table, the formula that rotates the first and second halves | torch.arange, torch.outer, .cos, .sin, torch.cat |
| Self-attention | The Q/K/V projection and the split into heads, applying RoPE, the output projection | The attention computation itself: F.scaled_dot_product_attention(is_causal=True) (softmax(QKᵀ/√d)V and the causal mask. An implementation optimized for GPUs) |
| SwiGLU FFN | The structure (three linear layers — gate, up, down — and the multiplication) | nn.Linear, the activation function F.silu |
| Transformer block and the GPT itself | The Pre-Norm residual structure, stacking the layers, weight sharing between the embedding and the LM head, the initialization policy | nn.Embedding, nn.Linear, nn.ModuleList, the initializer nn.init.normal_ |
| Loss | Computing it 4 sequences at a time, because the logits over a 48k vocabulary are large | Cross-entropy F.cross_entropy, recomputation during backpropagation torch.utils.checkpoint.checkpoint |
| Training | The learning-rate schedule (warmup + cosine), deciding which parameters get weight decay, sampling the training data | Automatic differentiation (loss.backward()), optimization (torch.optim.AdamW), gradient clipping, bf16 mixed precision (torch.autocast) |
| Evaluation | Aggregating per-language perplexity (assigning a language to each token) | Per-token loss F.cross_entropy, torch.exp |
| Generation | The autoregressive loop, temperature, top-p (cutting off by cumulative probability) | F.softmax, .sort, .cumsum, .scatter_, torch.multinomial |
What this shows
- There are only three external libraries. Leaving the server aside, only PyTorch, NumPy, and SentencePiece are used. Preprocessing is written entirely with the standard library, apart from SentencePiece.
- The structure of the model and the assembly of its formulas are in-house. RMSNorm, RoPE, SwiGLU, the block structure, weight sharing, the learning-rate schedule, and top-p generation were written by hand. The ready-made layers PyTorch provides (
nn.Transformer,nn.MultiheadAttention,nn.LayerNorm) are not used. - What is left to PyTorch is the underlying computation. That is: tensor operations, automatic differentiation, AdamW, cross-entropy, and mixed precision.
- From the standpoint of "from scratch", the dependence on a framework is greatest in two places.
- Training the tokenizer (SentencePiece). This was decided in review10.
- The core attention computation (
F.scaled_dot_product_attention). The Q, K, V projections and RoPE are in-house, but the computation of softmax(QKᵀ/√d)V and the causal mask uses PyTorch's implementation. It could be replaced with a few lines of in-house code, at a cost in speed and memory.
The second point, the attention computation, has never been explicitly confirmed as a decision. If it matters for the scope of "in-house implementation", it can be swapped for an in-house version and compared on speed and loss.
2-2. model.py — the GPT itself
| Call | Used for |
|---|---|
torch.nn.Module (inherited) | Base class of every layer (RMSNorm, Attention, SwiGLU, Block, GPT) |
torch.nn.Parameter | RMSNorm's learned scale |
torch.nn.Embedding | Token embedding (48k vocabulary × d_model) |
torch.nn.Linear ×6 | The Q/K/V projection, the attention output projection, the three SwiGLU layers, the LM head (all without bias) |
torch.nn.ModuleList | Lining up as many Transformer blocks as there are layers |
torch.nn.init.normal_ ×3 | Weight initialization (standard deviation 0.02; smaller, according to the number of layers, for projections added into the residual stream) |
torch.nn.functional.scaled_dot_product_attention | Attention with a causal mask |
torch.nn.functional.silu | SwiGLU's activation function |
torch.nn.functional.cross_entropy | The next-token prediction loss |
torch.nn.functional.softmax | Turning logits into probabilities during generation |
torch.utils.checkpoint.checkpoint | Splits LM head + loss into chunks and recomputes them during backpropagation (so the logits are not all kept; this fits training into 12GB of VRAM) |
torch.ones / torch.zeros_like | RMSNorm's initial value / the container for reordering probabilities in top-p |
torch.arange ×2, torch.outer | RoPE's frequency table (position × frequency) |
torch.rsqrt | The inverse square root in RMSNorm |
torch.cat ×2 | Joining the rotated first and second halves in RoPE / appending generated tokens |
torch.multinomial | Sampling during generation |
torch.no_grad (decorator) | No gradients are computed during generation |
Main tensor methods:
| Method | Used for |
|---|---|
.view, .transpose, .contiguous, .reshape | Splitting into heads and joining them back, reshaping for the loss |
.split, .chunk | Separating Q/K/V, RoPE's first and second halves, chunking the loss |
.float, .type_as | Promote to fp32 for the computation, then return to the original type (bf16) |
.pow, .mean, .cos, .sin, .unsqueeze | The RMSNorm and RoPE computations |
.sort, .cumsum, .scatter_, .argmax, .sum, .all | The top-p cut-off, greedy selection, the stop condition |
.new_zeros, .numel | The container for the loss total, parameter and token counts |
nn.Module methods: .apply (applies the initialization to every layer), .named_parameters, .parameters.
2-3. data.py — building training batches and evaluation
| Call | Used for |
|---|---|
torch.from_numpy ×3 | Turning a token sequence cut from a shard into a tensor |
torch.stack ×2 | Collecting the inputs (x) and the targets (y = the token one step ahead) into a batch |
torch.zeros ×2 | Containers for the per-language loss totals and counts (fp64) |
torch.autocast | bf16 mixed precision during evaluation |
torch.nn.functional.cross_entropy | Per-token loss (reduction="none", so it can be aggregated per language) |
torch.exp ×2 | Converting mean loss into perplexity |
torch.no_grad (decorator) | No gradients are computed during evaluation |
Main methods: the tensor's .to(device) (transfer to the GPU), .reshape, .sum, .numel; the model's .eval and .train (switching between evaluation and training).
How the position to cut from is chosen (a shard is picked in proportion to its token count, then ctx+1 tokens from a random position) and the step that counts EOS tokens to find each token's language are in-house logic.
2-4. train.py — the training loop
| Call | Used for |
|---|---|
torch.optim.AdamW | Optimization (β 0.9 / 0.95, weight decay 0.1; the fused implementation on CUDA) |
torch.autocast | bf16 mixed precision (on CUDA) |
torch.nn.utils.clip_grad_norm_ | Gradient clipping (limit 1.0). The returned gradient norm is logged |
torch.save, torch.load | Saving checkpoints and resuming (model weights, optimizer state, step count) |
torch.manual_seed | Fixing the random seed (reproducible initialization) |
torch.device, torch.cuda.is_available | Use CUDA when it is available |
Main methods:
| Method | Used for |
|---|---|
loss.backward() | Automatic differentiation (computes the gradient of every parameter) |
opt.step(), opt.zero_grad() | Updating the parameters, resetting the gradients |
opt.state_dict(), opt.load_state_dict() | Saving and restoring the optimizer state |
model.to(device), model.train(), model.state_dict(), model.load_state_dict() | Placing the model on the GPU, training mode, saving and restoring the weights |
.item() | Taking the loss out as a number (for the log) |
p.dim() | Deciding which parameters get weight decay (only weights with 2 or more dimensions) |
The learning-rate schedule (linear warmup → cosine decay, ending at 10%) is written as an in-house formula, without PyTorch's schedulers.
2-5. eval.py / sample.py
| File | Call | Used for |
|---|---|---|
| eval.py | torch.load, torch.device, torch.cuda.is_available | Loads a checkpoint and reports per-language perplexity on val / test (the computation is the evaluation in data.py) |
| sample.py | sentencepiece.SentencePieceProcessor | Loading the 48k model |
| sample.py | .encode, .bos_id / .decode | Turning the prompt into a token sequence / turning the generated tokens back into text |
| sample.py | torch.tensor, torch.manual_seed, torch.device, torch.cuda.is_available | Turning the prompt into a tensor, fixing the sampling seed |
config.py (the definitions of the S / M / L configurations) and fetch_shards.py (fetching shards from the API) contain no calls to a framework directly involved in LLM training.
3. Summary
| Module | Library | Core left to the library | Core written in-house |
|---|---|---|---|
zxell_prep/ | SentencePiece | BPE training, encoding, decoding | Apportioning the training file, normalization, deduplication, the per-language split, shuffling, sharding, verification |
zxell_model/ | PyTorch (+ SentencePiece only for generation) | Tensor operations, automatic differentiation, AdamW, the core attention computation, cross-entropy, mixed precision | The model's structure (RMSNorm, RoPE, SwiGLU, blocks, weight sharing), the learning-rate schedule, how batches are built, per-language evaluation, generation (top-p) |
From the standpoint of "from scratch", the dependence on a framework is greatest in two places:
- Training the tokenizer (
SentencePieceTrainer.train). Decided in review10. - The core attention computation (
F.scaled_dot_product_attention). The Q, K, V projections and RoPE are in-house, but softmax(QKᵀ/√d)V and the causal mask are PyTorch's implementation. The same computation could be replaced with a few lines of in-house code (matrix product + mask + softmax), at a cost in speed and memory.