ShayonSarker/KochiLLM-v1.0

KochiLLM v1.0

A 1.02B-parameter Mixture-of-Experts language model trained entirely from scratch on 2x Tesla T4 (~52 GPU-hours), including its own BPE tokenizer, data pipeline, and distributed training setup.

Weights status. This repository contains the model code, configuration, and documentation. Weights live under runs/ -- one folder per training run, with runs/index.json as the manifest. run-09 is the final, fully annealed checkpoint (step 5,499, val 3.96). Run notebooks/05_upload_hf.ipynb on Kaggle to (re)publish them.

Documentation

Document	Contents
ARCHITECTURE.md	Component-level reference: every tensor's shape and parameter count, MoE routing and load balancing, QK-Norm, RoPE, KV-cache/sliding-window interaction, the 768-GPU-sync optimisation ceiling
TRAINING.md	Two-phase recipe, schedule retargeting, FSDP memory arithmetic, checkpoint/resume protocol, full loss series, engineering lessons, cost-performance context
SAMPLES.md	Before/after generations on identical prompts, with a diagnosis of what the final epoch changed and what it did not
training_curve.csv	Machine-readable loss series with tokens seen and both LR schedules
config.json	Exact training config; MyCustomLLM(**config) builds 1,023,262,000 parameters

Configuration

Total parameters	1,023,262,000 (1.02B)
Active per token	634,259,000 (634M)
Layers / model dim	12 / 1536
Attention heads	24 query / 8 KV (GQA ratio 3), head_dim 64
Experts	4 routed (top-2) + 1 shared, SwiGLU (hidden = 2 x dim = 3072)
Attention window	alternating: full on even layers, 512-token sliding on odd
Normalization	RMSNorm (eps 1e-6) + QK-Norm on queries and keys
Position encoding	RoPE, theta 10000, precomputed to 2048 positions
Soft-capping	attention logits 50, final logits 30 (Gemma-2 style)
Multi-token prediction	n_predict=2; extra head discarded at inference
Weight tying	embed.weight is tied to output.weight
Tokenizer	custom BPE, 32,000 vocab, regex pre-tokenization
Context length	1024 tokens

Training at a glance

Corpus	FineWeb-Edu + Cosmopedia-v2, 4:1 mix, streamed
Train / val tokens	180,295,081 / 20,109,016
Training tokens	180,224,000 (0.9996 epochs, 0.176 tokens/param)
Steps	5,500, each consuming 32,768 tokens
Optimizer	AdamW, lr 1e-3 -> 1e-4 cosine, 200 warmup steps, weight decay 0.1
Gradient clip	1.0
Precision	fp16 + GradScaler (Tesla T4 has no bf16 support)
Parallelism	FSDP FULL_SHARD across 2 GPUs, one shard unit per block
Throughput	~37.9 s/step, ~865 tokens/s
Validation loss	3.96 final (annealed), 3.91 best, 5.64 at Phase-2 start

Architecture

Token IDs [B, 1024]
  -> Embedding 32000 x 1536 (tied to output head), x * sqrt(1536), dropout 0.05
  -> 12 x CustomBlock:
       RMSNorm
       Grouped Query Attention: 24Q / 8KV, RoPE, QK-Norm, softcap 50,
         causal mask (full on even layers, 512-token sliding window on odd)
       residual add
       RMSNorm
       MoE FFN: shared expert on every token + top-2 of 4 routed experts,
         gate weights renormalized over the selected 2
       residual add, dropout 0.05
  -> final RMSNorm
  -> output head 1536 x 32000 (tied), softcap 30
  -> cross-entropy + 0.01 * mean over layers of (4 * sum_i f_i * p_i)  [load balancing]
  -> MTP extra head (predicts +2), averaged with the main loss; training only

Loading (once weights are published)

This is not an AutoModel release: the architecture is custom code and the tokenizer is a custom pickle. Load it directly from this repo:

import json, torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from architecture import MyCustomLLM
from tokenizer import MyCustomTokenizer

REPO = 'ShayonSarker/KochiLLM-v1.0'
TAG  = 'run-09'                      # see runs/index.json for the final run

path = hf_hub_download(REPO, 'runs/%s/config.json' % TAG, token=...)
config = json.load(open(path))

