chore: import upstream snapshot with attribution

examples/ray/README.md

@@ -0,0 +1,133 @@
+# Ray-based Megatron RLHF examples (GKD & GRPO)
+
+GRPO/GKD on top of Megatron, orchestrated by Ray. The student/actor is
+trained with Megatron, generates completions with vLLM, and — for GKD — is distilled with a teacher model.
+
+## How to run
+
+```bash
+# via the helper scripts
+CUDA_VISIBLE_DEVICES=0,1,2,3 bash examples/ray/gkd/run.sh
+
+# or directly
+megatron rlhf --use_ray true --config examples/ray/gkd/rollout_colocate_teacher_colocate.yaml
+```
+
+The YAML is split into a top-level section (shared args) and per-role groups
+(`train`, `rollout`, and optionally `teacher`). Each group's `gpus:` field sets how
+many GPUs that role uses; `CUDA_VISIBLE_DEVICES` must expose at least the total
+number of GPUs the chosen placement needs (see below).
+
+The `gkd/` folder ships three ready-to-run configs. The file name encodes the two
+independent choices — **rollout placement** and **teacher mode**:
+
+| file | rollout | teacher |
+|------|---------|---------|
+| `rollout_colocate_teacher_colocate.yaml` | colocate (shares train GPUs) | colocated (shares train GPUs) |
+| `rollout_separate_teacher_colocate.yaml` | separate (own GPUs) | colocated (shares train GPUs) |
+| `rollout_colocate_teacher_standalone.yaml` | colocate (shares train GPUs) | standalone vLLM replicas (own GPUs) |
+
+---
+
+## 1. GPU placement: colocate vs separate
+
+This is controlled by `colocate_groups` plus each role's `gpus`.
+
+| Placement | `colocate_groups` | GPUs needed | When to use |
+|-----------|-------------------|-------------|-------------|
+| **colocate** | `[[train, rollout]]` | `train.gpus` — all roles in the group **must** set the same `gpus` (one shared set) | default; fewer GPUs, train and rollout time-share the same devices |
+| **separate** | *omit* | `train.gpus + rollout.gpus` (disjoint sets) | more GPUs, rollout overlaps with training |
+
+- **colocate** — train and rollout live on the *same* devices and take turns.
+  Set `offload_model`/`offload_optimizer`
+  (+ `offload_teacher_model` for GKD) and `sleep_level: 1` so the idle role releases
+  GPU memory to the active one.
+  Example: `train.gpus=4`, `rollout.gpus=4`, `colocate_groups: [[train, rollout]]`
+  → 4 GPUs total, with TP2 giving **DP2**.
+
+- **separate** — train and rollout occupy *disjoint* GPU sets; weights are pushed to
+  the rollout engine every step.
+
+---
+
+## 2. Teacher modes (GKD only)
+
+Pick exactly one. `gkd_logits_topk: K` selects top-k distillation;
+omit it for full-vocab distillation.
+
+| Mode | How to configure | top-k | full-vocab | Status |
+|------|------------------|:-----:|:----------:|--------|
+| **Colocated `teacher_model`** | set top-level `teacher_model:` (+ `offload_teacher_model: true`) | ✅ | ✅ | supported |
+| **Standalone teacher replicas** | add a `teacher:` group with `gpus`, `model`, and `vllm_engine_kwargs.max_logprobs` | ✅ | ❌ | supported |
+
+### 2a. Colocated teacher (`rollout_colocate_teacher_colocate.yaml`, `rollout_separate_teacher_colocate.yaml`)
+The teacher shares the **train** GPUs and is offloaded to CPU between teacher
+forwards. It is the only mode that supports full-vocab distillation, and it works
+with both colocate and separate rollout placements.
+
+```yaml
+teacher_model: Qwen/Qwen3.5-4B
+offload_teacher_model: true
+gkd_logits_topk: 64      # omit for full-vocab
+```
+
+### 2b. Standalone teacher replicas (`rollout_colocate_teacher_standalone.yaml`)
+The teacher runs as its own set of Ray-managed vLLM replicas on **separate** GPUs and
+returns prompt top-k logprobs; the driver fetches them per step.
+
+```yaml
+gkd_logits_topk: 64                 # REQUIRED — replicas are top-k only
+# do NOT set top-level teacher_model here (that would also load a colocated teacher)
+teacher:
+  gpus: 4
+  model: Qwen/Qwen3.5-4B            # the teacher checkpoint these replicas serve
+  vllm_engine_kwargs: {"max_logprobs": 64}   # MUST be >= gkd_logits_topk
+```
+- `max_logprobs` must be `>= gkd_logits_topk`, or vLLM rejects the `prompt_logprobs`
+  request.
+- GPUs needed = colocated train+rollout set **+** `teacher.gpus`.
+
+---
+
+## 3. top-k vs full-vocab distillation
+
+- **top-k** (`gkd_logits_topk: K`): the teacher exposes only the top-K logprobs per
+  position. Much lower memory, works for every teacher mode.
+- **full-vocab** (omit `gkd_logits_topk`): distill the full vocabulary distribution.
+  Colocated teacher only, and **memory-heavy** (caches per-rank vocab-sharded
+  teacher logits). If you OOM: switch to top-k, lower `micro_batch_size`
+
+---
+
+## 4. OPSD (On-Policy / privileged Distillation)
+
+OPSD lets the teacher see a *different* (privileged) prompt than the student while
+scoring the **same** on-policy response — e.g. the teacher sees the problem + a
+reference solution. A dataset preprocessor (loaded via `external_plugins`) emits a
+per-row `teacher_prompt`; the loss aligns the shared response tokens by mask.
+
+```yaml
+external_plugins: examples/train/rlhf/opsd/opsd_plugin.py   # registers teacher_prompt
+teacher_model: Qwen/Qwen3.5-4B
+gkd_logits_topk: 64
+```
+- Supported in Ray with **top-k** (`gkd_logits_topk`) for both a **colocated teacher**
+  and **standalone teacher replicas** (`teacher.gpus > 0`).
+- No extra flag is needed: OPSD activates automatically when rows carry a non-empty
+  `teacher_prompt`; otherwise training falls back to plain GKD.
+
+---
+
+## 5. Things to know (common knobs & pitfalls)
+
+- **Sequence length**: the encoder budget is `max_length + max_completion_length`
+  (prompt is capped at `max_length`, the on-policy completion adds up to
+  `max_completion_length`). Size `vllm_max_model_len` accordingly.
+- **`padding_free: true`** packs a micro-batch into one sequence; pair with
+  `sequence_parallel: true` when `tensor_model_parallel_size > 1`.
+- **Parallelism / DP**: data parallel size = `gpus / (TP * PP * CP)`. e.g. 4 GPUs
+  with `tensor_model_parallel_size: 2` → DP2.
+- **Memory release (colocate)**: `offload_model`, `offload_optimizer`,
+  `offload_teacher_model`, and `sleep_level: 1` are what let colocated roles fit.
+- **GRPO specifics**: rewards via `reward_funcs` + `external_plugins`; sampling via
+  `num_generations` / `steps_per_generation`; no `teacher_*` settings.

