Megatron vs FSDP vs Archon Backend Trade-offs in AReaL: A Complete Technical Guide

AReaL provides three distinct training backends—MegatronEngine for large-scale tensor and pipeline parallelism with native FP8 support, FSDPEngine for memory-efficient mid-scale training via CPU offloading, and ArchonEngine for experimental research requiring custom pipeline schedules and Ulysses sequence parallelism.

AReaL (inclusionai/areal) is an open-source framework for large language model training that abstracts distributed strategies through three engine implementations. Each backend maps to different PyTorch distributed primitives, making the choice between Megatron vs FSDP vs Archon critical for hardware efficiency and model scalability.

Core Distributed Primitives and Architecture

Each engine in AReaL implements a distinct approach to process-group creation and tensor sharding.

MegatronEngine (Megatron-Core DDP)

MegatronEngine leverages Megatron-Core's DistributedDataParallel combined with a custom parallel-state utility (mpu). In areal/engine/megatron_engine.py, the engine calls mpu.initialize_model_parallel after dist.init_process_group to establish tensor-parallel (TP), pipeline-parallel (PP), and context-parallel (CP) groups【/cache/repos/github.com/inclusionai/areal/main/areal/engine/megatron_engine.py#L72-L80】.

This backend supports the full suite of Megatron parallelism: TP, PP, CP, Expert-Parallel (EP), and Virtual-Pipeline (VPP). It is the only backend that validates quantization_config and propagates FP8 fields directly to TransformerConfig【/cache/repos/github.com/inclusionai/areal/main/areal/engine/megatron_engine.py#L164-L172】.

FSDPEngine (PyTorch FSDP 2)

FSDPEngine targets PyTorch 2.4+ and utilizes FSDP 2 with DeviceMesh and DTensor. The engine constructs a DeviceMesh from ParallelHelper and extracts dp_group, sp_group, and tp_group directly from the mesh【/cache/repos/github.com/inclusionai/areal/main/areal/engine/fsdp_engine.py#L31-L40】.

This backend supports N-D parallelism (Data-Parallel, Sequence-Parallel, and Tensor-Parallel) but does not implement native pipeline parallelism. It offers CPU offload via CPUOffloadPolicy and a memory-efficient load mode that streams weights from disk to CPU before sharding【/cache/repos/github.com/inclusionai/areal/main/areal/engine/fsdp_engine.py#L86-L96】.

ArchonEngine (Pure-Torch DeviceMesh)

ArchonEngine implements a pure-torch native approach using DeviceMesh and a custom ArchonParallelDims configuration. In areal/experimental/engine/archon_engine.py, the engine constructs an ArchonParallelDims object that internally creates a world mesh and cached groups for each dimension【/cache/repos/github.com/inclusionai/areal/main/areal/experimental/engine/archon_engine.py#L95-L104】.

Archon supports DP, TP, CP (Ulysses Sequence Parallel), PP, EP, and ETP (Expert Tensor-Parallel). It provides an offload path via torch-memory-saver and lazy materialisation, plus a torch.compile-ready path via the enable_compile flag in _setup_parallelism【/cache/repos/github.com/inclusionai/areal/main/areal/experimental/engine/archon_engine.py#L54-L60】.

Memory Management and Offloading Capabilities

Memory efficiency varies significantly across the three backends.

MegatronEngine relies on tensor-parallel sharding within Megatron-Core's distributed data parallel. It does not provide automatic CPU offload, requiring instead that models fit within GPU memory or use Megatron's native activation checkpointing.

FSDPEngine provides the most mature offloading capabilities. It supports CPUOffloadPolicy for parameter offloading and implements a memory-efficient load mode that streams weights from disk to CPU before sharding to avoid OOM during initialization【/cache/repos/github.com/inclusionai/areal/main/areal/engine/fsdp_engine.py#L86-L96】.

ArchonEngine offers an offload path via torch-memory-saver and supports lazy materialisation. It also implements Ulysses sequence-parallel padding that reduces per-GPU memory overhead for long-context training【/cache/repos/github.com/inclusionai/areal/main/areal/experimental/engine/archon_engine.py#L71-L78】.

Model Initialization and Weight Loading

Each engine implements distinct strategies for model creation and checkpoint loading.

