Training Engine Configurations

This page documents the configuration structures for AReaL's training engines, including TrainEngineConfig, OptimizerConfig, and engine-specific settings for FSDP, Megatron, and Archon backends. These configurations control model training behavior, optimization parameters, parallelism strategies, and backend-specific features.

For information about inference engine configurations, see Inference Engine Configurations For parallelism strategy configuration, see allocation_mode Syntax For algorithm-specific configurations like PPO parameters, see Algorithm-Specific Configurations

---

Configuration Architecture

Training engine configurations in AReaL follow a hierarchical structure where TrainEngineConfig serves as the core configuration containing engine-specific sub-configurations.

Configuration Hierarchy

Sources: areal/api/cli_args.py889-1005 areal/api/cli_args.py306-375

---

TrainEngineConfig

TrainEngineConfig is the core configuration class for training engines, containing common parameters that apply across all backends as well as engine-specific sub-configurations.

Core Training Parameters

Parameter	Type	Default	Description
experiment_name	string	Required	Name of the experiment areal/api/cli_args.py892
trial_name	string	Required	Name of the trial within the experiment areal/api/cli_args.py895
path	string	""	Path to HuggingFace checkpoint or model identifier areal/api/cli_args.py898
attn_impl	string	"flash_attention_2"	Attention implementation. Choices: "flash_attention_2", "sdpa", "eager" areal/api/cli_args.py901
init_from_scratch	boolean	False	Initialize model weights randomly instead of loading pretrained weights areal/api/cli_args.py910
is_critic	boolean	False	Whether this engine is for a critic/reward model areal/api/cli_args.py913
temperature	float	1.0	Temperature for generation (if applicable) areal/api/cli_args.py916
mb_spec	MicroBatchSpec	default	Micro-batch specification for memory management areal/api/cli_args.py919
pad_to_maximum	boolean	False	Pad each micro-batch to maximum length (reduces fragmentation, slower) areal/api/cli_args.py922

Training Backend Settings

Parameter	Type	Default	Description
disable_dropout	boolean	False	Disable dropout layers during training areal/api/cli_args.py925
gradient_checkpointing	boolean	False	Enable gradient checkpointing to reduce memory usage areal/api/cli_args.py928
dtype	string	"bfloat16"	Parameter data type areal/api/cli_args.py931
grad_reduce_dtype	string	"float32"	Gradient reduction data type for distributed training areal/api/cli_args.py934
optimizer	OptimizerConfig | None	None	Optimizer configuration. None means no training (inference only) areal/api/cli_args.py937
weight_update_mode	string	"xccl"	Weight update backend. Choices: "disk", "xccl" areal/api/cli_args.py942

Engine-Specific Configurations

Parameter	Type	Default	Description
fsdp	FSDPEngineConfig	default	FSDP engine-specific settings areal/api/cli_args.py947
archon	ArchonEngineConfig	default	Archon engine-specific settings areal/api/cli_args.py950
megatron	MegatronEngineConfig	default	Megatron engine-specific settings areal/api/cli_args.py953

LoRA Configuration

Parameter	Type	Default	Description
use_lora	boolean	False	Enable LoRA (Low-Rank Adaptation) training. Supported with FSDP and Megatron areal/api/cli_args.py956
lora_rank	integer	32	LoRA rank parameter areal/api/cli_args.py959
lora_alpha	integer	16	LoRA alpha parameter areal/api/cli_args.py962
target_modules	list of string	[]	Target modules for LoRA adaptation areal/api/cli_args.py965
peft_type	string	"lora"	PEFT method type. Only LoRA is currently supported areal/api/cli_args.py968

Tree Training Configuration

Parameter	Type	Default	Description
enable_tree_training	boolean	False	Enable tree training for prefix-sharing efficiency areal/api/cli_args.py971

Sources: areal/api/cli_args.py889-1005

Configuration Flow Diagram

Sources: areal/api/cli_args.py889-1005 areal/engine/fsdp_engine.py218-222 areal/engine/megatron_engine.py168-173 areal/experimental/engine/archon_engine.py150-155

