juiceb0xc0de/gemma-4-e2b-it-cu128-torch291-py311

juiceboxdocks/gemma-4-e2b-it-base — Docker Image Card

Not a data dataset. This repo is a landing page for the Docker image juiceboxdocks/gemma-4-e2b-it-base. If you're here because you're struggling to get Gemma-4-E2B-it running with the right CUDA/PyTorch stack, just pull the image and skip the environment setup entirely.

Quick start

docker pull juiceboxdocks/gemma-4-e2b-it-base:cu128-torch291-py311

docker run --rm -it --gpus all --ipc=host \
  -v gemma-workspace:/workspace \
  -e HF_TOKEN \
  juiceboxdocks/gemma-4-e2b-it-base:cu128-torch291-py311

That drops you at a bash shell in /workspace with everything already installed. Nothing downloads automatically. The image contains no model weights.

GPaaS

Don't want to deal with Docker? Deploy directly on RunPod:

RunPod template: gemma-4-e2b-it-base cu128-torch291-py311

Image tags

Docker Hub: `juiceboxdocks/gemma-4-e2b-it-base`

What's in the image

Base: nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04

Stack:

Pinned source trees (editable installs):

Source trees live at /opt/src/ and are installed as editable (pip install -e).

Dev tools included: tmux, Git, nvcc, cuda-gdb, compute-sanitizer, Nsight Systems (nsys), btop, nvitop, W&B, modern hf CLI, Prometheus client, NVTX, libaio-dev, Ninja, numactl.

Caches all live under `/workspace` — mount a persistent volume there and Triton kernels, HuggingFace downloads, and pip cache survive restarts.

Explicitly NOT included: model weights, checkpoints, datasets, credentials, automatic downloads, automatic jobs.

Triton cache

TRITON_CACHE_DIR is set to /workspace/.cache/triton. With a persistent mount, compiled kernels survive container restarts so you skip recompilation after the first run.

Warm the cache without weights:

python /opt/image-tools/warm_triton_cache.py

Run it twice on the same GPU + volume and look for new_files: 0 on the second pass — that's a real cache hit. A100 (sm_80) and H100 (sm_90) compile different binaries; both coexist under the same cache root.

GPU smoke test

No weights required:

python /opt/image-tools/gpu_smoke.py

Reports GPU/driver/stack, imports the SAE trainer, registers the Triton attention adapter, and runs forward + backward comparisons against PyTorch references for:

• —Gemma full causal attention: D=512, 32 query / 4 KV heads
• —Gemma sliding-window attention: D=256, 32 query / 16 KV heads, window 512

Load Gemma-4-E2B-it (after pulling weights)

Once you're inside the container with HF_TOKEN set:

import torch
from transformers import AutoModelForImageTextToText, AutoProcessor

model = AutoModelForImageTextToText.from_pretrained(
    "google/gemma-4-E2B-it",
    dtype=torch.bfloat16,
    device_map="auto",
)
processor = AutoProcessor.from_pretrained("google/gemma-4-E2B-it")
print(type(model).__name__)

Needs ~4–5 GB VRAM in bf16. A single A100/H100/RTX 3090+ is fine.

Extend

FROM juiceboxdocks/gemma-4-e2b-it-base:cu128-torch291-py311

RUN python -m pip install --no-cache-dir jupyterlab

COPY . /workspace/my-project
WORKDIR /workspace/my-project
CMD ["python", "train.py"]

No VOLUME directive on /workspace means your downstream build-time writes persist into the layer.

Sources

• —GitHub (image repo): JuiceB0xC0de/gemma-4-e2b-it-cu128-torch291-py311
• —Docker Hub: juiceboxdocks/gemma-4-e2b-it-base
• —SAE atlas trained with this image: juiceb0xc0de/gemma-4-e2b-it-SAE