examples/ray/gkd/multi_turn_colocate.yaml

@@ -0,0 +1,68 @@
+# Ray Megatron GKD multi-turn — colocate mode, colocated teacher.
+# Tests multi-turn GKD on Ray-Megatron backend with MathTipsScheduler.
+rlhf_type: gkd
+
+model: Qwen/Qwen3.5-0.8B
+teacher_model: Qwen/Qwen3.5-2B
+gkd_logits_topk: 64
+enable_thinking: false
+
+dataset:
+  - 'AI-MO/NuminaMath-TIR#2000'
+dataset_num_proc: 4
+split_dataset_ratio: 0
+
+micro_batch_size: 1
+global_batch_size: 16
+num_train_epochs: 1
+logging_steps: 5
+seed: 42
+max_length: 2048
+max_completion_length: 1024
+truncation_strategy: delete
+padding_free: true
+sequence_parallel: true
+attention_backend: flash
+
+lr: 1e-5
+min_lr: 1e-7
+lr_warmup_fraction: 0.05
+
+temperature: 1.0
+lmbda: 1
+beta: 0.5
+sft_alpha: 0.0
+
+# Multi-turn configuration
+multi_turn_scheduler: math_tip_trick
+max_turns: 2
+
+finetune: true
+no_save_optim: true
+no_save_rng: true
+recompute_granularity: selective
+
+use_vllm: true
+
+colocate_groups: [[train, rollout]]
+offload_model: true
+offload_optimizer: true
+offload_teacher_model: true
+sleep_level: 1
+
+train:
+  gpus: 4
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 2
+  pipeline_model_parallel_size: 1
+  expert_model_parallel_size: 1
+  context_parallel_size: 1
+  output_dir: megatron_output/gkd_multi_turn_ray
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 2
+  vllm_gpu_memory_utilization: 0.4
+  vllm_max_model_len: 4096

examples/ray/gkd/rollout_colocate_teacher_colocate.yaml

