Spaces:
Runtime error
Runtime error
File size: 3,452 Bytes
5b5c422 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | #!/usr/bin/env bash
# Runtime setup for the stock pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel image.
# We avoid baking feather + mamba_ssm + htm_rust into a custom Docker image
# because build-time baking on HF's cpu-basic builder reliably corrupts CUDA
# state on h200 runtime ("Error 802: system not yet initialized" every time,
# even in a fresh python -c subprocess). Installing at runtime, on the h200
# itself, avoids that path and keeps CUDA healthy.
#
# Trade-off: ~5-8 min cold start per job vs ~1 min for a baked image. The
# training run is 12h long, so the overhead is negligible.
set -euo pipefail
echo "[runtime] $(date -u +%H:%M:%S) starting feather runtime setup on $(hostname)"
# 1. Confirm CUDA before we do anything else.
python -c 'import torch; assert torch.cuda.is_available(), "cuda unavailable at runtime start"; print("[runtime] cuda OK —", torch.cuda.get_device_name(0))'
# 2. Install system build deps (rustup/build-essential for htm_rust).
apt-get update -qq
apt-get install -y -qq --no-install-recommends git curl ca-certificates build-essential pkg-config libssl-dev
# Rust toolchain for htm_rust
curl -sSf https://sh.rustup.rs | bash -s -- -y --profile minimal --default-toolchain stable
export PATH=/root/.cargo/bin:$PATH
# 3. Install Python deps.
pip install --quiet --upgrade pip setuptools wheel
pip install --quiet \
maturin \
huggingface_hub \
requests \
pyarrow \
rustbpe \
pandas \
tiktoken \
pydantic \
ninja \
packaging \
einops
# 4. Install mamba_ssm + causal_conv1d (prebuilt wheels, matching torch2.6/cu12).
pip install --quiet \
'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'
# 5. Graft Mamba3 from main (pure Triton, not in v2.3.1 release).
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 the eager-init __init__.py with our minimal version.
cp /workspace/feather/hf_jobs/feather_h200_image/mamba_ssm_init.py "$SITE/__init__.py"
# 6. Confirm CUDA still works after all installs.
python -c 'import torch; assert torch.cuda.is_available(), "cuda broken by installs"; print("[runtime] cuda OK after deps —", torch.cuda.get_device_name(0))'
# 7. Build + install htm_rust with sm_90 PTX (h200 arch).
cd /workspace/feather
export HTM_CUDA_ARCH=sm_90
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH:-}
maturin build --release --features gpu --manifest-path htm_rust/Cargo.toml 2>&1 | tail -5
pip install --quiet htm_rust/target/wheels/htm_rust-*.whl
# 8. Sanity: cuda still alive after htm_rust install.
python -c 'import torch; assert torch.cuda.is_available(), "cuda broken by htm_rust"; import htm_rust; print("[runtime] htm_rust OK, cuda OK")'
echo "[runtime] $(date -u +%H:%M:%S) runtime setup complete"
|