Agentic RL Configuration#
This section covers configuration parameters specific to agentic and reasoning RL training (math reasoning, coding agents, multi-agent systems). These extend the shared configuration described in Basic Configuration.
runner#
runner:
enable_dynamic_batch_size: False
max_tokens_per_mbs: 2048
Parameter |
Description |
|---|---|
|
Whether to use dynamic batch size when training with Megatron. |
|
Upper limit of tokens in a Megatron micro-batch when dynamic batching is enabled. |
algorithm#
algorithm:
n_minibatches: 4
training_batch_size_per_gpu: 1
rollout_batch_size_per_gpu: null
recompute_logprobs: True
shuffle_rollout: False
clip_ratio_low: null
clip_ratio_high: null
sampling_params:
max_new_tokens: ${subtract:${runner.seq_length}, ${data.max_prompt_length}}
min_new_tokens: 1
Parameter |
Description |
|---|---|
|
Number of gradient updates per batch. |
|
Micro-batch size on each actor GPU. |
|
Inference micro-batch per GPU; |
|
Recompute log-probs in the training engine instead of trusting the rollout engine’s values. |
|
Shuffle rollout samples before the optimization step. |
|
Asymmetric PPO clip bounds; |
|
Max generated tokens; computed from |
|
Minimum generated tokens. |
Shared loss/advantage keys (loss_type, adv_type, kl_beta,
ratio_clip_eps, group_size, sampling_params.temperature, …) are
documented in Basic Configuration.
rollout#
rollout:
rollout_backend: sglang # [sglang, vllm]
enforce_eager: False
distributed_executor_backend: mp # ray or mp
disable_log_stats: False
detokenize: False
padding: null # tokenizer.pad_token_id if null
eos: null # tokenizer.eos_token_id if null
tensor_parallel_size: 1
pipeline_parallel_size: 1
return_logprobs: ${not:${algorithm.recompute_logprobs}}
validate_weight: False
validate_save_dir: null
print_outputs: False
max_running_requests: 64
cuda_graph_max_bs: 128
sglang:
attention_backend: triton # [flashinfer, triton]
decode_log_interval: 500000
use_torch_compile: False
torch_compile_max_bs: 128
vllm:
attention_backend: FLASH_ATTN # [FLASH_ATTN, XFORMERS]
enable_chunked_prefill: True
enable_prefix_caching: True
enable_flash_infer_sampler: True
max_num_batched_tokens: null
Parameter |
Description |
|---|---|
|
Generation backend to use ( |
|
If True, disable CUDA-graph capture to shorten warm-up. |
|
Backend for launching rollout workers ( |
|
Suppress periodic backend stats logging. |
|
Detokenize outputs for debugging (RL usually uses token ids only). |
|
Pad token id override; |
|
EOS token id override; |
|
TP degree inside the generation backend. See 5D Parallelism Configuration. |
|
PP degree inside the generation backend. See 5D Parallelism Configuration. |
|
Whether the engine returns log-probs; defaults to the negation of
|
|
Send full weights once for cross-check/validation. |
|
Directory to store weights for comparison when validation is enabled. |
|
Print token ids/texts from the engine for debugging. |
|
Max concurrent decode requests. |
|
Max batch size eligible for CUDA graph. |
SGLang backend (``rollout.sglang``):
Parameter |
Description |
|---|---|
|
Attention kernel backend ( |
|
Interval (in steps) for SGLang to log decode stats. |
|
Enable |
|
Max batch size eligible for |
vLLM backend (``rollout.vllm``):
Parameter |
Description |
|---|---|
|
Attention backend ( |
|
Enable chunked prefill. |
|
Enable prefix caching. |
|
Use FlashInfer for sampling. |
|
Maximum tokens batched together; |
rollout.group_name, rollout.gpu_memory_utilization, and rollout.model.*
are shared keys documented in Basic Configuration.
data#
data:
type: math
dataset_name: boba
max_prompt_length: 1024
filter_prompt_by_length: True
rollout_batch_size: 512
val_rollout_batch_size: null
num_workers: 2
prompt_key: prompt
answer_key: solutions
apply_chat_template: False
shuffle: True
validation_shuffle: True
seed: 1234
train_data_paths: ["/dataset/boba/AReaL-boba-106k.jsonl"]
val_data_paths: ["/dataset/boba/AReaL-boba-106k.jsonl"]
Parameter |
Description |
|---|---|
|
Dataset/task family (e.g., |
|
Dataset identifier used to select preprocessing/formatting. |
|
Maximum tokens allowed for prompts. |
|
Drop prompts longer than |
|
Global rollout batch size across engines. |
|
Global validation rollout batch size; |
|
Data-loader workers per actor rank. |
|
JSONL key that stores the prompt text. |
|
JSONL key that stores the reference answer/solution. |
|
Wrap prompts with the tokenizer chat template before tokenizing. |
|
Shuffle training data each epoch. |
|
Shuffle validation data. |
|
RNG seed for loaders and sampling. |
|
List of training JSONL file paths. |
|
List of validation JSONL file paths. |
actor#
actor:
training_backend: megatron
mcore_gpt: True
spec_name: decoder_gpt
offload_optimizer: True
offload_weight: True
offload_grad: True
enable_dp_load_balance: False
calculate_flops: False
seed: 1234
model:
precision: fp16
add_bias_linear: False
tensor_model_parallel_size: 1
pipeline_model_parallel_size: 1
activation: swiglu
sequence_parallel: True
recompute_method: block
recompute_granularity: full
recompute_num_layers: 20
seq_length: ${runner.seq_length}
encoder_seq_length: ${runner.seq_length}
normalization: rmsnorm
position_embedding_type: rope
apply_rope_fusion: True
bias_dropout_fusion: False
persist_layer_norm: False
bias_activation_fusion: False
attention_softmax_in_fp32: True
batch_p2p_comm: False
variable_seq_lengths: True
gradient_accumulation_fusion: False
moe_token_dispatcher_type: alltoall
use_cpu_initialization: False
optim:
optimizer: adam
bf16: False
fp16: True
lr: 2e-05
adam_beta1: 0.9
adam_beta2: 0.95
adam_eps: 1.0e-05
min_lr: 2.0e-6
weight_decay: 0.05
use_distributed_optimizer: True
overlap_grad_reduce: False
overlap_param_gather: False
optimizer_enable_pin: false
overlap_param_gather_with_optimizer_step: False
clip_grad: 0.8
loss_scale: 65536
lr_sched:
lr_warmup_fraction: 0.01
lr_warmup_init: 0.0
lr_warmup_iters: 0
max_lr: 2.0e-5
min_lr: 0.0
lr_decay_style: constant
lr_decay_iters: 10
tokenizer:
tokenizer_model: /path/to/model/DeepSeek-R1-Distill-Qwen-1.5B/
use_fast: False
trust_remote_code: True
padding_side: 'right'
megatron:
ddp_bucket_size: null
distributed_backend: nccl # 'nccl' or 'gloo'
distributed_timeout_minutes: 30
ckpt_format: torch
use_dist_ckpt: False
tp_comm_bootstrap_backend: nccl
tp_comm_overlap_cfg: null
use_hf_ckpt: True
use_profiler: False
ckpt_convertor: # config for the checkpoint convertor
model: DeepSeek-R1-Distill-Qwen-1.5B
hf_model_path: ${rollout.model.model_path}
save_path: ${runner.output_dir}/${runner.experiment_name}/converted_ckpts/actor
use_gpu_num: 0
use_gpu_index: null
process_num: 16
tensor_model_parallel_size: ${actor.model.tensor_model_parallel_size}
pipeline_model_parallel_size: ${actor.model.pipeline_model_parallel_size}
fsdp_config:
strategy: "fsdp"
sharding_strategy: "no_shard"
cpu_offload: False
offload_pin_memory: False
reshard_after_forward: True
enable_gradient_accumulation: True
forward_prefetch: False
limit_all_gathers: False
backward_prefetch: null
use_orig_params: False
use_liger_kernel: False
mixed_precision:
param_dtype: ${actor.model.precision}
reduce_dtype: ${actor.model.precision}
buffer_dtype: ${actor.model.precision}
amp_autocast:
enabled: False
precision: "bf16"
grad_scaler:
enabled: False
Top-level
Parameter |
Description |
|---|---|
|
Training backend ( |
|
Use the Megatron-Core GPT stack. |
|
Model spec/preset name (e.g., decoder-only GPT). |
|
Offload optimizer state to CPU to reduce GPU memory. |
|
Offload model weights to CPU when possible. |
|
Offload gradients to CPU to reduce GPU memory. |
|
Enable data-parallel load balancing. |
|
Compute and log FLOPs for profiling. |
|
Global seed for reproducibility. |
Model sub-section
Parameter |
Description |
|---|---|
|
Numerical precision for training (e.g., |
|
Add bias terms to linear layers. |
|
TP degree for the actor. |
|
PP degree for the actor. |
|
Activation function (e.g., |
|
Enable sequence parallelism (requires TP). |
|
Activation-recompute strategy (e.g., |
|
Recompute scope ( |
|
Number of layers to checkpoint/recompute. |
|
Decoder context length for training. |
|
Encoder length (mirrors |
|
Norm-layer type (e.g., |
|
Positional-embedding type (e.g., |
|
Use fused RoPE kernels if available. |
|
Fuse bias + dropout kernels. |
|
Persist LayerNorm params in higher precision. |
|
Fuse bias + activation kernels. |
|
Compute attention softmax in FP32 for stability. |
|
Batch P2P communications across layers. |
|
Allow variable sequence lengths per micro-batch. |
|
Fused gradient accumulation. |
|
MoE token dispatcher (e.g., |
|
Initialize weights on CPU to reduce GPU spikes. |
Optimizer
Parameter |
Description |
|---|---|
|
Optimizer choice ( |
|
Mixed-precision flags. |
|
Base learning rate. |
|
Adam hyper-parameters. |
|
Minimum LR (for schedulers that decay below the base LR). |
|
L2 weight decay. |
|
Use the Megatron distributed optimizer. |
|
Overlap gradient reduction with the backward pass. |
|
Overlap parameter all-gather with the forward pass. |
|
Pin optimizer memory. |
|
Overlap param gather with the optimizer step. |
|
Global gradient-clipping norm. |
|
Static FP16 loss scale (use |
LR schedule
Parameter |
Description |
|---|---|
|
Warm-up as a fraction of total iterations. |
|
Initial LR value during warm-up. |
|
Warm-up iterations (overrides the fraction when > 0). |
|
LR bounds for the scheduler. |
|
Decay policy (e.g., |
|
Total decay iterations. |
Tokenizer
Parameter |
Description |
|---|---|
|
Path/name of the tokenizer. |
|
Use the HuggingFace fast tokenizer. |
|
Allow custom tokenizer code. |
|
|
Megatron integration
Parameter |
Description |
|---|---|
|
DDP gradient bucket size. |
|
Distributed backend ( |
|
Backend communication timeout. |
|
Checkpoint format (e.g., |
|
Use distributed (sharded) checkpointing. |
|
Backend used for TP bootstrap (e.g., |
|
YAML path for TP comm/compute overlap. |
|
Convert/load from a HuggingFace checkpoint for training. |
|
Enable the Torch profiler during training (affects performance). |
Megatron checkpoint converter (``actor.megatron.ckpt_convertor``)
Parameter |
Description |
|---|---|
|
Model name for the converter metadata. |
|
Source HF model path. |
|
Target directory for the converted Megatron checkpoint. |
|
Number of GPUs to use for conversion. |
|
Specific GPU index to use. |
|
CPU processes for conversion work. |
|
TP degree for the converted checkpoint. |
|
PP degree for the converted checkpoint. |
FSDP integration (``actor.fsdp_config``)
Used when actor.training_backend is fsdp.
Parameter |
Description |
|---|---|
|
FSDP strategy: |
|
Sharding strategy: |
|
FSDP2: keep parameters on CPU, moving them to GPU only when needed. |
|
FSDP2: use pinned CPU memory (only when |
|
FSDP2: re-shard parameters after the forward pass to save GPU memory. |
|
Communicate/update only after the last micro-batch. Speeds up training at the cost of GPU memory. |
|
FSDP: prefetch the next all-gather during the forward pass (more memory, better overlap). |
|
FSDP: limit concurrent all-gathers (recommended when CPU/memory bound). |
|
FSDP: prefetch strategy in the backward pass ( |
|
FSDP: expose original (unflattened) parameters; better compatibility, more communication overhead. |
|
Use Liger kernels (currently Qwen2.5 / Qwen2.5-VL) to cut memory and speed up training. |
|
Parameter dtype. |
|
Reduction dtype. |
|
Buffer dtype. |
|
Enable automatic mixed-precision (AMP) training. |
|
Numerical precision used by AMP. |
|
Enable the gradient scaler. |
reward#
reward:
reward_type: math
reward_scale: 5.0
Parameter |
Description |
|---|---|
|
Which reward type to use for training (e.g., |
|
A correct answer receives |