@@ -0,0 +1,62 @@
+# Ray Megatron GKD — colocate mode (train + rollout share GPUs), colocated teacher.
+rlhf_type: gkd
+
+model: Qwen/Qwen3.5-2B
+teacher_model: Qwen/Qwen3.5-4B
+gkd_logits_topk: 64  # omit for full-vocab distillation
+
+dataset:
+  - 'AI-ModelScope/alpaca-gpt4-data-zh#2000'
+  - 'AI-ModelScope/alpaca-gpt4-data-en#2000'
+dataset_num_proc: 4
+split_dataset_ratio: 0
+
+micro_batch_size: 1
+global_batch_size: 16
+num_train_epochs: 1
+logging_steps: 1
+seed: 42
+max_length: 8192
+max_completion_length: 8192
+padding_free: true
+sequence_parallel: true
+attention_backend: flash
+
+lr: 1e-6
+min_lr: 1e-6
+lr_warmup_fraction: 0.0
+
+temperature: 1.0
+lmbda: 1
+beta: 0.5
+sft_alpha: 0.0
+
+finetune: true
+no_save_optim: true
+no_save_rng: true
+recompute_granularity: selective
+
+use_vllm: true
+
+colocate_groups: [[train, rollout]]
+offload_model: true
+offload_optimizer: true
+offload_teacher_model: true
+sleep_level: 1
+
+train:
+  gpus: 4
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 2  # TP2 -> DP2 on 4 GPUs
+  pipeline_model_parallel_size: 1
+  expert_model_parallel_size: 1
+  context_parallel_size: 1
+  output_dir: megatron_output/gkd_rollout_colocate_teacher_colocate
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 2
+  vllm_gpu_memory_utilization: 0.4
+  vllm_max_model_len: 16384

examples/ray/gkd/rollout_colocate_teacher_standalone.yaml

@@ -0,0 +1,68 @@
+# Ray Megatron GKD — standalone vLLM teacher replicas (top-k only).
+# gkd_logits_topk is REQUIRED; the teacher group below serves the teacher model and its max_logprobs must be >= gkd_logits_topk.
+# Do NOT set a top-level `teacher_model` here.
+rlhf_type: gkd
+
+model: Qwen/Qwen3.5-2B
+gkd_logits_topk: 64
+
+dataset: AI-ModelScope/alpaca-gpt4-data-en#2000
+dataset_num_proc: 4
+split_dataset_ratio: 0
+
+micro_batch_size: 2
+global_batch_size: 16
+num_train_epochs: 1
+logging_steps: 1
+seed: 42
+max_length: 8192
+max_completion_length: 8192
+padding_free: true
+sequence_parallel: true
+attention_backend: flash
+
+lr: 1e-6
+min_lr: 1e-6
+lr_warmup_fraction: 0.0
+
+temperature: 1.0
+lmbda: 1
+beta: 0.5
+sft_alpha: 0.0
+
+finetune: true
+no_save_optim: true
+no_save_rng: true
+recompute_granularity: selective
+
+use_vllm: true
+
+colocate_groups: [[train, rollout]]
+offload_model: true
+offload_optimizer: true
+sleep_level: 1
+
+train:
+  gpus: 4
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 2
+  pipeline_model_parallel_size: 1
+  expert_model_parallel_size: 1
+  context_parallel_size: 1
+  output_dir: megatron_output/gkd_rollout_colocate_teacher_standalone
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 1
+  vllm_gpu_memory_utilization: 0.4
+  vllm_max_model_len: 16384
+
+teacher:
+  gpus: 4
+  model: Qwen/Qwen3.5-4B          # teacher served by these replicas
+  vllm_tensor_parallel_size: 1
+  vllm_gpu_memory_utilization: 0.9
+  vllm_max_model_len: 16384
+  vllm_engine_kwargs: {"max_logprobs": 64}   # must be >= gkd_logits_topk

examples/ray/gkd/rollout_separate_teacher_colocate.yaml

@@ -0,0 +1,60 @@
+# Ray Megatron GKD — separate rollout (disjoint GPUs) + colocated teacher.
+# No `colocate_groups` => train and rollout use disjoint GPU sets
+rlhf_type: gkd
+
+model: Qwen/Qwen3.5-2B
+teacher_model: Qwen/Qwen3.5-4B
+gkd_logits_topk: 64
+
+dataset: AI-ModelScope/alpaca-gpt4-data-en#2000
+dataset_num_proc: 4
+split_dataset_ratio: 0
+
+micro_batch_size: 2
+global_batch_size: 16
+num_train_epochs: 1
+logging_steps: 1
+seed: 42
+max_length: 8192
+max_completion_length: 8192
+padding_free: true
+sequence_parallel: true
+attention_backend: flash
+
+lr: 1e-6
+min_lr: 1e-6
+lr_warmup_fraction: 0.0
+
+temperature: 1.0
+lmbda: 1
+beta: 0.5
+sft_alpha: 0.0
+
+finetune: true
+no_save_optim: true
+no_save_rng: true
+recompute_granularity: selective
+
+use_vllm: true
+
+# No colocate_groups -> train and rollout occupy separate GPU sets
+offload_model: true
+offload_optimizer: true
+offload_teacher_model: true
+
+train:
+  gpus: 4
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 2
+  pipeline_model_parallel_size: 1
+  expert_model_parallel_size: 1
+  context_parallel_size: 1
+  output_dir: megatron_output/gkd_rollout_separate_teacher_colocate
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 1
+  vllm_gpu_memory_utilization: 0.8
+  vllm_max_model_len: 16384

