Memory Optimizations

This page documents memory optimization techniques implemented in AReaL to enable training large language models within GPU memory constraints. It covers gradient checkpointing, CPU offloading, per-layer optimizer stepping, memory-efficient loading, and MoE-specific optimizations.

For general OOM troubleshooting guidance and configuration recommendations, see page 16.2 (Memory and OOM Issues). For micro-batching configurations, see page 2.7 (MicroBatchSpec and Data Configurations). For parallelism strategies that reduce memory per GPU, see page 8 (Parallelism and Distribution).

Memory Optimization Overview

AReaL provides multiple layers of memory optimization that can be composed to fit models into available GPU memory:

Optimization	Memory Savings	Performance Impact	Configuration
Gradient Checkpointing	~50% activation memory	~20% slower forward/backward	gradient_checkpointing: true
CPU Offload (Params)	~100% parameter memory	Slower forward/backward	offload_params: true
CPU Offload (Optimizer)	~2x parameter memory	Very slow optimizer step (without per-layer)	Automatic with offload_params
Per-Layer Optim Step	No additional savings	~10x faster than CPU Adam	per_layer_optim_step: true
Memory-Efficient Load	Peak init memory	Slower initialization	memory_efficient_load: true
Low-Precision Optimizer	~50% optimizer memory	Negligible	optimizer.type: adam_bf16
Micro-Batch Tuning	Linear with batch size	Affects throughput	mb_spec.max_tokens_per_mb

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L1-L201" min=1 max=201 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L44-L90" min=44 max=90 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>

Memory Hierarchy and Optimization Points

Title: "Memory Distribution and Optimization Mapping"

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L6-L40" min=6 max=40 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L44-L90" min=44 max=90 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>

Gradient Checkpointing

Gradient checkpointing (also called activation checkpointing) trades compute for memory by recomputing activations during the backward pass instead of storing them.

Configuration

In the actor or critic engine configuration:

Implementation

AReaL's gradient checkpointing is applied at the transformer layer level. In FSDPEngine, this is handled during model parallelization where the wrap_policy determines which modules (typically transformer layers) receive activation checkpointing <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/__init__.py#L62-L85" min=62 max=85 file-path="areal/engine/fsdp_utils/__init__.py">Hii</FileRef>. In the ArchonEngine, activation checkpointing is configured via ActivationCheckpointConfig.

Memory Impact: Reduces activation memory significantly by only storing activations for a subset of layers (checkpoints) and recomputing the rest on-demand <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L105-L110" min=105 max=110 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>.

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L105-L110" min=105 max=110 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/__init__.py#L62-L85" min=62 max=85 file-path="areal/engine/fsdp_utils/__init__.py">Hii</FileRef>

CPU Offloading

CPU offloading moves parameters and optimizer states from GPU to CPU memory, streaming them to GPU only when needed for computation.

Parameter Offloading

Enabled via offload_params: true in the engine configuration. This utilizes CPUOffloadPolicy in FSDP2 <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/__init__.py#L10-L13" min=10 max=13 file-path="areal/engine/fsdp_utils/__init__.py">Hii</FileRef>.

When enabled:

1. Parameters reside in CPU RAM.
2. During forward/backward passes, FSDP fetches parameters to GPU.
3. Gradients are computed on GPU and then reduced.
4. Optimizer states are managed on the CPU <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L167-L181" min=167 max=181 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>.

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L167-L181" min=167 max=181 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/__init__.py#L10-L13" min=10 max=13 file-path="areal/engine/fsdp_utils/__init__.py">Hii</FileRef>

Memory-Efficient Loading

This reduces peak GPU memory during initialization by loading weights on CPU first and broadcasting them after sharding <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L203-L212" min=203 max=212 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>.

Title: "Model Loading Data Flow"

For LLMs, fsdp2_load_full_state_dict uses StateDictOptions with broadcast_from_rank0=True to ensure only one rank needs to load the full checkpoint into CPU memory <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/__init__.py#L109-L147" min=109 max=147 file-path="areal/engine/fsdp_utils/__init__.py">Hii</FileRef>.

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L203-L223" min=203 max=223 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/__init__.py#L109-L147" min=109 max=147 file-path="areal/engine/fsdp_utils/__init__.py">Hii</FileRef>

Per-Layer Optimizer Step

Per-layer optimizer step is a critical optimization that enables practical use of CPU-offloaded optimizers. Without it, standard CPU Adam is extremely slow. With per-layer stepping, optimizer updates achieve significant speedups by streaming states per-layer to GPU <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L167-L181" min=167 max=181 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>.

