DOCSEARCHbyAlgolia
Submit Docs

Oat

repository·main·Indexed 20 days ago

https://github.com/sail-sg/oat

An efficient, distributed framework for running online LLM alignment algorithms using an Actor-Learner-Oracle architecture. It supports reinforcement learning and preference learning workflows, including DPO, IPO, SLiC, and SimPO, and formalizes alignment as a Contextual Dueling Bandit (CDB) problem. The framework integrates vLLM for response sampling, DeepSpeed for memory efficiency, and Mosec for remote reward oracle serving.

Tokens
10.7K
Snippets
28
Records
35
Agent score
71%

What's inside oat-llm

  1. How LLM alignment is modeled as Contextual Dueling Bandits (CDB)

    main

    Oat formalizes LLM alignment as a Contextual Dueling Bandit (CDB) problem. In this framework, an agent (the LLM policy, potentially with a built-in reward model) interacts with an environment (humans or preference oracles) to achieve one of two objectives:

    1. Explore & Exploit: Minimizing cumulative regret.
    2. Best Arm Identification (BAI): Minimizing anytime regret.

    This framework allows existing alignment paradigms to be viewed through the lens of bandit problems, enabling sample-efficient alignment approaches like Thompson sampling.

  2. How Oat's Actor-Learner-Oracle architecture works

    main

    Oat uses a distributed architecture designed for high-efficiency online LLM alignment:

    • Actor: Uses vLLM to accelerate online response sampling.
    • Learner: Uses DeepSpeed ZeRO strategies to optimize memory efficiency during training.
    • Oracle: A model-based oracle served via Mosec as a remote service. It supports dynamic batching, data parallelism, and pipeline parallelism. If your cluster is k8s-managed, you can host the reward oracle as a remote service and assign it a cluster IP for easier access.

  3. Oracle simulation types in Oat

    main

    Oat provides several ways to simulate preference, reward, or verification feedback depending on your compute resources:

    • Verifiable rewards: Uses rule-based functions (e.g., for math reasoning).
    • Lightweight reward models: Runs directly within the actor's process, suitable for quick testing on as few as two GPUs.
    • Remote reward models: Larger models served remotely to leverage additional compute and memory.
    • LLM-as-a-judge: Queries the OpenAI API for model-based pairwise ranking.

  4. How to extend oat with custom reward oracles

    main
    The oat framework is designed to be extensible for new reasoning tasks. To add support for a new task (e.g., a specific math or logic puzzle), you need to implement a new reward oracle. This is conceptually similar to adding a new environment in traditional Reinforcement Learning. Once the oracle is implemented (e.g., in oat/oracles/), it can be passed to the training scripts via the --oracle flag.

  5. Install Oat in editable mode for development

    main

    If you are developing on oat locally, clone the repository and install it in editable mode along with the required vllm version.

    git clone git@github.com:sail-sg/oat.git
cd oat
pip install vllm==0.8.4 && pip install -e .

  7. Install Oat via PyPI

    main

    To install oat-llm in a Python environment (recommended version 3.10), install vllm==0.8.4 first, then install oat-llm using pip.

    npip install vllm==0.8.4 && pip install -U oat-llm

  8. Use dry run mode to debug OOM issues

    main

    To prevent Out-of-Memory (OOM) errors during training, you can enable dry run mode. This mode replaces real training data with dummy data of specific lengths within the dataset's __getitem__ method, allowing you to test if your specified context length will fit in your GPU memory.

    This feature is supported for both SFT (Supervised Fine-Tuning) and RL (Reinforcement Learning) training workflows. It has been verified on A100-40G GPUs.

    # Example usage pattern (refer to specific example commands in the directory)
python train.py --dry_run --dry_run_prompt_len 2048 --dry_run_response_len 512

  9. Host the reward oracle as a remote service on Kubernetes

    main

    If you are using a Kubernetes-managed cluster, you can host the reward oracle as a remote service. This allows you to assign it a cluster IP for easier access by multiple experiment pods. You can scale this by repeating the server startup step to create multiple instances, enabling parallel experiments.

    Follow these steps to set up the remote reward oracle:

    1. Create the Kubernetes service: Apply the service definition to your cluster.
    2. Deploy the serving pod: Apply the k8s/serving.yaml configuration to start your job/pod. Note: Ensure you update the path to the readiness probe script within the YAML file.
    3. Start the remote server: Once inside the pod, execute the remote server module.
    4. Run the experiment: Start your experiment using the remote preference oracle and point it to the service URL.
    # 1) Create the service:
