docs: add Colab quickstart; feat: safer train data check and single-GPU generate fix

1 files changed, 44 insertions, 0 deletions

diff --git a/Group-Entropy-Equalization/README.md b/Group-Entropy-Equalization/README.md
index 804af95..33bd020 100644
--- a/Group-Entropy-Equalization/README.md
+++ b/Group-Entropy-Equalization/README.md
@@ -13,6 +13,50 @@ pip install -r requirements.txt
 
 ---
 
+### Colab Quickstart (single-GPU, no DeepSpeed)
+
+In Colab, use a smaller model first to verify end-to-end. Then scale up if VRAM allows.
+
+```bash
+!git clone https://github.com/YurenHao0426/gee.git
+%cd /content/gee/Group-Entropy-Equalization
+!pip -q install transformers==4.44.2 accelerate==0.33.0 peft==0.12.0 bitsandbytes==0.43.3 datasets==2.21.0 wandb==0.17.7 pyarrow==17.0.0
+```
+
+Create a small parquet if you don’t have one:
+
+```python
+import os, pandas as pd
+os.makedirs("dataset/1shot_rlvr", exist_ok=True)
+df = pd.DataFrame({"problem": [
+    "What is 2 + 2?",
+    "If x=3, compute x^2 + 2x + 1.",
+    "The doctor is a ____.",
+    "Factor 12.",
+    "What is 7*8?",
+]})
+df_big = pd.concat([df]*256, ignore_index=True).iloc[:1280]
+df_big.to_parquet("dataset/1shot_rlvr/pi1_r1280.parquet", index=False)
+```
+
+Run training (no DeepSpeed, no AMP to avoid Colab GradScaler quirks):
+
+```bash
+!python train.py \
+  --model_name Qwen2.5-1.5B \
+  --model_path Qwen/Qwen2.5-1.5B \
+  --train_data dataset/1shot_rlvr/pi1_r1280.parquet \
+  --effective_batch 4 --micro_batch_size 1 \
+  --temperature 0.5 --learning_rate 2e-5 --sample_temp 0.5 \
+  --max_steps 10 --log_steps 1 --save_steps 10 \
+  --run_name colab_em10 --wandb_project one-shot-em \
+  --no_deepspeed --mixed_precision no
+```
+
+Checkpoints are saved under `checkpoints/<model>/<run_name>/`.
+
+---
+
 ### Reproducing One-shot EM Training (SOTA)
 
 ```bash