Long-Context QA with Arctic RL

Long-Context QA with Arctic RL

This recipe implements GRPO training for Qwen3-32B on long-context multi-hop question answering, using Arctic RL with the ZoRRo trainer. The training uses pure GRPO without a frozen reference model (no KL anchoring).

Training Data

The training data is LoongRL-Train-Data, a 16K-context corpus that merges three QA sources:

Infrastructure

Topology: 4 nodes × 8 H200 GPUs (32 GPUs total)

Configuration:

• colocate=True (training + sampling share each GPU bundle)
• Deepspeed ZeRO stage-3 with CPU optimizer offload
• vLLM rollout (TP=2)
• Log-probs recomputed through the training engine (no frozen reference model)

1. Ray and Multi-Node Hostfile

Ray and DeepSpeed require a hostfile (an MPI-style file) to discover participating nodes and their GPU counts.

If your CSP doesn't provide one, create a file with this format:

10.1.1.1 slots=8
10.1.1.2 slots=8
10.1.1.3 slots=8
10.1.1.4 slots=8

The first column is node IP addresses; the second column is the number of GPUs on each node.

Export the path so the install and launcher steps can access it:

export JOB_HOSTFILE=/path/to/hostfile

You can also modify the HOSTFILE setting at the top of the launcher scripts instead.

2. Install Packages

Conda environments are node-local, so each node must have its own environment. Use ds_ssh (DeepSpeed's multi-node helper) to install on all nodes simultaneously.

Set up environment variables:

export CONDA_ENV=long_context_qa
CONDA_BASE=$(conda info --base)
ENV=$CONDA_BASE/envs/$CONDA_ENV/bin        # env's python/uv/ds_ssh/ray live here

Bootstrap the launching node:

conda create -y -n $CONDA_ENV python=3.12
$ENV/python -m pip install -q uv
$ENV/uv pip install --python $ENV/python deepspeed       # provides $ENV/ds_ssh

Clone repositories to shared storage:

git clone https://github.com/Snowflake-AI-Research/Arctic-Platform
git clone -b arctic_rl_share_v0.7.1 --single-branch https://github.com/Snowflake-AI-Research/verl
cd Arctic-Platform/recipes/rl/verl/long_context_qa

Create environment and install dependencies on all nodes:

# Create env on remaining nodes (idempotent on launching node)
$ENV/ds_ssh -f $JOB_HOSTFILE "[ -x $ENV/python ] || $CONDA_BASE/bin/conda create -y -n $CONDA_ENV python=3.12"
$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/python -m pip install -q uv"

# Install torch (CUDA 12.9) first
$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/uv pip install --python $ENV/python torch==2.10.0 --index-url https://download.pytorch.org/whl/cu129 -U"

# Install pinned requirements
$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/uv pip install --python $ENV/python -r $PWD/requirements.txt --override $PWD/overrides.txt"
$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/uv pip install --python $ENV/python -U pip wheel packaging setuptools"

Note: If using a different CUDA version, update the torch index URL and the cuda-bindings pin in requirements.txt. The recipe pins vllm==0.18.0, which is patched by arctic-inference; overrides.txt forces compatible transitive dependencies.

Install Flash Attention:

The launcher uses flash_attention_2 by default:

$ENV/ds_ssh -f $JOB_HOSTFILE "$ENV/uv pip install --python $ENV/python flash-attn --no-build-isolation"

Alternatively, download a prebuilt FA2 wheel from flash-attention releases.

To use flash_attention_3 instead (faster on Hopper), install the matching flash_attn_3 wheel, then enable it in the launcher by commenting out flash_attention_v=flash_attention_2 and uncommenting the GPU-type auto-selection block.

Install verl (Snowflake fork) on all nodes:

cd ../../../../../verl
grep -v flash-attn requirements.txt > requirements-no-fa.txt
$ENV/ds_ssh -f $JOB_HOSTFILE "cd $PWD && $ENV/uv pip install --python $ENV/python -r requirements-no-fa.txt && $ENV/uv pip install --python $ENV/python -e ."
cd -

Single-node setup: Skip ds_ssh and run the $ENV/uv pip install commands directly on your local node.

3. Data Preparation

The download_data.py script pulls the three dataset pairs from HuggingFace, prepends a system prompt requesting thinking in <think> tags and answers in \boxed{}, and generates train/test parquet files.

From the recipe directory:

$ENV/python download_data.py --output_dir /data/snowflakesql/long-context

Default arguments:

• --output_dir: /data/snowflakesql/long-context
• --test_ratio: 0.05
• --seed: 42

Output structure:

/data/snowflakesql/long-context/
├── hotpotqa/{train,test}.parquet
├── musique/{train,test}.parquet
├── 2wikimqa/{train,test}.parquet
└── merged/
    ├── train.parquet      # ~14k rows: all three tasks concatenated
    └── test.parquet       # ~750 rows

The training recipe consumes merged/train.parquet and merged/test.parquet by default. To train on a single task, point to that task's {train,test}.parquet instead.

4. Train

Start the Ray cluster across all nodes (requires $JOB_HOSTFILE exported from step 1):

bash ./restart_multi_ray.sh

Configure the launcher script. Edit environment variables in run_qwen3_32b_longcontext_grpo_arl.sh:

• HF_HOME - HuggingFace hub cache location (optional)
• VLLM_CACHE_ROOT - vLLM cache directory

Launch training:

bash run_qwen3_32b_longcontext_grpo_arl.sh \
    data.train_files=/data/snowflakesql/long-context/merged/train.parquet \
    data.val_files=/data/snowflakesql/long-context/merged/test.parquet

Alternatively, edit DATA_DIR in the script (defaults to /data/snowflakesql/long-context) and launch without overrides:

bash run_qwen3_32b_longcontext_grpo_arl.sh

Reward Scoring

Answer rewards are scored by reward.py (included with the recipe and auto-wired via custom_reward_function). The scorer extracts the model's \boxed{} answer and matches it against the ground truth.

Key Training Knobs

Edit these settings inside run_qwen3_32b_longcontext_grpo_arl.sh before launching.