tok = MyCustomTokenizer(vocab_size=config['vocab_size'])
tok.load('custom_bpe_tokenizer.pkl')

model = MyCustomLLM(**config)                 # re-ties embed.weight <-> output.weight
sd = load_file(hf_hub_download(REPO, 'runs/%s/model.safetensors' % TAG, token=...))
sd['output.weight'] = sd['embed.weight']      # no-op here: FSDP emits both, byte-identical
model.load_state_dict(sd, strict=False)
model.eval().cuda().to(torch.bfloat16)

from generate import generate
print(generate(model, tok, 'The theory of relativity was developed by',
               max_new_tokens=100, temperature=0.8, top_k=40, device='cuda'))

See SAMPLES.md for worked examples with real output.

Checkpoints

Every training run is published in its own folder on main. Folders are numbered chronologically by step count, not by Kaggle run name:

runs/
├── index.json                       manifest: step, tokens, val loss, source run
├── run-01/                          earliest run
│   ├── model.safetensors            bf16, ~2.05 GB
│   ├── config.json
│   └── metadata.json
├── ...
└── run-09/                          final run  <- the published model

from huggingface_hub import hf_hub_download
sd = hf_hub_download('ShayonSarker/KochiLLM-v1.0',
                     'runs/run-09/model.safetensors', token=...)

runs/index.json is the manifest:

{
  "model": "KochiLLM-v1.0",
  "final_run": "run-09",
  "runs": {
    "run-01": { "iter": 970,  "tokens_seen": 31801360, "val_loss": 3.8501,
                "val_loss_is_phase2": false, "source_run": "kochi-llm-trainning-1",
                "path": "runs/run-01" }
  }
}

val_loss_is_phase2 — checkpoints saved before the Phase-2 best-tracking fix carry the Phase-1 Alpaca loss (3.8501) in val_loss and have no latest_val_loss. Where that flag is false, the honest reading is "step N of a run whose Phase-2 validation loss was not separately recorded", not "this run scored 3.85".

Weights are bf16, cast down from the fp32 training checkpoints. Use them for inference; to resume training you would need the original fp32 checkpoint.

Repository contents

File	Purpose
config.json	Exact training config
architecture.py	Model definition: MoE, GQA, RoPE, QK-Norm, sliding window, MTP
tokenizer.py	Custom BPE: train / encode / decode / save / load
data_loader.py	Memmap-backed batching over tokenized .bin shards
generate.py	KV-cache sampling with temperature and top-k, plus a CLI
prepare_data.py	Phase-1 Alpaca tokenization and tokenizer training
train_custom.py	FSDP training loop, cosine schedule, checkpoint protocol
ARCHITECTURE.md	Component-level architecture reference
TRAINING.md	Training report
SAMPLES.md	Before/after generations
training_curve.png / .csv	Validation-loss curve and raw series
LICENSE	MIT
requirements.txt	torch, numpy at runtime

Limitations

• No factual knowledge. At 0.176 tokens/param the model learned fluency and style, not facts. Compute-optimal for 1B parameters is roughly 20 tokens/param (~20B tokens); this run is data-limited by about 100x. Generated names, dates, and attributions are plausible and wrong.
• Not instruction-tuned. Trained only on web pages and textbook prose, so it continues documents rather than answering questions. An SFT phase would teach the format of answering, not the truth.
• Short context. 1024 tokens.
• Noisy validation. Each eval averages only 10 random 1024-token batches; run-to-run noise is roughly +/- 0.05-0.1, so best and final differ by less than the noise floor.
• Evaluation was in-distribution. The holdout comes from the same corpora as training.
• No document-boundary masking. <|endoftext|> is inserted as a literal token, but the attention mask does not reset there, so sequences can attend across document boundaries.
• Optimizer state was not checkpointed across session boundaries, so AdamW moments reset ~4 times. The upgrade path is sharded FSDP optimizer-state serialization.

Source code

Full training code, all five Kaggle notebooks, and step-by-step reproduction instructions: https://github.com/Dadhichi-Sarker-Shayon/KochiLLM-v1.0

License

MIT — see LICENSE. Weights are released under the same terms.