README.md

# LaughLM

A high-performance **decoder-only transformer training system** built with **JAX + Flax** and optimized for **TPU training**.

LaughLM is designed as a **research-friendly yet production-capable framework** for experimenting with modern transformer architectures while maintaining high training throughput.

The system emphasizes:

- clean modular architecture
- hardware-efficient training
- reproducible experiments
- flexible configuration
- large-scale dataset streaming
- high MFU optimization on TPUs

---

# Features

- **Decoder-only GPT architecture**
- **JAX + Flax implementation**
- **TPU-optimized mixed precision training**
- **Flexible architecture selection**
- **Pre-tokenized memory-mapped datasets**
- **Multiple attention variants**
- **Multiple FFN architectures**
- **Weight tying support**
- **Orbax checkpointing**
- **Optax optimizers**
- **Config-driven experiments**

Supported architecture features:

- MHA / MQA / GQA attention
- RoPE positional encoding
- SwiGLU / GEGLU / GELU MLP
- RMSNorm / LayerNorm
- configurable residual scaling
- multiple LR schedulers
- masked weight decay

---

# Project Structure:
```text

.

├── configs

│   ├── gpu_test.yaml

│   └── test.yaml

├── LaughLM

│   ├── config

│   │   ├── loader.py

│   │   ├── schema.py

│   │   └── validation.py

│   ├── data

│   │   ├── domain_sampler.py

│   │   ├── memmap_loader.py

│   │   ├── shard_writer.py

│   │   ├── tokenizer.py

│   │   └── tokenizer_train.py

│   ├── model

│   │   ├── gpt.py

│   │   ├── layers

│   │   │   ├── attention.py

│   │   │   ├── mlp.py

│   │   │   ├── normalization.py

│   │   │   ├── positional.py

│   │   │   └── residual.py

│   │   ├── parameter_utils.py

│   │   └── transformer_block.py

│   ├── training

│   │   ├── checkpoint.py

│   │   ├── logger.py

│   │   ├── loss.py

│   │   ├── optimizer.py

│   │   ├── scheduler.py

│   │   ├── trainer.py

│   │   ├── train_state.py

│   │   └── train_step.py

│   └── utils

│       └── rng.py

├── LICENSE

├── log.txt

├── pyproject.toml

├── README.md

├── requirements.txt

└── scripts

    ├── build_shard.py

    └── train_gpu_test.py

```

---

# Installation

Clone the repository:

```bash

git clone https://github.com/your-org/LaughLM.git

cd LaughLM

```

Create environment:
```bash

python -m venv venv

source venv/bin/activate

```
Install dependencies:
```bash

pip install -r requirements.txt

```

For TPU environments install JAX:

```bash

pip install --upgrade "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

```

---

Configuration

Experiments are fully defined via YAML configs.

Example:

configs/test.yaml

Configuration sections include:

model architecture

optimizer

scheduler

runtime parameters

dataset sources

tokenizer settings

hardware configuration


Example snippet:
```yaml

model:

  d_model: 768

  num_layers: 12

  num_heads: 12

  vocab_size: 32000

  max_seq_len: 2048

```

---

Dataset Pipeline

LaughLM uses a pre-tokenized dataset pipeline for maximum throughput.

Training datasets are converted into binary token shards.

Advantages:

high throughput

minimal CPU overhead

memory-mapped streaming

scalable to large datasets



---

Step 1 — Train Tokenizer

Train a tokenizer using streaming datasets.
```bash

python -m LaughLM.data.tokenizer_train

```
Output:

tokenizer.json


---

Step 2 — Build Token Shards

Convert raw text into token shards.
```bash

python scripts/build_shard.py

```
Output:

dataset_shard.bin



Shards contain:



uint16 token stream





---



Step 3 — Training



Run training:

```bash

python scripts/train_gpu_test.py

```

Training automatically handles:



optimizer



scheduler



logging



checkpointing





Example output:



STEP   PROGRESS │ LOSS   PPL │ LR │ TOK/S │ MFU





---



Checkpointing



Checkpoints are saved using Orbax.



Default directory:



checkpoints/



Resume training automatically if checkpoints exist.





---



Benchmarking Performance



Benchmark raw training throughput:



python scripts/benchmark_train_step.py



This measures:



compile time



step time



tokens/sec



MFU





Example output:



Compile time: 18.2s

Step time: 0.048s

Tokens/sec: 430000





---



Monitoring



Training logger displays:



loss



perplexity



gradient norm



tokens/sec



MFU



ETA





Example:



STEP  PROGRESS │ LOSS │ LR │ TOK/S │ MFU │ ETA





---



Optimization Roadmap



LaughLM is designed to progressively reach high TPU utilization.



Target MFU:



50–60% MFU on TPU v5e



Optimization phases:



Phase	Goal



Baseline	establish benchmark

Data pipeline	remove input bottlenecks

Graph optimization	eliminate Python overhead

Kernel fusion	maximize MXU utilization

Flash attention	reduce memory traffic







---



Development Workflow



Recommended workflow:



1. Create branch

2. Implement change

3. Run benchmark

4. Compare tokens/sec

5. Merge if improvement



Example:

```bash

git checkout -b optimize_attention
```



---



Contributing



Pull requests should include:



clear description



performance impact



benchmark results







---



License



MIT License





---



Acknowledgements



LaughLM builds on ideas from:



GPT



LLaMA



PaLM



DeepSeek



MiniCPM





and the JAX / Flax ecosystem.





---



Future Work



Planned improvements:



Flash Attention



Activation checkpointing



MoE layers



PJIT sharding



distributed training