OpenThoughts-Agent: Train Small AI Models with HPC Scale

OpenThoughts-Agent is a research toolkit providing the recipes and infrastructure needed to train small, high-capability agent models. It offers comprehensive support for high-throughput data generation (Datagen), supervised fine-tuning (SFT), and reinforcement learning (RL). The system uses a modular HPC launcher to connect code directly with high-performance compute clusters, bypassing rigid integrations by using customizable generation scripts and training configurations compatible with Llama Factory, vLLM, and Harbor. Harbor manages containerized tool deployment across diverse HPC environments, enabling remote container runtimes without requiring a workflow re-architecture.

Installing OpenThoughts-Agent

Environment Setup

1. Start with a clean Python 3.12 virtual environment using conda, mamba, or a similar manager.
2. From the project root, install the core HPC and data infrastructure:

pip install .

OpenThoughts-Agent has a broad dependency footprint. Using a fast package manager like uv is recommended to maintain environment stability.

Optional Extensions

• HPC data generation runtime (requires a Ray cluster and vLLM service):

pip install .[datagen]

• SweSmith-specific data generation tools (adds custom utilities on top of the base datagen extension):

pip install .[datagen,datagen-swesmith]

Installing for SFT

1. Initialize the git submodule:

git submodule update --init --recursive sft/llamafactory

2. Install LLaMA Factory directly from the submodule. Note that OT-Agent does not provide a standalone [sft] extension.

cd sft/llamafactory
pip install -e .[train,liger-kernel,deepspeed]  # Select components as required
cd -

Training configurations are located in sft/lf_configs/**. Refer to sft/llamafactory/README.md for parameter details and dependency requirements.

Data Utilities

Dataset tool documentation is available in data/README.md. After installing the datagen extension, specific generators may require additional dependencies.

CPP and Compiler Notes

Several launch modes compile CUDA/C++ extensions (such as flash-infer, flash-attn, and triton) on the fly. The build process is sensitive to compiler and CUDA versions. Ensure your toolchain matches the PyTorch CUDA version, which you can verify with:

python - <<<'import torch; print(torch.version.cuda)'

Toolchain Configuration Options

1. Cluster Modules: If your HPC environment uses modules, load them directly.

module load gcc/14.2.0
module load cuda/12.8

2. Container Terminal: Bind your workspace to a pre-built CUDA image for a predictable toolchain.

singularity shell --nv \
  --bind $SCRATCH/ot-agent \
  $SCRATCH/cuda-img/cuda-cudnn-12.8-ubuntu22.sif

3. Conda Toolchain: If modules or containers are unavailable, install compilers and sysroot via mamba. Pin package versions to prevent ABI mismatches.

mamba install -c conda-forge c-compiler cxx-compiler -y
mamba install -c conda-forge gcc_linux-64 gxx_linux-64 sysroot_linux-64 -y
mamba install -c conda-forge libstdcxx-ng=12 libgcc-ng=12 gcc_impl_linux-64 \
    gxx_impl_linux-64 sysroot_linux-64 -y

4. Environment Variables: Map CUDA and GCC tools to their local paths.

GCC_ROOT="$(dirname "$(dirname "$(which gcc)")")"
export CUDA_HOME=/usr/local/cuda
export CPATH="$CUDA_HOME/include${CPATH:+:$CPATH}"
export LIBRARY_PATH="$CUDA_HOME/lib64${LIBRARY_PATH:+:$LIBRARY_PATH}"
export LD_LIBRARY_PATH="$GCC_ROOT/lib64:$GCC_ROOT/lib:$CUDA_HOME/lib64${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
export PATH="$CUDA_HOME/bin${PATH:+:$PATH}"

5. Manual Compilations: JIT components build automatically upon import once the toolchain is configured. However, packages like flash-attn often require manual compilation. Install these last.

UV_COMPILE_THREADS=4 MAX_JOBS=4 NVCC_THREADS=4 TORCH_CUDA_ARCH_LIST="9.0" \
  pip install -v --no-build-isolation "flash-attn==2.8.1"

Keys and API Configuration

Most scripts require credentials—including Hugging Face tokens, Daytona keys, W&B API keys, or Supabase credentials. Store these in a private .env file excluded from version control and point the system to it:

export DC_AGENT_SECRET_ENV=/secure/path/to/my_dc_agent_secrets.env

The file should contain standard exports (e.g., export HF_TOKEN=...). Launchers and helper scripts reference the DC_AGENT_SECRET_ENV variable.

Launching a Data Generation Job

1. Configure your cluster environment. For TACC/Vista systems, follow the checklist in hpc/README.md and use hpc/dotenv/tacc.env as a base.
2. Activate the environment and load the configuration:

source hpc/dotenv/<your-cluster>.env
eval "$DCFT_ACTIVATE_ENV"
cd "$DCFT"

3. Select or author a data generation script in data/.... The script must implement the BaseDataGenerator class (see data/generation/base.py for reference).
4. Execute the launcher from the login node:

python -m hpc.launch \
  --job_type datagen \
  --datagen_script data/<dataset>/generate.py \
  --datagen_target_repo <org/dataset-tasks> \
  --datagen_engine vllm_local \
  --datagen_extra_args "--stage both --limit 200" \
  --experiments_dir "$DCFT/experiments" \
  --time_limit 12:00:00

5. To generate traces, include:

--enable_trace_gen \
--trace_target_repo <org/dataset-traces> \
--trace_harbor_config path/to/harbor_job.yaml

6. The launcher generates sbatch scripts and configurations under $experiments_dir. Use --dry_run to verify the setup before submission.

Launching an SFT Job

1. Update and install the SFT submodule:

git submodule update --init --remote sft/llamafactory
cd sft/llamafactory
pip install -e .[train,liger-kernel,deepspeed]
cd -

2. Configure the cluster environment as you would for data generation.
3. Select a training configuration YAML from sft/lf_configs.
4. Run from the login node:

python -m hpc.launch \
  --job_type train \
  --train_config_path sft/lf_configs/<path-to-config>.yaml \
  --dataset <org/dataset> \
  --num_nodes 8 \
  --time_limit 24:00:00 \
  --experiments_dir "$DCFT/experiments"

5. Override LLaMA Factory arguments using --train_extra_args. The launcher will generate the job-specific YAML and sbatch script, submitting the job automatically.

Custom Cluster Configuration

1. Create a dotenv file in hpc/dotenv/ (e.g., my_cluster.env). Define the following:
   • DCFT: Path to the OpenThoughts-Agent directory.
   • DCFT_ACTIVATE_ENV: Command to activate the Python environment.
   • Storage paths: EXPERIMENTS_DIR, DATASETS_DIR, MODELS_DIR.
   • SIF paths for Apptainer images.
2. Define HPC_NAME and related fields in the dotenv or pass them as arguments (e.g., --account, --partition, --gpus_per_node).
3. Create an sbatch template in hpc/sbatch_data/. Use placeholders like {time_limit}, {job_name}, and {experiments_dir}, which the launcher populates during execution.
4. Add the required template to hpc/sbatch_data_requirements.json to enable validation.
5. Perform a dry run to verify the configuration:

python -m hpc.launch \
  --job_type datagen \
  --datagen_script data/<dataset>/generate.py \
  --datagen_target_repo test-org/test-dataset \
  --experiments_dir "$DCFT/experiments" \
  --dry_run

6. If the cluster requires specialized logic (such as proxy settings or unique login/compute node behavior), extend hpc/hpc.py or modify hpc/launch.py. The JURECA/JUWELS implementation is a useful reference for complex network configurations.