Add support for UV

.github/workflows/gpu_test.yaml

@@ -38,7 +38,8 @@ jobs:
       - name: Update pip
         run: python -m pip install --upgrade pip
       - name: Install torchforge
-        run: ./scripts/install.sh
+        # run: ./scripts/install.sh
+        run: pip install uv && uv pip install .
       - name: Run unit tests with coverage
         # TODO add all tests
         run: pytest tests/unit_tests --cov=. --cov-report=xml --durations=20 -vv

pyproject.toml

@@ -11,18 +11,22 @@ authors = [
 keywords = ["pytorch", "training", "llm"]
 dependencies = [
     # PyTorch
+    "torch==2.9.0",
     "torchdata>=0.8.0",
-    "torchtitan",
+    "torchtitan==0.2.0",
+    "torchmonarch==0.1.2",
+    "torchstore==0.1.1",
     # vLLM
-    # TODO: pin specific vllm version
-    #"vllm==0.10.0",
+    "vllm",
     # Hugging Face integrations
     "datasets>=2.21.0",
     "tokenizers",
     # Miscellaneous
     "omegaconf",
     "wandb",
     "hf_transfer",
+    "six",
+    "setuptools<80",
 ]
 dynamic = ["version"]
 
@@ -69,23 +73,18 @@ members = [
 ]
 
 # pytorch
-# TODO: get auto backend to work
 [[tool.uv.index]]
-name = "pytorch-nightly-cu129"
-url  = "https://download.pytorch.org/whl/nightly/cu129"
-#explicit = true
+name = "pytorch-cu128"
+url  = "https://download.pytorch.org/whl/cu128"
 
 # vllm
-# [[tool.uv.index]]
-# name = "vllm-nightly"
-# url = "https://wheels.vllm.ai/nightly"
-# explicit = true
+[[tool.uv.index]]
+name = "vllm-forge"
+url = "https://download.pytorch.org/whl/preview/forge"
 
 [tool.uv.sources]
-torchtitan = { index = "pytorch-nightly-cu129" }
-torch = { index = "pytorch-nightly-cu129" }
-torchstore = { git = "ssh://[email protected]/meta-pytorch/torchstore.git" }
-#vllm = { index = "vllm-nightly" }
+torch = { index = "pytorch-cu128" }
+vllm = { index = "vllm-forge" }
 
 [tool.uv]
 # TODO: revert to stricter default uv strategy