examples/ray/gkd/run.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# Ray Megatron GKD — default example (rollout colocate + colocated teacher).
+# Swap --config for another yaml in this folder for other placements/teacher modes.
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+megatron rlhf --use_ray true --config "$SCRIPT_DIR/rollout_colocate_teacher_colocate.yaml"

examples/ray/grpo/multi_turn/frozen_lake_colocate.sh

@@ -0,0 +1 @@
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 megatron rlhf --use_ray --config frozen_lake_colocate.yaml

examples/ray/grpo/multi_turn/frozen_lake_colocate.yaml

@@ -0,0 +1,92 @@
+# Ray + Megatron GRPO on FrozenLake (multi-turn, colocate mode).
+# Parameters aligned with examples/megatron/grpo/multi_turn/frozen_lake.sh
+rlhf_type: grpo
+
+model: Qwen/Qwen3.5-2B
+dataset: 'examples/megatron/grpo/multi_turn/frozen_lake.jsonl#1024'
+external_plugins: examples/megatron/grpo/multi_turn/frozen_lake_plugin.py
+load_from_cache_file: false
+split_dataset_ratio: 0
+dataset_num_proc: 4
+dataloader_num_workers: 4
+
+# Multi-turn frozen_lake env
+multi_turn_scheduler: gym_scheduler
+gym_env: frozen_lake
+use_gym_env: true
+max_turns: 10
+
+# Batch config
+micro_batch_size: 1
+global_batch_size: 64
+num_generations: 8
+steps_per_generation: 4
+train_iters: 120
+logging_steps: 1
+seed: 42
+
+# Length params
+max_length: 6120
+max_completion_length: 512
+vllm_max_model_len: 6632
+
+enable_thinking: false
+
+# Training config
+padding_free: true
+cross_entropy_loss_fusion: true
+gradient_accumulation_fusion: false
+attention_backend: flash
+recompute_granularity: selective
+
+# Optimizer
+lr: 5e-5
+bf16: true
+
+# Generation config
+temperature: 1.0
+top_p: 1.0
+top_k: 80
+
+# GRPO params
+beta: 0.001
+epsilon: 0.2
+epsilon_high: 0.2
+loss_type: grpo
+advantage_estimator: grpo
+importance_sampling_level: token
+dynamic_sample: false
+overlong_filter: true
+log_completions: true
+log_rollout_offpolicy_metrics: true
+
+# vLLM + colocate
+use_vllm: true
+colocate_groups: [[train, rollout]]
+offload_model: true
+offload_optimizer: true
+sleep_level: 1
+
+# Checkpointing
+eval_steps: 1000
+save_steps: 1000
+no_save_optim: true
+no_save_rng: true
+
+# Reporting
+report_to: swanlab
+
+train:
+  gpus: 8
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 1
+  pipeline_model_parallel_size: 1
+  context_parallel_size: 1
+  output_dir: megatron_output/ray_grpo_frozen_lake_colocate
+
+rollout:
+  gpus: 8
+  vllm_tensor_parallel_size: 1
+  vllm_gpu_memory_utilization: 0.5

examples/ray/grpo/opd_rl_colocate.yaml

@@ -0,0 +1,53 @@
+rlhf_type: grpo
+
+model: Qwen/Qwen3.5-2B
+teacher_model: Qwen/Qwen3.5-9B
+
+dataset: modelscope/gsm8k
+dataset_num_proc: 4
+split_dataset_ratio: 0
+
+micro_batch_size: 2
+global_batch_size: 16
+num_generations: 1
+steps_per_generation: 4
+num_train_epochs: 1
+logging_steps: 1
+seed: 42
+max_length: 2048
+max_completion_length: 2048
+padding_free: false
+cross_entropy_loss_fusion: true
+gradient_accumulation_fusion: false
+lr: 3e-5
+lr_warmup_fraction: 0.0
+attention_backend: flash
+temperature: 1.0
+beta: 0
+teacher_kl_coef: 1.0
+
+use_vllm: true
+
+colocate_groups: [[train, rollout]]
+offload_model: true
+offload_optimizer: true
+offload_teacher_model: true
+sleep_level: 1
+
+save_steps: 100
+no_save_optim: true
+no_save_rng: true
+
+train:
+  gpus: 4
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 1
+  output_dir: megatron_output/ray_opd_rl_colocate
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 1
+  vllm_gpu_memory_utilization: 0.4
+  vllm_max_model_len: 4096

