Upload Dockerfile with huggingface_hub

Dockerfile

@@ -0,0 +1,131 @@
+FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
+
+ARG HTM_CUDA_ARCH=sm_86
+ARG TORCH_CUDA_ARCH_LIST=8.6
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_NO_CACHE_DIR=1 \
+    PYTHONUNBUFFERED=1 \
+    CARGO_HOME=/root/.cargo \
+    RUSTUP_HOME=/root/.rustup \
+    HTM_CUDA_ARCH=${HTM_CUDA_ARCH} \
+    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
+    PATH=/root/.cargo/bin:${PATH}
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git curl ca-certificates build-essential pkg-config libssl-dev && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN curl https://sh.rustup.rs -sSf | bash -s -- -y --profile minimal --default-toolchain stable
+
+RUN pip install --upgrade pip setuptools wheel && \
+    pip install \
+      maturin \
+      huggingface_hub \
+      datasets \
+      requests \
+      pyarrow \
+      rustbpe \
+      pandas \
+      tiktoken \
+      pydantic \
+      ninja \
+      packaging \
+      einops
+
+# Mamba-3 fused CUDA kernel stack (mandatory — NO fallback allowed).
+#
+# We install PRE-BUILT manylinux wheels from the official state-spaces/mamba
+# and Dao-AILab/causal-conv1d GitHub releases. Compiling mamba_ssm from source
+# on HF Spaces' cpu-basic builder (~16GB RAM) OOMKills even with MAX_JOBS=1 —
+# nvcc on the templated selective-scan/chunk-scan kernels needs 8–12GB per TU.
+#
+# Wheel selection for base image pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel:
+#   - Python 3.11 (cp311)                       — matches PyTorch 2.6.0 image
+#   - CUDA 12.x wheels (cu12)                   — matches host CUDA 12.4
+#   - PyTorch 2.6 ABI (torch2.6)                — exact torch match
+#   - cxx11abiFALSE                             — standard PyTorch pip build
+#
+# Versions: mamba_ssm 2.3.1 (first stable with Mamba3 class) + causal_conv1d
+# 1.6.1.post4 (matching ABI). Both are CUDA-compiled, no build toolchain needed
+# on the Space builder.
+#
+# Step A: install the published v2.3.1 prebuilt wheel (compiled CUDA ops
+# for selective_scan, layernorm_gated, ssd_*, causal_conv1d, etc).
+RUN pip install \
+      'https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.1.post4/causal_conv1d-1.6.1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl' \
+      'https://github.com/state-spaces/mamba/releases/download/v2.3.1/mamba_ssm-2.3.1+cu12torch2.6cxx11abiFALSE-cp311-cp311-linux_x86_64.whl' && \
+    python -c "import importlib.metadata as m; print('installed mamba_ssm=' + m.version('mamba_ssm') + ' causal_conv1d=' + m.version('causal_conv1d'))"
+
+#
+# Step B: graft the Mamba3 class + its pure-Triton ops subtree from mamba-ssm
+# main. v2.3.1 is the latest release but Mamba3 landed post-release; the new
+# files under ops/triton/mamba3/ are ALL pure Python @triton.jit kernels with
+# zero compiled-CUDA dependencies (verified: every import in that subtree is
+# triton/torch/python — no .so files, no nvcc). So we install the v2.3.1 wheel
+# (for its compiled ops) and overlay the main-branch Mamba3 sources on top.
+#
+# This avoids the source-build OOM on the cpu-basic HF Space builder and the
+# missing-file error the smoke hit on the last attempt.
+# Download grafted mamba3 module + triton ops subtree
+RUN SITE=/opt/conda/lib/python3.11/site-packages/mamba_ssm && \
+    BASE=https://raw.githubusercontent.com/state-spaces/mamba/main && \
+    curl -fsSL "$BASE/mamba_ssm/modules/mamba3.py" -o "$SITE/modules/mamba3.py" && \
+    mkdir -p "$SITE/ops/triton/mamba3" && \
+    for f in __init__.py angle_dt.py mamba3_mimo_rotary_step.py mamba3_mimo_utils.py mamba3_siso_bwd.py mamba3_siso_combined.py mamba3_siso_fwd.py mamba3_siso_step.py utils.py; do \
+        curl -fsSL "$BASE/mamba_ssm/ops/triton/mamba3/$f" -o "$SITE/ops/triton/mamba3/$f"; \
+    done
+
+# Replace mamba_ssm/__init__.py with a minimal one that only imports Mamba3
+# (pure-Triton, works). The shipped __init__.py eagerly imports
+# selective_scan_cuda.so which has a libtorch C++ ABI mismatch on this base
+# image ("undefined symbol: _ZN3c107WarningC1E..."). Since training only needs
+# Mamba3 (grafted from main), we skip all compiled-CUDA imports.
+COPY mamba_ssm_init.py /opt/conda/lib/python3.11/site-packages/mamba_ssm/__init__.py
+
+# Structural check (no triton init — triton has no GPU on the builder)
+RUN SITE=/opt/conda/lib/python3.11/site-packages/mamba_ssm && \
+    test -f "$SITE/modules/mamba3.py" && \
+    test -f "$SITE/ops/triton/mamba3/mamba3_siso_combined.py" && \
+    test -s "$SITE/__init__.py" && \
+    echo "mamba3 graft + __init__ override verified"
+
+# Optional tilelang for MIMO path — pure-python, cheap; SISO Mamba3 works without.
+RUN pip install tilelang || echo "[dockerfile] tilelang optional install failed — continuing"
+
+# Triton version decision: FORCE 3.5.1 — the only version with both mamba3
+# APIs (set_allocator + tl.make_tensor_descriptor). torch 2.6's _inductor
+# imports AttrsDescriptor from triton.compiler.compiler which was removed in
+# triton 3.4+, but mamba_ssm/__init__.py shims AttrsDescriptor as a stub
+# before any torch._inductor import path runs, so the incompatibility is
+# neutralized. Build-time assert verifies mamba3's two required APIs.
+RUN pip install --force-reinstall --no-deps 'triton==3.5.1' && \
+    python -c "import triton; from triton import language as tl; \
+               assert hasattr(triton, 'set_allocator'), 'missing triton.set_allocator'; \
+               assert hasattr(tl, 'make_tensor_descriptor'), 'missing tl.make_tensor_descriptor'; \
+               print(f'triton={triton.__version__} set_allocator+make_tensor_descriptor OK, AttrsDescriptor shimmed in mamba_ssm/__init__.py')"
+
+WORKDIR /workspace
+COPY overlay /workspace/feather
+COPY overlay/scripts /app/scripts
+COPY entrypoint.py /app/entrypoint.py
+WORKDIR /workspace/feather
+
+RUN test -f /app/scripts/htm_gpu_micro_canary.py && \
+    python -m py_compile hydra/training.py prepare.py train.py /app/scripts/htm_gpu_micro_canary.py && \
+    bash -n scripts/run_domain_expanded_pretrain.sh
+
+RUN export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} && \
+    echo "building htm_rust GPU kernels for HTM_CUDA_ARCH=${HTM_CUDA_ARCH} TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}" && \
+    if maturin build --release --features gpu --manifest-path htm_rust/Cargo.toml; then \
+        pip install htm_rust/target/wheels/htm_rust-*.whl && \
+        python -c "import htm_rust; assert hasattr(htm_rust, 'HTMRegionGpu'), 'htm_rust missing HTMRegionGpu GPU binding'"; \
+    else \
+        echo "[dockerfile] htm_rust GPU wheel build failed; building CPU wheel so A10 compromise/fresh-eval jobs can still run with explicit CPU fallback" && \
+        rm -rf htm_rust/target/wheels && \
+        maturin build --release --manifest-path htm_rust/Cargo.toml && \
+        pip install htm_rust/target/wheels/htm_rust-*.whl && \
+        python -c "import htm_rust; assert hasattr(htm_rust, 'HTMRegion'), 'htm_rust missing CPU HTMRegion binding'"; \
+    fi
+
+CMD ["python", "/app/entrypoint.py"]