kubectl create -f k8s/rm-service.yaml

# 2a) Start your job/pod with `k8s/serving.yaml` applied.
# (Ensure you change the path to the readiness probe script in the yaml)

# 2b) Inside the pod, start the remote server:
MOSEC_LOG_LEVEL=debug python -m oat.oracles.remote.server

# 3) With this being set up, start your experiment:
python -m oat.experiment.main \
    --preference_oracle remote \
    --remote_rm_url http://remote-rm \
    # other flags...

  10. Configure direct optimizers in oat

    main

    You can run various direct preference optimizers by setting the --dap-algo flag. Supported algorithms include DPO, IPO, SLiC, and SimPO. When changing the algorithm, ensure you adjust the associated --beta hyper-parameter accordingly.

    python -m oat.experiment.main \
    --dap-algo IPO \
    --beta 0.1

  11. Use a locally hosted Mosec service as a preference oracle

    main

    To use a preference oracle served via Mosec on your local machine, follow these steps:

    1. Start the Mosec service: Run the server and specify the CUDA devices to use for the parallel workers. For example, to use the first 4 GPUs:
      MOSEC_LOG_LEVEL=debug python -m oat.oracles.remote.server --cuda-devices 0,1,2,3
    2. Run the experiment: In a new terminal, start your training script using --preference-oracle remote and point to the local URL using --remote-rm-url (e.g., http://0.0.0.0:8000).
    # 1. Start service
MOSEC_LOG_LEVEL=debug python -m oat.oracles.remote.server --cuda-devices 0,1,2,3

# 2. Run experiment
python -m oat.experiment.main \
    --flash-attn \
    --gradient-checkpointing \
    --rnd-seed \
    --gpus 8 \
    --dap-algo DPO \
    --beta 0.1 \
    --preference-oracle remote \
    --remote-rm-url http://0.0.0.0:8000 \
    --pretrain trl-lib/pythia-1b-deduped-tldr-sft \
    --prompt-data lkevinzc/tldr-with-sft-reference \
    --input-key prompt \
    --output-key pythia-1b-reference \
    --sync-params-every 1 \
    --max-train 50000 \
    --generate-max-length 53 \
    --train-batch-size 128 \
    --rollout-batch-size 128 \
    --rollout-batch-size-per-device 32 \
    --pi-buffer-maxlen-per-device 32 \
    --train-batch-size-per-device 8 \
    --eval-steps 20 \
    --use-wb \
    --wb-run-name 1b_skywork_dpo_online

  12. Scale up with a remote Mosec service

    main

    To train larger models, you can host a preference oracle on a separate machine or a Kubernetes cluster.

    • Kubernetes: Follow the project's K8s guide to serve a remote oracle at a service URL like http://remote-rm.
    • Direct IP: For a standalone remote machine, use its IP address (e.g., http://10.0.0.1:8000) with the --remote-rm-url flag.

    Set --preference-oracle remote to enable this mode.

    python -m oat.experiment.main \
    --flash-attn \
    --gradient-checkpointing \
    --rnd-seed \
    --gpus 8 \
    --dap-algo DPO \
    --beta 0.1 \
    --preference-oracle remote \
    --remote-rm-url http://remote-rm \
    --pretrain trl-lib/pythia-6.9b-deduped-tldr-sft \
    --prompt-data lkevinzc/tldr-with-sft-reference \
    --input-key prompt \
    --output-key pythia-6.9b-reference \
    --sync-params-every 1 \
    --max-train 50000 \
    --generate-max-length 53 \
    --train-batch-size 128 \
    --rollout-batch-size 128 \
    --rollout-batch-size-per-device 32 \
    --pi-buffer-maxlen-per-device 32 \
    --train-batch-size-per-device 8 \
    --eval-steps 20 \
    --use-wb \
    --wb-run-name 6.9b_skywork_dpo_online