Train in Google Colab

This guide shows how to train ULTRATHINK on a small setup in Colab with stable defaults.

1) Connect GPU runtime

Runtime → Change runtime type → Hardware accelerator: GPU (T4/A100 preferred)

2) Clone repo and install deps

!git clone https://github.com/vediyappanm/UltraThinking-LLM-Training.git
%cd UltraThinking-LLM-Training/deep

# Pin versions known to work well in Colab
!pip install --upgrade pip
!pip install torch --index-url https://download.pytorch.org/whl/cu121
!pip install "transformers>=4.41.0" "datasets==2.14.7" accelerate einops tqdm wandb

# Optional (FlashAttention). Skip if it fails on your GPU/driver.
# !pip install flash-attn --no-build-isolation

3) Environment tweaks for stability

import os
os.environ['TORCHDYNAMO_DISABLE'] = '1'
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

4) Sanity run on dummy data (quick)

!python train_ultrathink.py \
  --dataset dummy --train_samples 2000 --val_samples 200 \
  --vocab_size 50257 --hidden_size 384 --num_layers 4 --num_heads 6 --num_kv_heads 6 \
  --intermediate_size 1536 --max_seq_length 256 \
  --batch_size 4 --gradient_accumulation_steps 8 \
  --learning_rate 5e-4 --use_amp --gradient_checkpointing \
  --num_epochs 1 \
  --output_dir ./outputs/sanity_dummy

5) Small real-data run (C4 streaming)

!python train_ultrathink.py \
  --dataset c4 --dataset_subset en --streaming \
  --enable_dre --dre_warmup_steps 500 \
  --amp_warmup_steps 200 \
  --tokenizer_name gpt2 \
  --vocab_size 50257 \
  --hidden_size 384 --num_layers 4 --num_heads 6 --num_kv_heads 6 \
  --intermediate_size 1536 --max_seq_length 512 \
  --batch_size 1 --gradient_accumulation_steps 64 \
  --learning_rate 5e-5 --weight_decay 0.1 \
  --warmup_steps 2000 \
  --use_amp --gradient_checkpointing \
  --eval_frequency 1 \
  --output_dir ./outputs/ultrathink_c4_seq512_sdpa_warmup

6) Monitoring

Watch the printed average loss to trend downward.
Optional: add --use_wandb --run_name ultrathink_colab to log to Weights & Biases.

7) Switching datasets

See docs/datasets.md for details. Examples:

# Wikipedia streaming
!python train_ultrathink.py --dataset wikipedia --dataset_subset 20231101.en --streaming --tokenizer_name gpt2

# Custom JSONL
!python train_ultrathink.py --dataset custom --data_path /content/data.jsonl --text_column text --tokenizer_name gpt2

# Mixed datasets (50/50 wikitext + openwebtext)
!python train_ultrathink.py --mix_datasets "wikitext:0.5,openwebtext:0.5" --tokenizer_name gpt2 --streaming

Tips

If you hit non-finite losses early, lower LR (2e-5), keep --max_seq_length 512, or increase warmups.
FlashAttention may not install on all Colab GPUs; SDPA fallback is already enabled in the model.