---

OptimizerConfig

OptimizerConfig specifies optimizer type, learning rate, scheduling, and related hyperparameters. It is referenced by TrainEngineConfig.optimizer.

Optimizer Type and Learning Rate

Parameter	Type	Default	Description
type	string	"adam"	Optimizer type. Choices: "adam", "sgd", "adam_bf16". For FSDP, adam_bf16 enables memory-efficient BF16 optimizer states via AnyPrecisionAdamW areal/api/cli_args.py309-315 areal/engine/fsdp_utils/optimizer.py85
lr	float	0.001	Learning rate areal/api/cli_args.py318
weight_decay	float	0.01	Weight decay coefficient areal/api/cli_args.py321

Adam-Specific Parameters

Parameter	Type	Default	Description
beta1	float	0.9	Adam beta1 parameter. Only effective for adam/adam_bf16 areal/api/cli_args.py324
beta2	float	0.999	Adam beta2 parameter. Only effective for adam/adam_bf16 areal/api/cli_args.py327
eps	float	1e-8	Adam epsilon parameter. Only effective for adam/adam_bf16 areal/api/cli_args.py330

Learning Rate Scheduling

Parameter	Type	Default	Description
lr_scheduler_type	string	"constant"	Learning rate scheduler type. Choices: "linear", "cosine", "constant" areal/api/cli_args.py333-337
warmup_steps_proportion	float	0.001	Proportion of training steps for warmup areal/api/cli_args.py340
min_lr_ratio	float	0.0	Minimum learning rate ratio after annealing areal/api/cli_args.py343

Optimizer State Management

Parameter	Type	Default	Description
offload	boolean	False	Enable optimizer state offloading to CPU areal/api/cli_args.py346

Mixed Precision Training (Loss Scaling)

Parameter	Type	Default	Description
initial_loss_scale	float	4294967296 (2^32)	Initial loss scaling factor areal/api/cli_args.py349
min_loss_scale	float	1.0	Minimum loss scaling factor areal/api/cli_args.py352
loss_scale_window	float	5	Window size for loss scaling adjustment areal/api/cli_args.py355
hysteresis	integer	2	Hysteresis (scaling factor) for loss scaling areal/api/cli_args.py358

Gradient Clipping

Parameter	Type	Default	Description
gradient_clipping	float	1.0	Gradient clipping threshold areal/api/cli_args.py361

Sources: areal/api/cli_args.py306-375

---

FSDP Engine Configuration

FSDP (Fully Sharded Data Parallel) is PyTorch's native training backend supporting N-D parallelism. FSDPEngineConfig controls FSDP-specific behaviors.

FSDPEngineConfig

Parameter	Type	Default	Description
wrap_policy	FSDPWrapPolicy | None	None	FSDP wrap policy specifying model layers to wrap. None defaults to wrapping transformer decoder layers areal/api/cli_args.py391
offload_params	boolean	False	Whether to offload FSDP parameters to CPU areal/api/cli_args.py396
memory_efficient_load	boolean	False	Enable memory-efficient model loading areal/api/cli_args.py399
shard_vision_across_sp	boolean	False	Shard vision encoder across SP ranks by image areal/api/cli_args.py408

FSDPWrapPolicy

Parameter	Type	Default	Description
transformer_layer_cls_to_wrap	list of string | None	None	List of transformer layer names for FSDP to wrap areal/api/cli_args.py381

Sources: areal/api/cli_args.py388-417 areal/api/cli_args.py378-385

FSDP Model Parallelization

Sources: areal/engine/fsdp_engine.py218-222 areal/engine/fsdp_utils/parallel.py86

---

Megatron Engine Configuration

Megatron-LM is NVIDIA's training framework supporting pipeline parallelism and expert parallelism. MegatronEngineConfig controls Megatron-Core specific features.

DistributedDataParallelConfig

Configuration for Megatron's DistributedDataParallel wrapper.

