Phoenix & Grok ML Setup

Phoenix & Grok ML Setup

The Phoenix ranking and retrieval models are built using JAX and Haiku, leveraging the transformer architecture originally developed for Grok-1. This guide will walk you through setting up the Python environment, installing the required machine learning frameworks, and verifying your installation.

Prerequisites

Before beginning, ensure your environment meets the following requirements:

• Python 3.10+: JAX and Haiku perform best on recent Python versions.
• CUDA 12.0+ (Optional but Recommended): For high-performance ranking, an NVIDIA GPU with updated drivers is required.
• Virtual Environment: It is highly recommended to use venv or conda to manage dependencies.

1. Create a Virtual Environment

Initialize a clean environment to avoid dependency conflicts with other Python projects.

python3 -m venv venv
source venv/bin/activate

2. Install JAX

JAX is the core engine for Phoenix. The installation command varies based on whether you are using a CPU or a GPU.

For GPU (NVIDIA):

pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

For CPU only:

pip install --upgrade "jax[cpu]"

3. Install Haiku and Core Dependencies

Phoenix uses Haiku as a neural network library on top of JAX. You will also need numpy for data handling and pytest for running verification tests.

pip install dm-haiku numpy pytest

4. Configure the Grok Transformer

The Phoenix model imports the Transformer implementation from the grok module. Ensure that the grok/ directory from the Grok-1 open source release is present in your python path.

If you are working within the x-algorithm repository, the paths are pre-configured. If you are importing these modules into a custom project, set your PYTHONPATH:

export PYTHONPATH=$PYTHONPATH:/path/to/x-algorithm/phoenix:/path/to/x-algorithm/grok

5. Initialize the Retrieval Model

To use the Phoenix retrieval system in your own scripts, you must define a PhoenixRetrievalModelConfig. This configuration specifies the transformer dimensions and embedding sizes used during candidate retrieval.

from phoenix.recsys_retrieval_model import PhoenixRetrievalModelConfig
from grok import TransformerConfig

# Define the configuration
config = PhoenixRetrievalModelConfig(
    emb_size=64,
    history_seq_len=128,
    candidate_seq_len=32,
    model=TransformerConfig(
        emb_size=64,
        widening_factor=2,
        key_size=32,
        num_q_heads=2,
        num_kv_heads=2,
        num_layers=1,
    )
).initialize()

# Create the model instance
model = config.make()

6. Verify Installation

The repository includes a suite of tests to ensure the JAX operations and attention masks are functioning correctly. Run these tests to confirm your environment is ready.

# Run model architecture tests
pytest phoenix/test_recsys_model.py

# Run retrieval model tests
pytest phoenix/test_recsys_retrieval_model.py

Troubleshooting

• JAX/CUDA Mismatch: If you see PTX toolchain not found, ensure your LD_LIBRARY_PATH includes the paths to your CUDA installation.
• ModuleNotFoundError: 'grok': Ensure the grok/ directory is in the same parent folder as phoenix/ or explicitly added to your PYTHONPATH.
• Memory Errors: The Grok-based transformer can be memory-intensive. If running on a local GPU, try reducing the batch_size in your inference runners.