Writing a transformer

Phase 1:

• Embedding transformer ---Showcase attention , apply triton optimization (FP32 - FP16, etc)

dataset

a) Tinystories dataset

1. wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt
2. wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt

b) OpenWebText

1. wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_train.txt.gz , gunzip owt_train.txt.gz
2. wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_valid.txt.gz , gunzip owt_valid.txt.gz

Transformer architecture in short

Step	Component	Input	Output	Precision
1-2	DataLoader + Tokenizer	Raw text corpus lines	idx: [B, T], targets: [B, T]	int64
3	CombinedEmbeddingPipeline	idx: [B, T]	x: [B, T, D]	float32
4-6	TransformerBlock × layers	x: [B, T, D]	x: [B, T, D]	float32
7	LayerNormFromScratch (final)	x: [B, T, D]	x: [B, T, D]	float32
8	self.lm_head projection	x: [B, T, D]	logits: [B, T, V]	float32
9	nn.functional.cross_entropy	logits: [B*T, V], targets: [B*T]	scalar loss	float32

• Typical shapes in the current test path: B = 3, T = 128, D = 64 for the embedding check, and vocabulary size V = 256 for the byte-level tokenizer setup.

• Vision transformer (ViT)