Parameter	Type	Default	Description
grad_reduce_in_fp32	boolean	True	Reduce gradients in FP32 precision areal/api/cli_args.py573
overlap_grad_reduce	boolean	False	Overlap gradient reduction with computation areal/api/cli_args.py574
overlap_param_gather	boolean	False	Overlap parameter gather with computation areal/api/cli_args.py575
align_param_gather	boolean	False	Align parameter gather operations areal/api/cli_args.py576
use_distributed_optimizer	boolean	True	Use Megatron's distributed optimizer areal/api/cli_args.py577
bucket_size	integer | None	None	Bucket size for gradient reduction areal/api/cli_args.py579

MegatronEngineConfig

Parameter	Type	Default	Description
wrap_with_ddp	boolean	True	Wrap model with DistributedDataParallel areal/api/cli_args.py695
ddp	DistributedDataParallelConfig	default	DDP configuration areal/api/cli_args.py704
virtual_pipeline_parallel_size	integer	1	Virtual pipeline parallel size for interleaved schedule areal/api/cli_args.py707
bridge_type	string	"mbridge"	Bridge type for weight loading. areal/api/cli_args.py692

Gradient Checkpointing Options

Only effective when TrainEngineConfig.gradient_checkpointing=True.

Parameter	Type	Default	Description
recompute_granularity	string | None	"full"	Recomputation granularity areal/api/cli_args.py741
recompute_method	string | None	"uniform"	Recomputation method areal/api/cli_args.py746
recompute_num_layers	integer | None	1	Number of layers to recompute areal/api/cli_args.py751

Sources: areal/api/cli_args.py692-772

---

Archon Engine Configuration

Archon is AReaL's experimental torch-native training backend. ArchonEngineConfig controls Archon-specific behaviors.

ArchonEngineConfig

Parameter	Type	Default	Description
attn_type	string	"varlen"	Attention backend type. Choices: "varlen", "sdpa", "tree" areal/api/cli_args.py423-427
enable_compile	boolean	True	Enable torch.compile for TransformerBlocks areal/api/cli_args.py436
pp_schedule	string	"Interleaved1F1B"	Pipeline parallel schedule areal/api/cli_args.py467-472
pp_layers_per_stage	integer | None	None	Number of transformer layers per virtual pipeline stage areal/api/cli_args.py475

Sources: areal/api/cli_args.py420-566

Archon Pipeline Parallelism Flow

Sources: areal/experimental/engine/archon_engine.py183-187 areal/experimental/engine/archon_runner.py56

---

FP8 Training Configuration

FP8EngineConfig encapsulates FP8 (8-bit floating point) training parameters. Currently supported by the Megatron engine.

Parameter	Type	Default	Description
mode	string	"e4m3"	FP8 precision mode areal/api/cli_args.py590-593
recipe	string	"delayed"	FP8 scaling recipe areal/api/cli_args.py596-601
param	boolean	False	Keep parameters in FP8 precision to save memory areal/api/cli_args.py612
direct_convert	boolean	True	Use direct FP8 conversion during weight updates areal/api/cli_args.py686

Sources: areal/api/cli_args.py587-689

---

Scheduling Configuration

Scheduling configurations control how training workers are allocated across the cluster.

SchedulingSpec

Parameter	Type	Default	Description
cpu	integer	8	CPU cores required per GPU areal/api/cli_args.py802
gpu	integer	0	GPU units required areal/api/cli_args.py805
mem	integer	32	RAM (GB) required per GPU areal/api/cli_args.py808
task_type	string	"worker"	Choices: "worker", "engine" areal/api/cli_args.py817-821

SchedulingStrategy

Parameter	Type	Default	Description
type	string	"separation"	Choices: "separation", "colocation" areal/api/cli_args.py783-787
target	string | None	None	Role to colocate with areal/api/cli_args.py790

Sources: areal/api/cli_args.py780-886

Scheduling Strategies Diagram

Sources: areal/api/cli_args.py780-796