LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)
Pretraining

Data Pipeline

Premium

Understand how Parquet shards become input_ids, labels, and attention_mask

This section ties together the previous chapters: the pretraining data has already been split into Parquet shards, the tokenizer can already turn text into token ids, and the model has already been defined. What's missing in between is the data pipeline.

In cookllm-bento, the process of generating a pretraining batch can be summarized as:

Pretrain Data Pipeline

How a Parquet row becomes a training batch for next-token prediction.

1
Parquet shard
read text column from row groups
list[str]
2
Batch tokenization
prepend BOS, encode texts in parallel
token ids
3
Truncate to max_length
keep the first 512 tokens by default
sample
4
input_ids / labels
same tokens; shifting happens inside the model
tensors
5
Batch padding
pad input_ids, mask labels with -100
attention_mask
6
PretrainModule.training_step()
forward batch and log train/loss
loss

This pipeline is handled by two files:

LevelFileRole
DataModulesrc/datamodule/pretrain_datamodule.pyParse data paths, load the tokenizer, split train/val, create the DataLoader
Datasetsrc/dataset/pretrain.pyStream-read Parquet, batch tokenization, generate a single training sample

Configuration Entrypoint

The core fields of the data config used by the pretraining example are as follows:

Log in to continue reading

This is premium content. Please log in to access the full article.

configs/data/data_pretrain_sample.yaml
data:
  data_path: pretrain_data/fineweb_shards
  tokenizer_path: tokens
  max_length: 512
  batch_size: 96
  num_workers: 4
  tokenizer_batch_size: 128
  tokenizer_num_threads: 8
  val_files: 10
ConfigMeaning
data_pathThe pretraining Parquet shard directory
tokenizer_pathThe tokenizer file directory; by default reads tokens/ under the project root
max_lengthHow many tokens at most to keep per sample
batch_sizeHow many samples the DataLoader returns each time
num_workersThe number of DataLoader workers
tokenizer_batch_sizeThe number of texts sent to the tokenizer at a time
tokenizer_num_threadsThe number of internal parallel threads for the tokenizer
val_filesHow many Parquet files to randomly keep as the validation set

Model Architecture

Read BentoLM's structure and parameter size from bento_29m.yaml

Training Loop

Take apart LightningCLI, PretrainModule, the optimizer, and the scheduler

Table of Contents

Configuration Entrypoint
Data Directory
File Split
Streaming Reads
Sample Shape
Batch Padding
Data Throughput Test
Entering the Training Loop