Skip to content

Commit b7c9c31

Browse files
uralikuralik
authored
pixi manifest and install command (#726)
1 parent 759ab71 commit b7c9c31

File tree

3 files changed

+149
-1
lines changed

3 files changed

+149
-1
lines changed

‎README.md‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ You can also find our notebook tutorials (coming soon)
3232

3333
torchforge requires PyTorch 2.9.0 with [Monarch](https://github.com/meta-pytorch/monarch), [vLLM](https://github.com/vllm-project/vllm), and [torchtitan](https://github.com/pytorch/torchtitan).
3434

35+
### Conda
36+
3537
Install torchforge with:
3638

3739
```bash
@@ -40,7 +42,7 @@ conda activate forge
4042
./scripts/install.sh
4143
```
4244

43-
### ROCm Installation
45+
#### ROCm Installation
4446

4547
ROCm users can install with the dedicated script:
4648

@@ -61,6 +63,20 @@ The install script installs system dependencies along with torchforge. Note that
6163

6264
Optional: By default, the packages installation uses conda. If you want to install system packages on the target machine instead of conda, you can pass the `--use-sudo` flag to the installation script: `./scripts/install.sh --use-sudo`.
6365

66+
### Pixi
67+
68+
Pixi combines benefits of uv with access to conda forge for system dependencies. [pixi.toml](./pixi.toml) provides a manifest with build tasks with `install` as a the combined install all task.
69+
70+
Install pixi:
71+
```
72+
curl -fsSL https://pixi.sh/install.sh | bash
73+
```
74+
75+
Install torchforge with pixi:
76+
```
77+
pixi run install
78+
```
79+
6480
> **Note:** We are actively working on enabling pure `uv` installation. Currently, Conda is the recommended approach. `uv` support is not fully working at the moment but is being tracked in [issue #494](https://github.com/meta-pytorch/torchforge/issues/494).
6581
6682
After install, you can run the following command and should see output confirming GRPO training is running (you need a minimum 3 GPU devices):

‎pixi.toml‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
# Pixi configuration for Forge
2+
# This replaces the conda-based installation with pixi environment manager
3+
#
4+
# Copyright (c) Meta Platforms, Inc. and affiliates.
5+
# All rights reserved.
6+
7+
[workspace]
8+
name = "forge"
9+
version = "0.1.0"
10+
description = "A PyTorch native platform for post-training generative AI models"
11+
channels = ["conda-forge", "pytorch", "nvidia"]
12+
platforms = ["linux-64"]
13+
14+
[system-requirements]
15+
linux = "5.10"
16+
cuda = "12.8"
17+
18+
# System dependencies from conda-forge
19+
[dependencies]
20+
python = "3.10.*"
21+
openssl = "*"
22+
# RDMA and InfiniBand libraries (equivalent to rdma-core, libibverbs)
23+
rdma-core = "*"
24+
# Git for installing from git repos
25+
git = "*"
26+
# Use uv for fast Python package installation
27+
uv = ">=0.1.0"
28+
29+
# Development dependencies (equivalent to [dev] extras)
30+
[feature.dev.dependencies]
31+
pre-commit = "*"
32+
pytest = "*"
33+
pytest-cov = "*"
34+
pytest-timeout = "*"
35+
tensorboard = "*"
36+
tomli = ">=1.1.0"
37+
anyio = "*"
38+
pytest-asyncio = "*"
39+
40+
# Documentation dependencies (equivalent to [docs] extras)
41+
[feature.docs.dependencies]
42+
sphinx = "==7.2.6"
43+
matplotlib = "*"
44+
45+
[tasks]
46+
# Installation tasks using uv for faster package management
47+
# uv will use the configuration from pyproject.toml [tool.uv] section
48+
install-pytorch = """
49+
uv pip install torch==2.9.0 --index-url https://download.pytorch.org/whl/cu128
50+
"""
51+
52+
# CRITICAL: Install vLLM dependencies BEFORE installing vLLM itself
53+
install-vllm-reqs = """
54+
uv pip install -r .github/packaging/vllm_reqs_12_8.txt
55+
"""
56+
57+
install-vllm-deps = """
58+
uv pip install six && \
59+
uv pip install "setuptools<80"
60+
"""
61+
62+
# Install vLLM after all dependencies are resolved
63+
install-vllm = """
64+
uv pip install vllm --no-cache-dir --index-url https://download.pytorch.org/whl/preview/forge
65+
"""
66+
67+
install-torchstore = """
68+
uv pip install "git+https://github.com/meta-pytorch/torchstore.git@no-monarch-2025.12.17"
69+
"""
70+
71+
install-forge = """
72+
uv pip install -e ".[dev]"
73+
"""
74+
75+
# Full installation task - follows exact sequence from install.sh
76+
install = { depends-on = ["install-pytorch", "install-vllm-reqs", "install-vllm-deps", "install-vllm", "install-torchstore", "install-forge"] }
77+
78+
# Test installation task
79+
test-install = """
80+
python -c "import torch; print(f'PyTorch {torch.__version__} (CUDA: {torch.cuda.is_available()})')" && \
81+
python -c "import vllm; print('vLLM imported successfully')" && \
82+
python -c "import forge; print('forge imported successfully')" || true
83+
"""
84+
85+
# Define environments
86+
[environments]
87+
default = { solve-group = "default" }
88+
dev = { features = ["dev"], solve-group = "default" }
89+
docs = { features = ["docs"], solve-group = "default" }
90+
91+
[target.linux-64.activation]
92+
scripts = ["scripts/pixi_cuda_env.sh"]

‎scripts/pixi_cuda_env.sh‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
#!/bin/bash
2+
# Copyright (c) Meta Platforms, Inc. and affiliates.
3+
# All rights reserved.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
# CUDA environment activation script for Pixi
9+
# This script is automatically sourced when the pixi environment is activated
10+
11+
# CUDA environment variables
12+
export CUDA_VERSION=12.8
13+
export NVCC=/usr/local/cuda-${CUDA_VERSION}/bin/nvcc
14+
export CUDA_NVCC_EXECUTABLE=/usr/local/cuda-${CUDA_VERSION}/bin/nvcc
15+
export CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}
16+
export PATH="${CUDA_HOME}/bin:$PATH"
17+
export CUDA_INCLUDE_DIRS=$CUDA_HOME/include
18+
export CUDA_CUDART_LIBRARY=$CUDA_HOME/lib64/libcudart.so
19+
20+
# DO NOT set LD_LIBRARY_PATH globally - it breaks system tools
21+
# Instead, we use python wrapper functions below to set it only for Python processes
22+
23+
# Define python wrappers that set LD_LIBRARY_PATH only for the launched process
24+
# Priority: system CUDA driver (/usr/lib64) > CUDA toolkit > conda/pixi libs
25+
# Use system CUDA driver (newer) instead of compat (outdated stub that may be incompatible)
26+
python() {
27+
LD_LIBRARY_PATH="/usr/lib64:/usr/local/cuda-12.8/lib64:${CONDA_PREFIX}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" \
28+
command python "$@"
29+
}
30+
python3() {
31+
LD_LIBRARY_PATH="/usr/lib64:/usr/local/cuda-12.8/lib64:${CONDA_PREFIX}/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" \
32+
command python3 "$@"
33+
}
34+
35+
# Export functions to subshells when possible (best-effort, shell-dependent)
36+
if [ -n "${BASH_VERSION:-}" ]; then
37+
export -f python python3 2>/dev/null || true
38+
elif [ -n "${ZSH_VERSION:-}" ]; then
39+
typeset -fx python python3 >/dev/null 2>&1 || true
40+
fi

0 commit comments

Comments
 (0)
⚡