Fix GPU-less CPU use on Linux with specific Torch indexes

This commit is contained in:
Luke Stanley
2025-10-21 19:52:21 +00:00
committed by Andrej Karpathy
parent 32571664b1
commit dd9387b362
3 changed files with 29 additions and 13 deletions

View File

@@ -14,7 +14,7 @@ NANOCHAT_BASE_DIR="$HOME/.cache/nanochat"
mkdir -p $NANOCHAT_BASE_DIR mkdir -p $NANOCHAT_BASE_DIR
command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
[ -d ".venv" ] || uv venv [ -d ".venv" ] || uv venv
uv sync uv sync --extra cpu
source .venv/bin/activate source .venv/bin/activate
if [ -z "$WANDB_RUN" ]; then if [ -z "$WANDB_RUN" ]; then
WANDB_RUN=dummy WANDB_RUN=dummy

View File

@@ -44,19 +44,35 @@ python_files = ["test_*.py"]
python_classes = ["Test*"] python_classes = ["Test*"]
python_functions = ["test_*"] python_functions = ["test_*"]
# target torch to cuda 12.8 # target torch to cuda 12.8 or CPU
[tool.uv.sources] [tool.uv.sources]
torch = [ torch = [
{ index = "pytorch-cpu", marker = "sys_platform != 'linux'" }, { index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }, { index = "pytorch-cu128", extra = "gpu" },
] ]
[[tool.uv.index]] [[tool.uv.index]]
name = "pytorch-cpu" name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu" url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true explicit = true
[[tool.uv.index]] [project.optional-dependencies]
name = "pytorch-cu128" cpu = [
url = "https://download.pytorch.org/whl/cu128" "torch>=2.8.0",
explicit = true ]
gpu = [
"torch>=2.8.0",
]
[tool.uv]
conflicts = [
[
{ extra = "cpu" },
{ extra = "gpu" },
],
]

View File

@@ -23,7 +23,7 @@ command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
# create a .venv local virtual environment (if it doesn't exist) # create a .venv local virtual environment (if it doesn't exist)
[ -d ".venv" ] || uv venv [ -d ".venv" ] || uv venv
# install the repo dependencies # install the repo dependencies
uv sync uv sync --extra gpu
# activate venv so that `python` uses the project's venv instead of system python # activate venv so that `python` uses the project's venv instead of system python
source .venv/bin/activate source .venv/bin/activate