weather-llm-scaledup-training

WeatherLLM (~1.2B)

A ~1.2B parameter decoder-only language model trained from scratch to answer historical U.S. weather and climate questions. Built end-to-end — no pretrained weights borrowed from elsewhere.

Example queries it is designed to answer:

• "What is January historically like in Arizona?"
• "How do summers in Florida compare to Texas?"
• "Does it snow in Denver in March?"
• "What was the average temperature in Houston in the 1980s?"

---

Architecture

LLaMA-style decoder-only transformer:

---

Training Plan

---

Data Sources

Each NOAA document is a natural-language sentence generated from the raw data — e.g. "Houston Intercontinental AP recorded a July average high of 93.2°F in 2018."

Processed training data is publicly available at: AuraWorxAI/weatherLLMscaleup-data

---

Pretraining Curriculum

Training uses a three-phase curriculum to gradually shift focus from general English to dense climate facts:

Each phase resumes from the previous phase's best checkpoint.

---

Repository Structure

configs/
  model_1b.yaml            # Model architecture (~1.2B params)
  pretrain_phase_a.yaml    # Phase A training config (30K steps)
  pretrain_phase_b.yaml    # Phase B training config (50K steps)
  pretrain_phase_c.yaml    # Phase C training config (15K steps)
  sft_default.yaml         # SFT training config
  data_default.yaml        # Data pipeline settings

scripts/
  download_noaa.py         # Download NOAA GHCN-Daily station CSVs
  preprocess_noaa.py       # Convert GHCN-Daily → weather.jsonl
  download_normals.py      # Download NOAA Climate Normals archives
  preprocess_normals.py    # Convert Normals → normals.jsonl
  download_nclimdiv.py     # Download NOAA nClimDiv files
  preprocess_nclimdiv.py   # Convert nClimDiv → nclimdiv.jsonl
  preprocess_wiki.py       # Download + chunk Wikipedia → wikipedia.jsonl
  merge_corpus.py          # Merge sources into phase corpora
  build_sft_dataset.py     # Build instruction/response SFT pairs
  train_tokenizer.py       # Train SentencePiece tokenizer
  train_pretrain.py        # Causal LM pretraining
  train_sft.py             # Supervised fine-tuning
  push_data.py             # Upload processed data to HuggingFace
  push_checkpoint.py       # Upload training checkpoint to HuggingFace
  push_model.py            # Upload final model to HuggingFace

src/weather_llm/
  data/                    # Data builders (NOAA parsers, merge logic, SFT generation)
  models/                  # Model factory (LlamaForCausalLM wrapper)
  tokenization/            # SentencePiece tokenizer wrapper
  training/                # Trainer setup, callbacks, logging

data/raw/                  # Downloaded NOAA files (gitignored)
data/processed/            # Generated JSONL corpora (gitignored)
artifacts/tokenizer/       # Trained tokenizer files
artifacts/checkpoints/     # Training checkpoints (gitignored)

---

Reproducing: Phase 1 — Data Pipeline

Requires a CPU instance with ~200 GB disk and internet access.

git clone https://github.com/AuraWorx/weatherLLM_ScaledUp.git
cd weatherLLM_ScaledUp
pip install -e .

# Download and process each source (can run steps 1-6 in parallel)
python scripts/download_noaa.py          # ~25 GB, 30-60 min
python scripts/preprocess_noaa.py        # → data/raw/noaa/weather.jsonl

python scripts/download_normals.py       # ~500 MB, <5 min
python scripts/preprocess_normals.py     # → data/raw/normals/normals.jsonl

python scripts/download_nclimdiv.py      # ~20 MB, fast
python scripts/preprocess_nclimdiv.py    # → data/raw/nclimdiv/nclimdiv.jsonl

python scripts/preprocess_wiki.py        # ~1-2 hrs → data/raw/wikipedia/wikipedia.jsonl

# Merge into three pretrain corpora
python scripts/merge_corpus.py --profile phase_a
python scripts/merge_corpus.py --profile phase_b
python scripts/merge_corpus.py --profile phase_c

# Build SFT instruction/response dataset
python scripts/build_sft_dataset.py

# Push processed data to HuggingFace
huggingface-cli login
python scripts/push_data.py --repo_id AuraWorxAI/weatherLLMscaleup-data

Pushes only training files (train_phase_{a,b,c}.jsonl, val_phase_{a,b,c}.jsonl, sft_train.jsonl, sft_eval.jsonl) — raw NOAA files are excluded. Add --private if the dataset should stay private.

---

Repository Layout

---

Next Steps

Phase 2 — Tokenizer Retrain

Train a new BPE tokenizer on the full Phase B corpus so the vocabulary reflects weather/climate terminology (temperature units, station names, month patterns).

python scripts/train_tokenizer.py --corpus data/processed/data/train_phase_b.jsonl

Phase 3 — Three-Phase Pretraining

Pull training data from HuggingFace on the GPU instance (multi-GPU recommended):

# Example: pull Phase A files (repo layout uses data/<filename> in the dataset)
huggingface-cli download your-username/weatherllm-data \
    --repo-type dataset --local-dir data/processed/data/

Train each phase sequentially. Use configs/model_1b.yaml (~1.2B LLaMA-style). Phase A is a scratch run (no checkpoint). For Phase B/C, set resume_from_checkpoint in the YAML to the best checkpoint-* from the previous phase, or pass --resume_from_checkpoint on the CLI.

