| FROM nvidia/cuda:12.9.1-devel-ubuntu24.04 |
| ENV DEBIAN_FRONTEND=noninteractive |
| RUN apt-get update && apt-get install -y --no-install-recommends python3 python3-dev python3-venv git ca-certificates build-essential |
| RUN mkdir -p /home/dev/pretrain_suffix_min/vendor && python3 -m venv /home/dev/venvs/main |
| ENV PATH=/home/dev/venvs/main/bin:$PATH |
| ENV HF_HOME=/home/dev/.cache/huggingface |
| ENV TRITON_CACHE_DIR=/home/dev/.cache/triton |
| ENV TORCHINDUCTOR_CACHE_DIR=/home/dev/.cache/torchinductor |
| RUN pip install --no-cache-dir setuptools wheel packaging |
| RUN pip install --no-cache-dir torch==2.13.0+cu129 --index-url https://download.pytorch.org/whl/cu129 |
| RUN git clone https://github.com/fla-org/flash-linear-attention /home/dev/pretrain_suffix_min/vendor/flash-linear-attention && git -C /home/dev/pretrain_suffix_min/vendor/flash-linear-attention checkout 7b5138a92133132cfc5f80455c9c0920d4fe0278 |
| RUN git clone https://github.com/EleutherAI/lm-evaluation-harness /home/dev/pretrain_suffix_min/vendor/lm-evaluation-harness && git -C /home/dev/pretrain_suffix_min/vendor/lm-evaluation-harness checkout d6de81643928d653435c431bae19945d41d32520 |
| RUN pip install --no-cache-dir --no-build-isolation -e /home/dev/pretrain_suffix_min/vendor/flash-linear-attention -e /home/dev/pretrain_suffix_min/vendor/lm-evaluation-harness transformers==4.57.6 datasets==5.0.1 numpy==2.5.3 pytest==9.1.1 matplotlib==3.11.2 |
| |
| ENV MAX_JOBS=8 |
| ENV MAMBA_FORCE_BUILD=TRUE |
| ENV CAUSAL_CONV1D_FORCE_BUILD=TRUE |
| ENV TORCH_CUDA_ARCH_LIST=9.0 |
| RUN pip install --no-cache-dir ninja einops && pip install --no-cache-dir --no-build-isolation causal-conv1d==1.5.2 |
| RUN git clone https://github.com/state-spaces/mamba /home/dev/pretrain_mamba2_outer_vendor && git -C /home/dev/pretrain_mamba2_outer_vendor checkout d7b1ceb3c367ec9022925e812f507bcf706937c6 |
| RUN pip install --no-cache-dir --no-build-isolation -e /home/dev/pretrain_mamba2_outer_vendor |
| RUN mkdir -p /home/dev/pretrain_mamba2_outer |
| COPY source.tar.gz /home/dev/source.tar.gz |
| RUN tar -xzf /home/dev/source.tar.gz -C /home/dev/pretrain_mamba2_outer |
| WORKDIR /home/dev/pretrain_mamba2_outer |
| ENV PYTHONPATH=/home/dev/pretrain_mamba2_outer |
| ENV FLA_CI_ENV=0 |
| ENV CUBLAS_WORKSPACE_CONFIG=:4096:8 |
| CMD ["python", "-u", "-m", "pretraining_mamba2_outer.debug_gate", "--out", "/home/dev/mamba2-debug"] |
|
|