MegatronEngine uses MBridge to stitch HuggingFace and Megatron models. It loads weights on rank 0 then performs an all-gather to distribute them across tensor-parallel ranks【/cache/repos/github.com/inclusionai/areal/main/areal/engine/megatron_engine.py#L42-L59】.

FSDPEngine builds the model on a meta device, then materializes it after parallelization. This approach avoids OOM during initialization by only allocating memory after sharding decisions are applied【/cache/repos/github.com/inclusionai/areal/main/areal/engine/fsdp_engine.py#L85-L100】.

ArchonEngine follows a similar meta-tensor approach, building a meta model structure then materializing and loading weights after parallelisation. It supports both standard checkpoint loading and experimental weight synchronization utilities【/cache/repos/github.com/inclusionai/areal/main/areal/experimental/engine/archon_engine.py#L62-L70】【/cache/repos/github.com/inclusionai/areal/main/areal/experimental/engine/archon_engine.py#L85-L99】.

Advanced Features: FP8, LoRA, and Pipeline Parallelism

Feature parity varies across quantization, parameter-efficient fine-tuning, and pipeline strategies.

FP8 Quantization is only natively supported by MegatronEngine, which validates quantization_config and propagates FP8 fields to TransformerConfig【/cache/repos/github.com/inclusionai/areal/main/areal/engine/megatron_engine.py#L164-L172】. FSDP and Archon currently rely on torch.amp or manual quantization implementations.

LoRA/PEFT is supported across all three backends. MegatronEngine wraps LoRA via mbridge.AutoBridge and get_peft_model. FSDPEngine provides direct integration in _apply_peft_wrapper, which adds LoraConfig during initialization【/cache/repos/github.com/inclusionai/areal/main/areal/engine/fsdp_engine.py#L80-L88】. ArchonEngine supports LoRA via create_optimizer and create_lr_scheduler utilities.

Pipeline Parallelism is fully implemented in MegatronEngine and ArchonEngine, but absent from FSDPEngine. Megatron provides mature pipeline scheduling with virtual pipeline parallelism, while Archon offers experimental custom schedules like dual_pipe_v【/cache/repos/github.com/inclusionai/areal/main/areal/experimental/engine/archon_engine.py#L54-L60】. ArchonEngine also uniquely supports Ulysses sequence parallelism (CP) for long-context training.

Decision Matrix: When to Use Each Backend

Choose your backend based on model scale, hardware constraints, and experimental requirements.

Scenario	Recommended Backend	Rationale
Very large LLMs (≥70B parameters) requiring TP + PP + MoE	MegatronEngine	Mature Megatron-Core pipeline, high-throughput tensor sharding, and native FP8 support provide optimal throughput for massive models.
Mid-scale clusters (4-8 GPUs) with memory constraints	FSDPEngine	CPU offload policies and memory-efficient weight loading minimize GPU memory footprint while maintaining decent scaling via N-D parallelism.
Research prototypes requiring custom parallel dimensions	ArchonEngine	Pure-torch implementation with ArchonParallelDims enables Ulysses SP, custom pipeline schedules, and torch.compile integration for experimental workflows.
FP8 quantization requirements	MegatronEngine	Only backend that validates quantization_config and propagates FP8 fields automatically to TransformerConfig.
Avoiding custom Megatron model classes	FSDPEngine or ArchonEngine	Both rely on HuggingFace AutoModel* with thin wrappers, eliminating the need to write Megatron-specific model definitions.

Configuration Examples

MegatronEngine with Tensor and Pipeline Parallelism

from areal.api.cli_args import TrainEngineConfig, MegatronEngineConfig, ParallelStrategy
from areal.engine.megatron_engine import MegatronEngine

cfg = TrainEngineConfig(
    path="meta-llama/Meta-Llama-3-8B",
    dtype="bfloat16",
    megatron=MegatronEngineConfig(
        tensor_parallel_size=4,
        pipeline_parallel_size=2,
        virtual_pipeline_parallel_size=1,
        use_fp8=True,
    ),
    parallel_strategy=ParallelStrategy(
        data_parallel_size=1,
        tensor_parallel_size=4,
        pipeline_parallel_size=2,
    ),
)

engine = MegatronEngine(cfg)
engine.create_process_group()
engine.initialize(None, ft_spec)

Source references: Process-group initialization at areal/engine/megatron_engine.py lines 72-80; FP8 configuration handling at lines 164-172.