# Phase A: general English + domain intro (~30K steps), random init
torchrun --standalone --nproc_per_node=8 scripts/train_pretrain.py \
    --pretrain_config configs/pretrain_phase_a.yaml \
    --model_config configs/model_1b.yaml

# Phase B: weather-heavy mix (~50K steps) — set resume to Phase A best checkpoint, e.g.:
#   in configs/pretrain_phase_b.yaml: resume_from_checkpoint: "artifacts/checkpoints/pretrain/phase_a/checkpoint-30000"
torchrun --standalone --nproc_per_node=8 scripts/train_pretrain.py \
    --pretrain_config configs/pretrain_phase_b.yaml \
    --model_config configs/model_1b.yaml

# Phase C: dense climate facts (~15K steps) — resume from Phase B best checkpoint similarly
torchrun --standalone --nproc_per_node=8 scripts/train_pretrain.py \
    --pretrain_config configs/pretrain_phase_c.yaml \
    --model_config configs/model_1b.yaml

Phase 4 — SFT

torchrun --standalone --nproc_per_node=8 scripts/train_sft.py \
    --sft_config configs/sft_default.yaml

Phase 5 — Evaluation

python scripts/eval_prompts.py --model_dir artifacts/checkpoints/sft/run1

---

Model Architecture

Training objective:

---

Publishing to HuggingFace

Hub mirrors (AuraWorxAI):

New GPU instance (train without re-running the data pipeline):

git clone https://github.com/AuraWorx/weatherLLM_ScaledUp.git
cd weatherLLM_ScaledUp
pip install -r requirements.txt
pip install -r requirements-gpu.txt   # PyTorch CUDA wheels; see requirements-gpu.txt header
pip install -e .

# Processed corpora (large download)
hf download AuraWorxAI/weatherLLMscaleup-data --repo-type dataset --local-dir data/processed/data

# Tokenizer files (mirrored on the code repo)
hf download AuraWorxAI/weather-llm-scaledup-training --include "artifacts/tokenizer/*" --local-dir .

Then start Phase A with torchrun as in Phase 3 — Three-Phase Pretraining.

---

Quick Smoke Test

Validates the full pipeline wiring on a tiny corpus without GPU:

python scripts/bootstrap_minimal_corpus.py
python scripts/train_pretrain.py --pretrain_config configs/pretrain_smoke.yaml --model_config configs/model_1b.yaml
python scripts/train_sft.py --sft_config configs/sft_smoke.yaml

---

Reproducing: Phase 2 — Tokenizer

python scripts/train_tokenizer.py \
  --corpus data/processed/data/train_phase_b.jsonl \
  --out_dir artifacts/tokenizer \
  --vocab_size 32000

Trains a 32K Unigram SentencePiece tokenizer with split_digits=True and byte_fallback=True — same design as LLaMA-2's tokenizer, but trained on weather domain text so climate terms are efficiently encoded.

---

Reproducing: Phase 3 — Pretraining

Requires a GPU instance (recommended: 2× H100 80GB). First download the processed data:

python -c "
from huggingface_hub import snapshot_download
snapshot_download(repo_id='AuraWorxAI/weatherLLMscaleup-data', repo_type='dataset', local_dir='data/processed/data', local_dir_use_symlinks=False)
"

Install Flash Attention 2 for 2–4× training speedup on H100:

pip install flash-attn --no-build-isolation

Train each phase sequentially:

# Phase A — 30,000 steps (~10 hrs on 2× H100)
torchrun --nproc_per_node=2 scripts/train_pretrain.py \
  --pretrain_config configs/pretrain_phase_a.yaml \
  --model_config configs/model_1b.yaml

# Phase B — 50,000 steps (~17 hrs on 2× H100)
# Resume from Phase A best checkpoint
torchrun --nproc_per_node=2 scripts/train_pretrain.py \
  --pretrain_config configs/pretrain_phase_b.yaml \
  --model_config configs/model_1b.yaml \
  --resume_from_checkpoint artifacts/checkpoints/pretrain/phase_a/checkpoint-NNNNN

# Phase C — 15,000 steps (~5 hrs on 2× H100)
# Resume from Phase B best checkpoint
torchrun --nproc_per_node=2 scripts/train_pretrain.py \
  --pretrain_config configs/pretrain_phase_c.yaml \
  --model_config configs/model_1b.yaml \
  --resume_from_checkpoint artifacts/checkpoints/pretrain/phase_b/checkpoint-NNNNN

To find the best checkpoint after each phase:

cat artifacts/checkpoints/pretrain/phase_a/trainer_state.json | grep best_model_checkpoint

To back up a checkpoint to HuggingFace before starting the next phase:

python scripts/push_checkpoint.py \
  --repo_id your-org/weatherllm-checkpoints \
  --checkpoint_dir artifacts/checkpoints/pretrain/phase_a/checkpoint-NNNNN

---

Reproducing: Phase 4 — SFT

After pretraining completes:

python scripts/build_sft_dataset.py
torchrun --nproc_per_node=2 scripts/train_sft.py --sft_config configs/sft_default.yaml

The SFT dataset contains instruction/response pairs grounded in actual station data, with three question styles per example (formal, informal, abbreviated) and city-level Q&A for 52 major US metros.

---

Tokenizer Details