Architecture

The PerLayerOptimWrapper class implements pipelined per-layer optimizer stepping <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L445-L487" min=445 max=487 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.

Key Components:

• PerLayerOptimWrapper: Orchestrates the layer-by-layer update loop by identifying FSDPModule instances <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L445-L487" min=445 max=487 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.
• OptimKernel / AdamKernel: GPU-accelerated logic for single layers <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L348-L443" min=348 max=443 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.
• ParamTransferState: Manages async H2D/D2H transfers for a layer <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L327-L337" min=327 max=337 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.

Configuration

Title: "Pipelined Per-Layer Optimizer Execution"

The wrapper identifies layers by finding FSDPModule instances in the model hierarchy <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L489-L546" min=489 max=546 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>. It uses pinned memory to ensure non-blocking transfers <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L548-L566" min=548 max=566 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L320-L757" min=320 max=757 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L167-L181" min=167 max=181 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/tests/test_per_layer_optim_step.py#L124-L144" min=124 max=144 file-path="tests/test_per_layer_optim_step.py">Hii</FileRef>

Optimizer Precision Optimizations

AnyPrecisionAdamW

AReaL includes AnyPrecisionAdamW, which allows storing momentum and variance in bfloat16 to save 50% of optimizer state memory while using Kahan summation to maintain numerical stability <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L44-L90" min=44 max=90 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.

Memory Impact:

• Momentum (exp_avg): bfloat16 (2 bytes/param) <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L144-L144" min=144 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.
• Variance (exp_avg_sq): bfloat16 (2 bytes/param) <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L147-L147" min=147 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.
• Kahan Compensation (Optional): bfloat16 (2 bytes/param) <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L150-L153" min=150 max=153 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>.

Lightweight Optimizers

Users can switch to SGD or AdamW_bf16 via configuration:

SGD and AdamW_bf16 use less memory than the default AdamW <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L189-L201" min=189 max=201 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>.

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L189-L201" min=189 max=201 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/engine/fsdp_utils/optimizer.py#L44-L189" min=44 max=189 file-path="areal/engine/fsdp_utils/optimizer.py">Hii</FileRef>

Micro-Batch Size Tuning

The MicroBatchSpec controls the granularity of training steps <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/utils/data.py#L18-L18" min=18 file-path="areal/utils/data.py">Hii</FileRef>.

max_tokens_per_mb: This is the primary parameter for controlling training memory. It limits the number of tokens in a single forward/backward pass. If a global batch contains more tokens, it is split into multiple micro-batches <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L87-L95" min=87 max=95 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>.

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/docs/en/best_practices/handling_oom.md?plain=1#L87-L95" min=87 max=95 file-path="docs/en/best_practices/handling_oom.md">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/utils/data.py#L18-L18" min=18 file-path="areal/utils/data.py">Hii</FileRef>

MoE-Specific Optimizations

For Mixture-of-Experts (MoE) models, AReaL provides MoEWeightConverter to handle the transition between sharded 3D expert weights and standard 2D HuggingFace formats without incurring full-model memory overhead on a single rank <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/experimental/models/archon/moe_weight_converter.py#L46-L61" min=46 max=61 file-path="areal/experimental/models/archon/moe_weight_converter.py">Hii</FileRef>.

Key Features:

• DTensor-Aware Splitting: calculate_indices_from_placements determines which experts belong to which GPU based on DeviceMesh sharding <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/experimental/models/archon/moe_weight_converter.py#L124-L146" min=124 max=146 file-path="areal/experimental/models/archon/moe_weight_converter.py">Hii</FileRef>.
• Stateful Conversion: MoEConversionState tracks shapes and placements to ensure consistency between saving and loading <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/experimental/models/archon/moe_weight_converter.py#L19-L37" min=19 max=37 file-path="areal/experimental/models/archon/moe_weight_converter.py">Hii</FileRef>.

Sources: <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/areal/experimental/models/archon/moe_weight_converter.py#L1-L180" min=1 max=180 file-path="areal/experimental/models/archon/moe_weight_converter.py">Hii</FileRef>, <FileRef file-url="https://github.com/inclusionAI/AReaL/blob/2e12c19b/tests/experimental/archon/torchrun/run_ep_tests.py#L42-L123" min=42 max=123 file-path="tests/experimental/archon/torchrun/run_ep_tests.py">Hii</FileRef>