FSDPEngine with CPU Offload and Memory-Efficient Loading

from areal.api.cli_args import TrainEngineConfig, ParallelStrategy
from areal.engine.fsdp_engine import FSDPEngine

cfg = TrainEngineConfig(
    path="EleutherAI/gpt-neox-20b",
    dtype="float16",
    fsdp=dict(
        offload_params=True,
        memory_efficient_load=True,
    ),
    parallel_strategy=ParallelStrategy(
        data_parallel_size=2,
        tensor_parallel_size=2,
        context_parallel_size=1,
    ),
)

engine = FSDPEngine(cfg)
engine.create_process_group()
engine.initialize(None, ft_spec)

Source references: DeviceMesh creation at areal/engine/fsdp_engine.py lines 31-40; memory-efficient load implementation at lines 86-96.

ArchonEngine with Ulysses Sequence Parallelism

from areal.api.cli_args import TrainEngineConfig, ParallelStrategy
from areal.experimental.engine.archon_engine import ArchonEngine

cfg = TrainEngineConfig(
    path="meta-llama/Meta-Llama-3-8B",
    dtype="bfloat16",
    archon=dict(
        pp_schedule="dual_pipe_v",
        offload_params=True,
    ),
    parallel_strategy=ParallelStrategy(
        data_parallel_size=1,
        tensor_parallel_size=4,
        pipeline_parallel_size=2,
        context_parallel_size=2,
    ),
)

engine = ArchonEngine(cfg)
engine.create_process_group()
engine.initialize(None, ft_spec)

Source references: ArchonParallelDims construction at areal/experimental/engine/archon_engine.py lines 95-104; torch.compile integration at lines 54-60.

CLI Backend Selection

Switch backends without code changes using AReaL's CLI:

# Megatron backend with tensor and pipeline parallelism

areal train --engine megatron --config configs/megatron.yaml

# FSDP backend with CPU offloading

areal train --engine fsdp --config configs/fsdp.yaml

# Archon backend with experimental features

areal train --engine archon --config configs/archon.yaml

Summary

• MegatronEngine provides the most mature support for massive models (70B+ parameters) requiring tensor parallelism, pipeline parallelism, and FP8 quantization, implemented via Megatron-Core's DistributedDataParallel and custom parallel-state utilities.

• FSDPEngine offers the best memory efficiency for mid-scale clusters through CPU offloading and memory-efficient weight streaming, leveraging PyTorch 2.4+ DeviceMesh and DTensor primitives.

• ArchonEngine delivers maximum flexibility for research prototypes, supporting Ulysses sequence parallelism, custom pipeline schedules like dual_pipe_v, and torch.compile integration through pure-torch DeviceMesh and ArchonParallelDims.

Frequently Asked Questions

Which AReaL backend supports FP8 quantization?

MegatronEngine is currently the only backend that natively supports FP8 training. It validates quantization_config and automatically propagates FP8 fields to TransformerConfig during initialization. Both FSDPEngine and ArchonEngine currently rely on torch.amp or require manual quantization implementations.

Can I use pipeline parallelism with FSDP in AReaL?

No. FSDPEngine does not implement pipeline parallelism; it supports only data-parallel, sequence-parallel, and tensor-parallel dimensions via N-D parallelism. For pipeline parallelism, use MegatronEngine (for standard schedules) or ArchonEngine (for custom schedules like dual_pipe_v).

How does ArchonEngine differ from FSDPEngine if both use DeviceMesh?

While both use PyTorch DeviceMesh, ArchonEngine implements a custom ArchonParallelDims abstraction that provides finer-grained control over parallelism dimensions, including Ulysses sequence parallelism and expert tensor-parallelism for MoE models. FSDPEngine uses standard FSDP2 sharding with DTensor and focuses on memory efficiency through CPU offloading rather than custom parallel strategies.

Which backend should I choose for a 70B parameter model on 64 GPUs?

Use MegatronEngine. For models with 70B or more parameters, the combination of tensor parallelism, pipeline parallelism, and native FP8 support in MegatronEngine provides the highest throughput. The MegatronEngineConfig allows fine-grained control over tensor_parallel_size, pipeline_parallel_size, and virtual_pipeline_parallel_size to optimize the 64-GPU cluster layout.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}