LittleBit: recipe and throughput (weights coming soon)

README.md

@@ -0,0 +1,94 @@
+---
+license: cc-by-nc-4.0
+base_model: Qwen/Qwen3-8B
+base_model_relation: quantized
+library_name: transformers
+tags: [littlebit, littlebit-2, sub-1-bit, quantization, qat, qwen3]
+---
+
+![LittleBit.](assets/banner.png)
+
+# LittleBit
+
+**Big models. Little bits.** Sub-1-bit Qwen3 models built with LittleBit-2:
+binarized latent factorization, recovered by distillation from the original model.
+
+> **Status: weights coming soon.** This repo holds the training recipe and measured
+> throughput. Model weights and eval results will be added here when the first full
+> training run finishes. Follow [@LittleBit_llm](https://x.com/LittleBit_llm) for updates.
+
+> Independent community project. Not affiliated with or endorsed by Samsung Research.
+> Built on the LittleBit method and code by Lee, Kim, You & Kim ([SamsungLabs/LittleBit](https://github.com/SamsungLabs/LittleBit)).
+
+## Planned releases
+
+| Model | Base | Target bpw | Status |
+|---|---|---|---|
+| littlebit-qwen3-4b | Qwen/Qwen3-4B | 0.55 | planned |
+| littlebit-qwen3-8b | Qwen/Qwen3-8B | 0.55 | planned |
+| littlebit-qwen3-14b | Qwen/Qwen3-14B | 0.55 | planned |
+
+Bits per weight apply to linear layers. Embeddings and `lm_head` stay BF16.
+
+## Method
+
+Each linear layer `W` is approximated as `sign(U) · diag(h·g·ℓ) · sign(V)ᵀ`:
+low-rank latent factors binarized to ±1, plus three thin learned scale vectors.
+
+1. **Latent factorization:** SVD splits each linear layer into rank-r factors sized to the bit budget.
+2. **Joint-ITQ rotation (LittleBit-2):** aligns the factors with the binary hypercube before training. It folds into the factors, so it adds no inference cost.
+3. **SmoothSign binarization:** a smooth surrogate gradient keeps the sign step trainable.
+4. **Residual compensation:** a second binarized path learns what the first one missed.
+5. **Distillation:** quantization-aware training on C4 + WikiText-2 (seq len 2048), with the BF16 model as teacher (logit KL + layer-to-layer MSE).
+
+## Measured throughput
+
+Measured on RunPod with the recipe in [`recipe/`](recipe/). One step = 4 sequences × 2048 tokens.
+One epoch of the C4-shard-0 + WikiText-2 mix is about 20,750 steps with the Qwen3 tokenizer.
+
+| Model | GPU | Sec / step | Peak VRAM | One-time init (SVD + Joint-ITQ) | Est. 1 epoch |
+|---|---|---:|---:|---:|---:|
+| Qwen3-0.6B @ 0.55 bpw | 1× H100 80GB | 1.03 | — | ~2.3 min | ~6 h |
+| Qwen3-8B @ 0.55 bpw | 1× H200 141GB | 2.84 | ~107 GB | ~14 min | ~16.5 h |
+
+Qwen3-8B does not fit on a single 80 GB GPU: about 3.7B latent parameters are trainable,
+and their optimizer state alone exceeds the memory. Use a 141 GB GPU or ≥ 2 GPUs with ZeRO-3.
+
+## Recipe
+
+[`recipe/`](recipe/) runs the official LittleBit code on a RunPod GPU pod:
+
+- `setup.sh`: clones SamsungLabs/LittleBit at a pinned commit, applies the patches, and installs dependencies (`transformers==4.51.*`, DeepSpeed).
+- `train.sh`: runs QAT. Defaults: Qwen3-8B, 0.55 bpw, LittleBit-2 init, SmoothSign, residual. Override settings with env vars (`MODEL_ID`, `EFF_BIT`, `EPOCHS`, `NUM_GPUS`, …).
+- `eval.sh`: measures WikiText-2/C4 perplexity and zero-shot accuracy (lm-eval).
+- `zero3_nooffload.json`: multi-GPU ZeRO-3 config without CPU offload.
+- `patches/teacher-on-gpu.patch`: keeps the teacher on GPU instead of ZeRO-3 CPU offload (`--teacher_offload False`).
+- `patches/eval-import-fix.patch`: fixes a circular import between lm-eval, transformers, and DeepSpeed in `eval.py`.
+
+```bash
+bash recipe/setup.sh
+MODEL_ID=Qwen/Qwen3-8B EFF_BIT=0.55 EPOCHS=1 bash recipe/train.sh
+CKPT=/workspace/outputs/littlebit-qwen3-8b-0.55bpw bash recipe/eval.sh
+```
+
+## License
+
+CC BY-NC 4.0 (non-commercial), inherited from the LittleBit code. Released weights
+are also subject to the base model's license (Qwen3: Apache 2.0).
+
+## Citation
+
+```bibtex
+@inproceedings{lee2025littlebit,
+  title     = {LittleBit: Ultra Low-Bit Quantization via Latent Factorization},
+  author    = {Lee, Banseok and Kim, Dongkyu and You, Youngcheon and Kim, Youngmin},
+  booktitle = {NeurIPS},
+  year      = {2025}
+}
+@inproceedings{lee2026littlebit2,
+  title     = {LittleBit-2: Maximizing the Spectral Energy Gain in Sub-1-Bit LLMs via Latent Geometry Alignment},
+  author    = {Lee, Banseok and Kim, Youngmin},
+  booktitle = {ICML},
+  year      = {2026}
+}
+```

recipe/eval.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# Evaluate a trained LittleBit checkpoint (perplexity + zero-shot).
+#   CKPT=/workspace/outputs/Qwen2.5-0.5B-1.0bpw bash eval.sh
+# Set FP_BASELINE=1 with CKPT=<hf model id> to score the original FP16 model.
+set -euo pipefail
+
+WORKDIR="${WORKDIR:-/workspace}"
+CKPT="${CKPT:?set CKPT to a checkpoint dir or HF model id}"
+ZEROSHOT="${ZEROSHOT:-boolq,piqa,hellaswag,winogrande,arc_easy,arc_challenge,openbookqa}"
+
+export HF_HOME="$WORKDIR/hf_cache"
+# shellcheck disable=SC1091
+source "$WORKDIR/venv/bin/activate"
+cd "$WORKDIR/LittleBit"
+mkdir -p "$WORKDIR/logs"
+
+EXTRA=()
+if [ "${FP_BASELINE:-0}" = "1" ]; then EXTRA+=(--fp_eval True); fi
+
+CUDA_VISIBLE_DEVICES=0 python eval.py \
+  --model_id "$CKPT" \
+  --seqlen 2048 \
+  --ppl_task wikitext2,c4 \
+  --zeroshot_task "$ZEROSHOT" \
+  "${EXTRA[@]}" 2>&1 | tee "$WORKDIR/logs/eval-$(basename "$CKPT")-$(date +%Y%m%d-%H%M).log"

recipe/patches/eval-import-fix.patch

@@ -0,0 +1,20 @@
+diff --git a/eval.py b/eval.py
+index 446b1ed..84b3355 100644
+--- a/eval.py
++++ b/eval.py
+@@ -4,6 +4,15 @@ import json
+ import torch
+ import torch._dynamo
+ import torch.nn as nn
++
++# Import deepspeed before lm_eval/transformers (as main.py does). Otherwise
++# transformers.modeling_utils imports deepspeed mid-init and deepspeed's
++# import of transformers.models.opt hits a circular import.
++try:
++    import deepspeed  # noqa: F401
++except ImportError:
++    pass
++
+ from lm_eval import evaluator
+ from tqdm import tqdm
+ from transformers import AutoTokenizer, AutoModelForCausalLM

recipe/patches/teacher-on-gpu.patch

@@ -0,0 +1,44 @@
+diff --git a/main.py b/main.py
+index d5f6810..d30ecb2 100644
+--- a/main.py
++++ b/main.py
+@@ -83,6 +83,9 @@ def get_args():
+ 
+     parser.add_argument("--use_itq", type=str2bool, default=False)
+     parser.add_argument("--itq_n_iter", type=int, default=50)
++    parser.add_argument("--teacher_offload", type=str2bool, default=True,
++                        help="True: teacher via DeepSpeed ZeRO-3 with CPU param offload (upstream). "
++                             "False: full bf16 teacher copy resident on each GPU (faster, needs VRAM).")
+ 
+     args = parser.parse_args()
+ 
+@@ -142,7 +145,29 @@ def load_student_model(args, device_map, torch_dtype):
+     return model
+ 
+ 
++def load_teacher_on_gpu(args, torch_dtype):
++    local_rank = int(os.environ.get("LOCAL_RANK", 0))
++    device = torch.device(f"cuda:{local_rank}")
++
++    teacher_model = AutoModelForCausalLM.from_pretrained(
++        args.model_id,
++        torch_dtype=torch_dtype,
++        low_cpu_mem_usage=True,
++        trust_remote_code=True,
++    ).to(device)
++    teacher_model.eval()
++    for param in teacher_model.parameters():
++        param.requires_grad = False
++    teacher_model.config.use_cache = False
++
++    logger.info(f"Teacher resident on {device} (no offload)")
++    return teacher_model
++
++
+ def load_teacher_model(args, num_gpus, torch_dtype, config_path="configs/zero3_inference.json"):
++    if not args.teacher_offload:
++        return load_teacher_on_gpu(args, torch_dtype)
++
+     with open(config_path, 'r') as f:
+         config = json.load(f)
+ 

recipe/setup.sh

@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+# One-time setup for LittleBit on a RunPod GPU pod.
+# Run from the pod terminal:  bash /workspace/runpod/setup.sh
+set -euo pipefail
+
+WORKDIR="${WORKDIR:-/workspace}"
+REPO_DIR="$WORKDIR/LittleBit"
+VENV_DIR="$WORKDIR/venv"
+
+# Keep every cache on the persistent volume so it survives pod restarts.
+export HF_HOME="$WORKDIR/hf_cache"
+export PIP_CACHE_DIR="$WORKDIR/pip_cache"
+
+echo ">> GPU check"
+nvidia-smi --query-gpu=name,memory.total --format=csv || { echo "No NVIDIA GPU found"; exit 1; }
+
+echo ">> System packages (mpi4py needs OpenMPI headers)"
+apt-get update -qq && apt-get install -y -qq libopenmpi-dev git >/dev/null
+
+if [ ! -d "$REPO_DIR" ]; then
+  echo ">> Cloning LittleBit"
+  git clone https://github.com/SamsungLabs/LittleBit.git "$REPO_DIR"
+  # Pin the commit the patches were written against.
+  git -C "$REPO_DIR" checkout -q 42d658b
+fi
+
+echo ">> Applying local patches"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+for p in "$SCRIPT_DIR"/patches/*.patch; do
+  if git -C "$REPO_DIR" apply --reverse --check "$p" 2>/dev/null; then
+    echo "   already applied: $(basename "$p")"
+  else
+    git -C "$REPO_DIR" apply "$p" && echo "   applied: $(basename "$p")"
+  fi
+done
+
+echo ">> Python venv (inherits the pod's preinstalled CUDA torch)"
+python3 -m venv --system-site-packages "$VENV_DIR"
+# shellcheck disable=SC1091
+source "$VENV_DIR/bin/activate"
+pip install -q --upgrade pip
+
+# requirements.txt pins torch==2.6.0, which would replace the pod's CUDA build.
+# Install everything except torch, then the versions the README asks for.
+grep -viE '^torch(==|$)' "$REPO_DIR/requirements.txt" > /tmp/req-notorch.txt
+pip install -q -r /tmp/req-notorch.txt
+pip install -q "transformers==4.51.*" deepspeed wandb
+
+echo ">> Sanity check"
+python - <<'PY'
+import torch, transformers, deepspeed
+print("torch", torch.__version__, "cuda", torch.version.cuda, "available", torch.cuda.is_available())
+print("transformers", transformers.__version__, "deepspeed", deepspeed.__version__)
+for i in range(torch.cuda.device_count()):
+    print(" gpu", i, torch.cuda.get_device_name(i))
+assert torch.cuda.is_available(), "CUDA not available"
+PY
+
+cat <<EOF
+
+Setup done. Next:
+  export HF_TOKEN=hf_xxx        # needed for gated models like Llama-2
+  export WANDB_API_KEY=xxx      # optional; without it, REPORT=tensorboard is used
+  bash $WORKDIR/runpod/train.sh
+EOF

recipe/train.sh

@@ -0,0 +1,70 @@
+#!/usr/bin/env bash
+# LittleBit QAT training on RunPod. Override any setting via env vars, e.g.
+#   MODEL_ID=Qwen/Qwen3-14B EFF_BIT=0.55 EPOCHS=5 bash train.sh
+set -euo pipefail
+
+WORKDIR="${WORKDIR:-/workspace}"
+REPO_DIR="$WORKDIR/LittleBit"
+
+MODEL_ID="${MODEL_ID:-Qwen/Qwen3-8B}"
+EFF_BIT="${EFF_BIT:-0.55}"                  # bits per weight: 1.0, 0.55, 0.3, 0.1 ...
+EPOCHS="${EPOCHS:-1.0}"                     # paper uses 5.0
+BATCH="${BATCH:-2}"                         # Qwen3's 151k vocab makes KD logits large
+GRAD_ACCUM="${GRAD_ACCUM:-2}"
+LR="${LR:-4e-05}"
+USE_ITQ="${USE_ITQ:-True}"                  # True = LittleBit-2 init (Joint-ITQ)
+TEACHER_OFFLOAD="${TEACHER_OFFLOAD:-False}" # False = teacher stays on GPU (faster, patched)
+NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}"
+MODEL_SLUG="$(basename "$MODEL_ID" | tr '[:upper:]' '[:lower:]')"
+RUN_NAME="${RUN_NAME:-littlebit-${MODEL_SLUG}-${EFF_BIT}bpw}"
+SAVE_DIR="${SAVE_DIR:-$WORKDIR/outputs}"
+
+if [ -z "${REPORT:-}" ]; then
+  if [ -n "${WANDB_API_KEY:-}" ]; then REPORT=wandb; else REPORT=tensorboard; fi
+fi
+
+export HF_HOME="$WORKDIR/hf_cache"
+# shellcheck disable=SC1091
+source "$WORKDIR/venv/bin/activate"
+cd "$REPO_DIR"
+mkdir -p "$WORKDIR/logs"
+LOG="$WORKDIR/logs/train-$RUN_NAME-$(date +%Y%m%d-%H%M).log"
+
+ARGS=(
+  --model_id "$MODEL_ID"
+  --dataset c4_wiki
+  --data_root "$WORKDIR/data"
+  --save_dir "$SAVE_DIR"
+  --f_name "$RUN_NAME"
+  --run_name "$RUN_NAME"
+  --num_train_epochs "$EPOCHS"
+  --per_device_train_batch_size "$BATCH"
+  --gradient_accumulation_steps "$GRAD_ACCUM"
+  --lr "$LR"
+  --warmup_ratio 0.02
+  --report "$REPORT"
+  --quant_func SmoothSign
+  --quant_mod LittleBitLinear
+  --residual True
+  --eff_bit "$EFF_BIT"
+  --kv_factor 1.0
+  --min_split_dim 8
+  --l2l_loss_scale 10.0
+  --use_itq "$USE_ITQ"
+  --teacher_offload "$TEACHER_OFFLOAD"
+)
+# Multi-GPU: ZeRO-3 sharding. Default keeps optimizer/params on GPU (fast);
+# set DS_CONFIG=configs/zero3.json for upstream CPU offload (slow, less VRAM).
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+DS_CONFIG="${DS_CONFIG:-$SCRIPT_DIR/zero3_nooffload.json}"
+if [ "$NUM_GPUS" -gt 1 ]; then
+  ARGS+=(--ds_config_path "$DS_CONFIG")
+fi
+
+# Always go through the deepspeed launcher: it sets LOCAL_RANK, which main.py
+# needs to place the student on a concrete GPU (plain `python -m main` passes
+# device_map="auto" to model.to(), which fails).
+echo ">> $RUN_NAME on $NUM_GPUS GPU(s), log: $LOG"
+deepspeed --num_gpus="$NUM_GPUS" main.py "${ARGS[@]}" 2>&1 | tee "$LOG"
+
+echo ">> Checkpoint: $SAVE_DIR/$RUN_NAME"

recipe/zero3_nooffload.json

@@ -0,0 +1,45 @@
+{
+    "bf16": {
+        "enabled": true
+    },
+
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+          "lr": "auto",
+          "betas": "auto",
+          "eps": "auto",
+          "weight_decay": "auto"
+        }
+    },
+
+    "scheduler": {
+        "type": "WarmupCosineLR",
+        "params": {
+            "total_num_steps": "auto",
+            "warmup_min_ratio": 0,
+            "warmup_num_steps": "auto",
+            "cos_min_ratio": 0.0001,
+            "warmup_type": "linear"
+        }
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 5e8,
+        "stage3_prefetch_bucket_size": 5e8,
+        "stage3_param_persistence_threshold": 1e6,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}