docs/colab.md

# 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
```bash

!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
```bash

import os

os.environ['TORCHDYNAMO_DISABLE'] = '1'

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

```

## 4) Sanity run on dummy data (quick)
```bash

!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)
```bash

!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:
```bash

# 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.