examples/ray/grpo/ray_grpo_colocate.yaml

@@ -0,0 +1,56 @@
+rlhf_type: grpo
+
+model: Qwen/Qwen2.5-VL-3B-Instruct
+dataset: AI-ModelScope/clevr_cogen_a_train#2000
+external_plugins: examples/train/grpo/plugin/plugin.py
+reward_funcs: [external_r1v_acc, format]
+overlong_filter: false
+dataset_num_proc: 4
+split_dataset_ratio: 0
+
+micro_batch_size: 2
+global_batch_size: 16
+num_generations: 8
+steps_per_generation: 4
+num_train_epochs: 1
+logging_steps: 1
+seed: 42
+max_length: 4096
+max_completion_length: 4096
+padding_free: false
+cross_entropy_loss_fusion: true
+gradient_accumulation_fusion: false
+lr: 5e-5
+lr_warmup_fraction: 0.0
+min_lr: 5e-5
+attention_backend: flash
+temperature: 1.0
+beta: 0
+epsilon: 0.2
+epsilon_high: 0.28
+loss_type: grpo
+advantage_estimator: grpo
+importance_sampling_level: token
+log_rollout_offpolicy_metrics: true
+system: examples/train/grpo/prompt.txt
+
+use_vllm: true
+
+colocate_groups: [[train, rollout]]
+offload_model: true
+offload_optimizer: true
+sleep_level: 1
+
+train:
+  gpus: 4
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 1
+  output_dir: megatron_output/ray_grpo_colocate
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 1
+  vllm_gpu_memory_utilization: 0.4
+  vllm_max_model_len: 8192

examples/ray/grpo/ray_grpo_separate.yaml

@@ -0,0 +1,53 @@
+rlhf_type: grpo
+
+model: Qwen/Qwen2.5-VL-3B-Instruct
+dataset: AI-ModelScope/clevr_cogen_a_train#2000
+external_plugins: examples/train/grpo/plugin/plugin.py
+reward_funcs: [external_r1v_acc, format]
+overlong_filter: false
+dataset_num_proc: 4
+split_dataset_ratio: 0
+load_from_cache_file: false
+
+micro_batch_size: 2
+global_batch_size: 16
+num_generations: 8
+steps_per_generation: 4
+num_train_epochs: 1
+train_iters: 100
+logging_steps: 1
+seed: 42
+max_length: 4096
+max_completion_length: 4096
+padding_free: false
+cross_entropy_loss_fusion: true
+gradient_accumulation_fusion: false
+lr: 5e-5
+lr_warmup_fraction: 0.0
+min_lr: 5e-5
+attention_backend: flash
+temperature: 1.0
+beta: 0
+epsilon: 0.2
+epsilon_high: 0.28
+loss_type: grpo
+advantage_estimator: grpo
+importance_sampling_level: token
+log_rollout_offpolicy_metrics: true
+system: examples/train/grpo/prompt.txt
+
+use_vllm: true
+
+train:
+  gpus: 4
+  tuner_type: lora
+  lora_rank: 8
+  lora_alpha: 32
+  tensor_model_parallel_size: 1
+  output_dir: megatron_output/ray_grpo_separate
+
+rollout:
+  gpus: 4
+  vllm_tensor_parallel_size: 1
+  vllm_gpu_memory_utilization: 0.8
+  vllm_max_model_len: 8192

examples/ray/grpo/run_colocate.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+# Ray-based Megatron GRPO — colocate mode
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+megatron rlhf --use_ray true --config "$SCRIPT_DIR/ray_grpo_colocate.yaml"

examples/ray/grpo/run_opd.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+# Ray Megatron OPD-RL (On-Policy Distillation as RL) — colocate mode + colocated teacher.
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+megatron rlhf --use_ray true --config "$SCRIPT_DIR/opd_rl_colocate.yaml"

examples/ray/grpo/run_separate.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+# Ray-based Megatron GRPO — separate mode
+export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+megatron rlhf --use_ray true --config "$SCRIPT_DIR/ray_grpo_separate.yaml"