AXLearn: Modular Large Model Training on Heterogeneous Infrastructure

ArXiv: 2507.05411

🎯 Pitch

Adding features like Rotary Position Embeddings to other training systems requires modifying hundreds of modules, but AXLearn achieves the same in just 10 lines by enforcing strict encapsulation—yielding constant code complexity regardless of system scale. This modular design also generalizes to inference out-of-the-box, achieving 2.8× higher throughput than vLLM on TPUs, while matching or beating the training speed of hardware-specialized alternatives.

1. Executive Summary

AXLearn introduces a modular, hardware-agnostic training system for large deep learning models that enforces strict encapsulation across all software components, enabling rapid model development without increasing system complexity as modules scale. The paper evaluates AXLearn's modularity against state-of-the-art systems (Megatron-LM, DeepSpeed, TorchTitan, MaxText, Flax, Praxis) and its training performance across GPUs, TPUs, and AWS Trainium2 on Llama2-7B, Llama2-70B, and Qwen-3 30B-A3B models. AXLearn achieves constant LoC-complexity—requiring zero changes to existing interfaces when integrating features like Rotary Position Embeddings or Mixture of Experts across hundreds of modules (versus O(NM) complexity and hundreds to thousands of lines of code in other systems)—while matching or exceeding the training throughput of hardware-optimized alternatives, reaching 66.2% MFU on TPU v5p-512 and 3.0M tokens/s on H100 GPUs for Llama2-7B. The paper also demonstrates that AXLearn's modular design extends to inference, achieving 2.8× higher throughput than vLLM on TPUs for Llama2-7B without specialized inference engineering, establishing that strict encapsulation need not trade off performance—though achieving peak hardware utilization still requires backend-specific kernel implementations.

2. Context and Motivation

The Core Problem: Modularity and Hardware Flexibility Are in Tension in Large Model Training

The fundamental challenge this paper addresses is deceptively familiar to anyone who has built production software: how do you design a system where adding new features doesn't require changing existing code? In most software domains, modularity through strict encapsulation is a solved problem — you define an interface, implement it, and swap implementations without touching callers. But in large-scale deep learning training, the paper argues, this principle has been systematically neglected, with profound consequences for the velocity of ML research and the operational flexibility of organizations that deploy models at scale.

The gap is not that existing systems lack any modularity. Megatron-LM, DeepSpeed, TorchTitan, MaxText, and others all allow some degree of configuration and reuse. The gap is that none of them enforce strict encapsulation, meaning that changes to one component routinely propagate through the module hierarchy, requiring modifications to parent modules, sibling modules, and configuration structures far removed from the change itself. The paper formalizes this as a complexity scaling problem: in existing systems, the lines of code (LoC) required to integrate a feature like Rotary Position Embeddings (RoPE) or Mixture of Experts (MoE) scales with $O(N)$ or $O(NM)$, where $N$ is the number of modules in the system and $M$ is the number of feature variants. AXLearn's claim is that this complexity should be constant — $O(1)$ — and that achieving this requires a systematic commitment to encapsulation that goes beyond what any prior training framework has attempted.

Why This Problem Matters: Velocity at Scale

The practical significance of this gap becomes clear when we examine the operational context the paper describes (Section 7.4). Apple trains "thousands of models involving hundreds of engineers" across "tens of different hardware clusters." In this environment, model engineers iterate on architecture variants, hyperparameters, and training methods daily. Every line of code that must be changed to accommodate a new feature is a line of code that can introduce bugs, create merge conflicts across a large team, and slow the pace of experimentation.

The paper quantifies this through a specific, concrete metric: LoC-complexity, defined as the asymptotic LoC changes required to re-parameterize a system's API to support a new feature, measured across existing modules (not the new feature implementation itself). This is a clever framing because it isolates the integration cost of a feature from its implementation cost. The implementation of RoPE or MoE is roughly comparable across systems (up to a constant), so the differentiating factor is how much existing code must be touched to wire it in.

The numbers in Table 2 are striking. In a production setting with 20 model variants and 10 attention variants, integrating a single RoPE variant requires an estimated 400 LoC changes in Megatron-LM, 320 in DeepSpeed, 240 in TorchTitan, 600 in Flax, 300 in Praxis, and 200 in MaxText. For MoE, the situation is even worse in some systems: the paper estimates 4,000 LoC for DeepSpeed, driven by the need to subclass each model from a custom DSMoETransformerModelBase and re-implement most methods. In AXLearn, both features require 0 LoC changes to existing interfaces — a 10-line config modifier suffices, and the same snippet configures over 1,000 experiments internally.

This is not merely a developer convenience. The paper argues that the LoC-complexity of existing systems fundamentally limits the rate of experimentation at organizations that maintain large codebases with many model variants. If a researcher wants to test a new attention mechanism across 20 model architectures, the $O(N)$ complexity means touching every model's implementation — a task that may take days or weeks of careful, error-prone refactoring. The $O(1)$ complexity of AXLearn means the same change can be expressed once and applied universally through a configuration modifier. For an organization shipping AI features to over a billion users (as the paper notes in Section 7.4), this velocity translates directly to competitive advantage.

The Hardware-Agnostic Imperative

The second major gap the paper addresses is hardware lock-in. The paper is unusually frank about the business reality: "As one of the largest technology companies, Apple cannot feasibly rely on a single hardware platform for all of our machine learning workloads" (Section 2.2). This is not a theoretical concern — it reflects supply chain constraints, pricing dynamics, and the strategic risk of depending on a single vendor. The paper notes that AWS Trainium2 "did not exist when we first began development," yet AXLearn was able to support it as one of the first deep learning systems at scale precisely because of its hardware-agnostic design.

Existing systems largely optimize for specific backends. Megatron-LM, the paper notes, "has a vested interest in optimizing for Nvidia GPUs" with carefully tuned CUDA optimizations that "do not directly apply to other hardware." Haiku, Flax, Pax, and MaxText are "mostly optimized for Google TPUs." Table 1 makes this explicit: of the nine systems surveyed, only AXLearn supports GPU, TPU, and Trainium with full 3D parallelism and modular configuration.

The paper's insight is that hardware agnosticism is not just about choosing a portable compiler (JAX/XLA in this case) — it requires that all optimizations be expressible as configuration, not code. If parallelism strategies, rematerialization policies, quantization settings, and kernel selection are baked into model implementations as conditional logic (the pattern in most existing systems), then supporting a new hardware backend means modifying every model implementation that needs to run on it. AXLearn's solution is the "mesh rule" concept (Section 4.2, Appendix A): a mapping from hardware instance type patterns to config modifiers that apply transformations like "use FSDP within TPU v5e slices and data parallelism across slices, with activation offloading and INT8 training," all expressed in roughly 10 lines of configuration. Because modules are encapsulated and unaware of parallelism details, these rules can be applied universally without touching model code.

Where Prior Approaches Fall Short

The paper identifies specific failure modes in existing systems that prevent them from achieving constant complexity, and it's worth understanding the mechanism behind each:

Subtyping violations of encapsulation (Section 2.1). The dominant design pattern in existing ML frameworks is subtyping: a layer inherits from a base class, overrides methods, and accesses parent module state through instance attribute traversal. The paper demonstrates through the MoE case study why this pattern scales poorly. In DeepSpeed, replacing a feed-forward network with MoE requires modifying not just the FFN layer but also its parent layer (which must instantiate the new subtype), and that parent's parent, and so on up the hierarchy. The paper formalizes this: "by induction, it's easy to see how such a change can compound to changes to multiple modules across the subtype hierarchy." The QwenV2 to QwenV2MoE transition in DeepSpeed requires over 200 LoC, not 4 as might be expected from a naively modular design.

The deeper issue is that subtyping couples the interface of a module to the interfaces of its children, because parent modules must be aware of child types at instantiation time. Composition with strict encapsulation breaks this coupling: the parent only knows that a child conforms to a Config interface, and child implementations are swappable without the parent's awareness.

Config flattening limits composability (Section 4.1, 7.1). TorchTitan and DeepSpeed adopt "flat" configuration layouts where all parameters — model, optimizer, checkpointer, trainer — live in a single monolithic config class. The paper argues this has "significant ramifications in extensibility." When a new feature like RoPE is added, every model's ModelArgs subclass must be modified to include the new parameters, and every attention layer that consumes those parameters must be updated to handle them. The analysis in Appendix B traces this through TorchTitan: each model has its own Attention implementation that conditions on RoPE configs to instantiate the appropriate embedding layer, meaning that both the config class and the attention layer change for each model variant — yielding $O(NM)$ complexity.

AXLearn's alternative is hierarchical, partially-specified configs. A TransformerLayer config contains a feed_forward: FeedForwardLayer.Config field, but does not specify the FFN's hyperparameters directly. Instead, those are resolved at instantiation time when the parent sets input_dim on the child. This means swapping FeedForwardLayer for MoELayer requires changing only the child config reference — the parent TransformerLayer config is untouched.

State management breaks modularity (Section 4.3). JAX's functional programming model requires explicit state management: parameters, PRNG keys, and summaries must be threaded through function calls as inputs and outputs. In libraries without an abstraction for this (early Flax, for example), users must manually pass state through the module hierarchy. This means that when a new layer is added or removed, the state-passing code in every intermediate module must be updated — another $O(N)$ scaling violation.

AXLearn's InvocationContext solves this by maintaining a stack of module states that is transparently pushed and popped during module invocation, decoupling state management from the module hierarchy itself. A module can retrieve shared state (like tied weights) by traversing the context stack rather than reaching through parent instance attributes, preserving encapsulation.

Parallelism and memory optimizations are entangled with model code (Section 4.2). In most systems, parallelism strategy selection requires modifying model implementations. Megatron-LM, for instance, propagates tensor-parallelism-specific arguments through its module hierarchy, meaning that switching parallelism strategies involves code changes across the model. The paper notes that "different parallelization strategies meant that logic had to be rewritten, with inevitable interactions across modules" (Section 7.4). AXLearn separates parallelism specification into config modifiers that the composer applies to the layer graph, without the layer implementations themselves being aware of the parallelism strategy.

How This Paper Positions Itself

The paper positions AXLearn not as a competitor to Megatron-LM or DeepSpeed on raw performance optimization, but as the only system that achieves constant complexity scaling while matching state-of-the-art performance. This is a deliberate tradeoff: the paper acknowledges that Megatron-LM achieves slightly higher throughput on H100 GPUs "because PyTorch currently has finer-grained scheduling capability over XLA" (Section 7.2), but argues that "this is a trade-off we are willing to take" because XLA enables hardware agnosticism.

The theoretical contribution is the LoC-complexity framework itself — a formal way to quantify the modularity of a training system by measuring the asymptotic integration cost of new features. The paper argues that this metric explains why composition (AXLearn's approach) should be preferred over subtyping (the dominant approach), and provides a rigorous basis for comparing systems that goes beyond "rules of thumb" about good design.

The practical contribution is demonstrating that the combination of (1) strict encapsulation, (2) GSPMD/XLA for compiler-injected parallelism, (3) hierarchical configs with programmatic modification, and (4) transparent state management through InvocationContext can together achieve constant $O(1)$ complexity without sacrificing performance. The paper shows this across four hardware backends with competitive MFU numbers (Table 3), weak scaling to 32,768 TPUs with near-linear efficiency (Figure 4), and even inference workloads where the modular design yields 2.8× throughput over vLLM on TPUs (Table 4).

The paper also positions itself within the broader trajectory of ML infrastructure evolution. Section 7.4 describes the transition from PyTorch (imperative, mutable state) to JAX/XLA (functional, compiler-driven), and the design decisions that made that transition viable — particularly InvocationContext to restore imperative-style programming within a functional framework. This narrative positions AXLearn as a synthesis: it adopts the compiler-first approach of JAX/XLA (enabling hardware agnosticism) while providing abstractions that preserve the developer experience of imperative frameworks, all within a strictly encapsulated architecture that existing systems have not achieved.

3. Technical Approach

3.1 Reader Orientation

AXLearn is a large-scale deep learning training system built on JAX and XLA that enforces strict encapsulation across every component, from neural network layers to input pipelines to checkpointers. The system solves the problem that adding a new feature (like Rotary Position Embeddings or Mixture of Experts) to existing training frameworks requires modifying code scattered across dozens or hundreds of modules, because those frameworks violate encapsulation through subtyping and config flattening — AXLearn reduces this integration cost from $O(N)$ or $O(NM)$ lines of code changes to $O(1)$ by treating every module as a black box with a well-defined configuration interface, composing them hierarchically, and applying system-wide transformations (parallelism, rematerialization, hardware-specific optimizations) as config modifiers that traverse the hierarchy externally rather than being baked into module implementations.

3.2 Big-Picture Architecture (Diagram in Words)

The AXLearn system processes a training job through seven major stages, flowing from user specification to distributed execution:

1. User Script — a Python file where the engineer defines the training configuration by composing hierarchical Config objects for the trainer, model, input pipeline, and learning algorithm. The script never directly implements parallelism, memory optimization, or hardware-specific logic.

2. AXLearn Composer (Config Generator) — reads the user's hierarchical config, applies hardware-specific "mesh rules" (config modifiers keyed by accelerator type) to inject parallelism sharding annotations, rematerialization policies, quantization settings, and custom kernel selections into the layer graph, then materializes the complete JAX program with all compiler annotations.

3. AOT Compilation — the JAX program undergoes Ahead-of-Time compilation locally (on CPU if desired), allowing the user to check memory utilization, FLOP counts, and OOM errors without touching accelerator hardware.

4. JAX Compiler (XLA) — the annotated JAX program is compiled by XLA into hardware-specific accelerator programs (e.g., CUDA kernels for GPUs, Pallas kernels for TPUs, NKI kernels for Trainium), with the compiler handling sharding propagation and operator fusion automatically.

5. Accelerator Programs — the compiled executable runs on each accelerator in the distributed cluster, with communication collectives (all-reduce, all-gather) injected by GSPMD based on the sharding annotations.

6. AXLearn Runtime — orchestrates execution across the cluster (typically via Kubernetes), providing asynchronous checkpointing to cloud storage (S3, GCS), monitoring via JAX's profiler and a configurable watchdog that detects stalled hosts, and fault tolerance through slice-level hot-swap that replaces failed nodes with over-provisioned spares.

7. InvocationContext Stack — a transparent runtime mechanism that maintains a stack of module states during forward and backward passes. When a parent module invokes a child, a context is pushed containing the child's parameters, a split PRNG key, and a fresh output collection; when the child returns, the context is popped and child outputs merge into the parent's collection. This allows modules to be implemented in imperative style while preserving JAX's functional purity requirements.

The critical architectural invariant is that modules never reference their parents or siblings directly — all inter-module communication flows through the InvocationContext stack, which is programmatically traversable for features like tied weights, and all system-level transformations (parallelism, remat, kernel selection) are applied by the composer via config modifiers that operate on the config tree from the outside, not from within module implementations.

3.3 Roadmap for the Deep Dive

• First, the hierarchical configuration system — how modules define their interfaces, how configs are partially specified and resolved, and how config traversal enables $O(1)$ feature integration. This is the foundation because everything else (parallelism, remat, hardware optimization) operates by modifying these configs externally.

• Second, the config modifier mechanism — how mesh rules and replace_config traversals apply system-wide transformations without touching module code. This is the mechanism that makes $O(1)$ complexity concrete.

• Third, parallelism and sharding — how AXLearn separates parallelism specification (what the user writes) from parallelism implementation (what the compiler injects), and why this separation requires config-based sharding annotations rather than code-based sharding logic.

• Fourth, memory optimization — how rematerialization and optimizer state offloading are configured via tagged points in the module hierarchy, enabling per-hardware strategies without modifying model code.

• Fifth, hardware-specific kernel selection — how AXLearn dispatches to backend-specific attention kernels (cuDNN, Pallas, NKI) through drop-in layer replacements expressed as mesh rules, and why this requires layers to be sharding- and remat-aware.

• Sixth, AOT compilation — the ahead-of-time compilation workflow that catches memory and sharding errors locally, and why sharing the exact codepath between AOT and actual training is the property that makes this reliable.

• Seventh, the InvocationContext state management system — how AXLearn reconciles JAX's functional purity requirements with the imperative programming style users expect, and specifically how the context stack enables encapsulation-preserving features like tied weights and summary collection.

3.4 Detailed, Sentence-Based Technical Breakdown

This is a systems design paper whose core idea is that strict encapsulation across all components of a deep learning training system — combined with config-based external application of parallelism, memory, and hardware optimizations — can reduce feature integration complexity from $O(N)$ or $O(NM)$ to $O(1)$ without sacrificing training performance.

Hierarchical Configuration System

The configuration system is the central mechanism through which AXLearn achieves constant integration complexity. Instead of a flat config file that enumerates every hyperparameter across every component, AXLearn represents a training job as a hierarchical tree of partially-specified configuration objects, where each node in the tree corresponds to a Module and its Config subclass, and parent nodes resolve unspecified parameters of their children at instantiation time rather than at config definition time.

Module and Config structure. Every module in AXLearn follows a consistent two-class pattern (Section 4.1). A module is defined as a Python class inheriting from Module, containing an inner Config class that inherits from Module.Config. The Config class declares child modules as typed attributes:

class TransformerLayer(Module):
    class Config(Module.Config):
        self_attention: AttentionLayer.Config
        feed_forward: FeedForwardLayer.Config

The key design decision here is that TransformerLayer.Config declares that it has a feed_forward child, and what interface that child must satisfy (the FeedForwardLayer.Config type), but does not specify the child's internal hyperparameters — hidden dimension, number of layers, activation function, or any other implementation detail. This is what the paper means by "strict encapsulation": the parent's config is independent of the child's internal structure, and conversely the child's config can be modified or swapped entirely without the parent being aware.

Partial specification and resolution. Configs are often initially partially specified. For example, a FeedForwardLayer config might declare that its hidden_dim should be computed as a function of an (as yet unknown) input_dim:

cfg.feed_forward.hidden_dim = lambda input_dim: 4 * input_dim

The input_dim is set at instantiation time by the parent, which knows the output dimension of the previous layer:

class TransformerLayer(Module):
    def __init__(self, cfg):
        cfg.feed_forward.set(input_dim=cfg.input_dim)
        self._add_child("feed_forward", cfg.feed_forward)

This two-phase initialization — declare configs with potentially unresolved references, then resolve them top-down at instantiation — means that configs "often only need to be specified once, commonly at the layers near the root of the tree" (Section 4.1). A researcher defining a new model variant can reuse entire subtrees of configuration (e.g., a standard attention mechanism, a standard optimizer) and only specify the components that differ, with dimensional consistency enforced by the parent's resolution logic rather than by manual bookkeeping.

Contrast with config flattening. The paper explicitly contrasts this with the "flat" config approach used by TorchTitan and DeepSpeed. In a flat config, all hyperparameters across all components live in a single ModelArgs or Config dataclass. When a new feature like MoE is added, every model's monolithic config class must be modified to include the new MoE-specific fields (e.g., num_experts, expert_capacity, routing_strategy), and every layer that consumes those fields must be updated to propagate them. The paper traces this through TorchTitan: "each model has its own Attention implementation... conditions on the value of RoPE configs (e.g. rope scaling) to decide which child RoPE layer to instantiate" (Appendix B). This coupling between config structure and the number of model variants is what produces the $O(N)$ scaling — each of the 20 model variants in the paper's production estimate must independently update its config class and its layer implementations.

In AXLearn, adding MoE requires only that a new MoELayer.Config exists with an interface compatible with FeedForwardLayer.Config (same input/output dimensional interface). The parent TransformerLayer.Config already declares feed_forward: FeedForwardLayer.Config — the MoELayer.Config can be directly substituted because the type annotation only constrains the interface, not the concrete implementation. No change to TransformerLayer or any ancestor is needed.

Config traversal as a universal modification primitive. Because configs form a tree and nodes are encapsulated, arbitrary modifications can be applied by traversing the tree and applying a transformation function at each node. The paper calls such a function a "config modifier" and provides a concrete implementation:

def replace_config(cfg, tgt, new_cfg):
    def enter_fn(child):
        for key, value in child.items():
            if isinstance(value, tgt.Config):
                new_cfg.set(**value.items())
                child.set(key, new_cfg)
    cfg.visit(enter_fn=enter_fn)

This function visits every node in the config tree, checks whether the node's value is an instance of the target config type, and if so, replaces it with a new config that inherits any settings from the original via set(**value.items()) before applying overrides. The paper emphasizes that this "roughly 10-line code snippet is used to apply MoE to over 1,000 experiment configs, without any additional changes to other modules" (Section 4.1).

What makes this powerful is that it operates externally to any module implementation. The replace_config function doesn't need to know anything about TransformerLayer, LLaMAModel, or any specific architecture — it only needs to know the type of config to replace (FeedForwardLayer) and the replacement (MoELayer). This is composition operating on the configuration graph rather than on the class hierarchy, which is why it achieves $O(1)$ complexity: the modifier is written once and applies to any config tree containing the target type, regardless of the tree's depth or the number of distinct model architectures.

Mesh Rules: Hardware-Dependent Optimization as Configuration

The paper introduces "mesh rules" as the mechanism for applying hardware-specific optimizations without modifying model code (Section 4.2, Appendix A). A mesh rule is a mapping from a regular expression matching accelerator instance types to a list of config modifiers to apply. When the AXLearn composer prepares a JAX program for a specific hardware target, it evaluates the mesh rules, finds the first matching rule, and applies the associated modifiers to the config tree.

Concrete example. Appendix A provides a complete example. When training on TPU v5e slices, the following modifiers are applied:

[("tpu-v5e-256-*",
  [MeshShapeModifier.default_config().set(
      mesh_shape=mesh(data=-1, fsdp=256)),
   RematSpecModifier.default_config().set(
      remat_policies={
          "model.decoder.transformer.layer":
          RematSpec(policy=offload_dots)}),
   INT8ConfigModifier.default_config()])]

When training on H100 GPUs instead, a different rule fires:

[("gpu-H100-*",
  [MeshShapeModifier.default_config().set(
      mesh_shape=mesh(fsdp=-1, model=8)),
   RematSpecModifier.default_config().set(
      remat_policies={
          "model.decoder.transformer.layer":
          RematSpec(policy=save_qkvoflash)}),
   FP8ConfigModifier.default_config().set(
      fp8_amax_history_length=128)])]

Each modifier in the list corresponds to a different system concern: MeshShapeModifier configures the GSPMD device mesh (how accelerators are partitioned into data-parallel, tensor-parallel, and pipeline-parallel groups), RematSpecModifier configures which activations to save versus recompute during backpropagation, and INT8ConfigModifier/FP8ConfigModifier configure quantization. The paper notes that "these configs are all that are necessary to apply per-target optimizations, allowing users to scale training on different platforms with ease" (Appendix A).

Why this requires encapsulation. The mesh rule mechanism only works because AXLearn layers are already "sharding and remat aware" — meaning they have internal hooks where sharding annotations and rematerialization policies can be attached, but those hooks are exposed through the config interface, not through the Python class hierarchy. If layers instead baked parallelism logic into their implementations (as DeepSpeed and Megatron-LM do, propagating tensor-parallelism arguments through init signatures), then switching between mesh(data=-1, fsdp=256) (FSDP within TPU v5e slices) and mesh(fsdp=-1, model=8) (8-way tensor parallelism within H100 nodes) would require modifying the layer implementations themselves — exactly the $O(N)$ scaling the paper seeks to avoid.

The paper argues that this design is only possible because of GSPMD (Xu et al., 2021), which separates the specification of sharding from the implementation of communication. GSPMD allows the user (or the config system) to annotate tensors with sharding constraints, and the XLA compiler automatically inserts the necessary collective communication operations (all-reduce, all-gather, reduce-scatter) between operations. Without GSPMD, the system would need to manually insert communication primitives in the forward and backward passes, which would couple the parallelism strategy to the layer implementation code. The paper describes this as a "bet on the compiler-first approach" made in late 2021 (Section 7.4).

Config-Based Parallelism and Sharding

Section 4.2 describes how AXLearn "natively integrates parallelism support in every relevant layer for all common parallelism strategies." This is a stronger claim than it might appear — it means that the layer library includes layers like Linear, Attention, Embedding, and LayerNorm that are implemented once and can be used with any combination of data parallelism, fully-sharded data parallelism (FSDP), tensor model parallelism, pipeline parallelism, sequence parallelism, and expert parallelism, without any parallelism-specific code in the layer's forward method.

The mechanism: sharding propagation via annotations. When the composer applies a MeshShapeModifier (from a mesh rule or directly from user config), it annotates specific tensors in the layer graph with GSPMD sharding constraints. For example, in FSDP mode, the weight tensors of linear layers are annotated as sharded across the data-parallel dimension, and the XLA compiler automatically inserts all-gather operations before each forward pass and reduce-scatter after each backward pass. The layer's forward implementation remains unchanged — it simply calls jnp.dot(input, weight), and the compiler handles the distributed execution.

Granular user control. The paper emphasizes that users have "granular control over how specific parameters in specific layers are partitioned" (Section 4.2). This granularity is essential for advanced parallelism strategies: for instance, in a transformer with tensor model parallelism, the attention QKV projection might be sharded along the head dimension across tensor-parallel workers, while the output projection is sharded along the hidden dimension, and the feed-forward layers use yet another sharding scheme. In Megatron-LM, these choices are hard-coded into the layer implementations through ColumnParallelLinear and RowParallelLinear subclasses. In AXLearn, they are expressed as sharding annotations on the standard Linear layer's config, with different annotations automatically selected by the mesh rule for the target hardware.

Contrast with PyTorch and Flax. The paper notes that "in Flax or PyTorch, sharding is not a native concept in the layer library, and thus code changes may be necessary depending on parallelism strategy" (Section 4.2). This is a crucial point: in PyTorch FSDP, the sharding is applied at the torch.distributed level by wrapping modules, not by annotating individual tensors. This means that if you want to switch from FSDP to tensor model parallelism, you must change the model code to use different linear layer variants. In JAX with GSPMD, the same Linear module works for both because the sharding is a property of the tensor annotations, not the module implementation.

Memory Optimizations: Rematerialization and Offloading

AXLearn provides two memory optimization mechanisms, both configurable through the config tree without modifying layer implementations (Section 4.2).

Rematerialization (activation checkpointing). During backpropagation, intermediate activations from the forward pass are needed to compute gradients. Storing all activations in HBM (high-bandwidth memory) is prohibitively expensive for large models, so training systems selectively recompute some activations during the backward pass rather than storing them. The paper calls this "rematerialization" (or "remat"), following JAX terminology.

AXLearn's approach is to tag specific points in the module hierarchy with named identifiers, which users can target in rematerialization policies:

"In AXLearn, common remat points (such as the attention QKV projections and output) are 'tagged' with names, such that users can selectively target remat points to decide which activations to save in accelerator memory, offload to CPU memory (if supported by hardware), or recompute."

The RematSpecModifier in the mesh rule example above shows this in action. The modifier specifies a remat_policies dictionary mapping path patterns (like "model.decoder.transformer.layer") to RematSpec objects that encode the policy. The policy=offload_dots setting means that dot product activations should be offloaded to host memory rather than recomputed (which the paper notes is beneficial on TPU v5e where HBM is limited). The policy=save_qkvoflash setting means that query, key, value, and output projection activations should be saved in HBM (which is beneficial on H100 GPUs where recomputing them is expensive relative to the memory cost).

The paper also notes that "users can also employ programmatic remat strategies, such as only saving the output of linear layers." This flexibility is possible because the remat policy is applied externally to the config tree — the layer implementations themselves contain only the named tags, not the decision of what to save versus recompute.

Optimizer state offloading. For very large models (hundreds of billions of parameters) on hardware with limited HBM (like TPU v5e), even sharded optimizer states may exceed memory capacity. AXLearn supports offloading optimizer states to CPU memory, which the paper describes as "essential for training models with more than hundreds of billions of parameters on certain platforms like TPU v5e where HBM is limited, and where sharding beyond a certain point can be inefficient." This is configured through the same mesh rule mechanism, though the paper provides fewer details on the implementation.

The key property is that both memory optimizations are strictly config modifications — they do not require changes to layer implementations because the layers expose tagged points that the external optimization system targets. This is another instance of the $O(1)$ complexity principle: adding a new rematerialization strategy for a new hardware platform requires adding a mesh rule entry with a different RematSpecModifier, not modifying every layer that might need a different policy.

Hardware-Dependent Kernel Selection

Section 4.2 describes how AXLearn achieves high performance across GPU, TPU, and Trainium while maintaining a single codebase. The mechanism relies on two properties: (1) specific layer variants (like FlashAttention) that are drop-in replacements for the default implementation, and (2) mesh rules that select which variant to use based on the target hardware.

Attention kernel dispatch. The paper provides a concrete example with the attention layer:

"AXLearn provides a FlashAttention (Dao et al., 2022; Dao, 2023) layer, which can be used as a drop-in replacement for the default attention layer. Like above, this can be expressed as a config modifier in a mesh rule."

Behind the scenes, the FlashAttention layer dispatches to different kernels depending on the backend. The paper lists three dispatch paths:

• GPU: uses cuDNN (Chetlur et al., 2014) "when possible, falling back to a custom Pallas (pallas, 2025) kernel for cases like block-sparse attention where cuDNN is not supported."
• AWS Trainium: uses "the Nki kernel from AWS Neuron Toolkit (AWS, 2025)."
• TPU: uses "the SplashAttention Pallas kernel in JAX (Bradbury et al., 2018)."

The critical design point is that the user never writes if backend == "gpu": use_cudnn() elif backend == "tpu": use_splash_attention(). Instead, they write a mesh rule that says "when on GPU, replace the default attention config with FlashAttention.Config," and the FlashAttention layer internally handles the backend dispatch. This means that adding a new attention kernel for a new hardware backend requires modifying only the FlashAttention layer's dispatch logic, not every model implementation that uses attention — another $O(1)$ rather than $O(N)$ change.

Why this requires careful layer design. The paper notes that this drop-in replacement pattern only works because "AXLearn layers are already sharding and remat aware" (Section 4.2). If the default attention layer baked in assumptions about sharding (e.g., manually inserting all-reduce operations for tensor parallelism), a drop-in replacement would need to replicate those assumptions or the parallelism would break. Because AXLearn layers delegate sharding to GSPMD annotations applied externally, the replacement layer inherits the same sharding behavior automatically.

Quantization as layer replacement. The paper extends the same pattern to quantization:

"All components are implemented as strictly encapsulated modules. This allows expressing optimizations like quantization as a replacement of DotGeneral layers (XLA, 2025) with their quantization-aware equivalents."

DotGeneral is the XLA operation that represents general tensor contractions (matrix multiplications, convolutions, etc.). By replacing the standard DotGeneral layer with an INT8 or FP8 variant through a config modifier (as shown in the mesh rule example with INT8ConfigModifier and FP8ConfigModifier), quantization can be applied uniformly across the entire model without modifying any layer that uses matrix multiplication.

Ahead-of-Time (AOT) Compilation

Section 4.2 describes AOT compilation as a workflow that catches errors locally before distributed execution:

"AXLearn provides native support for JAX Ahead-of-Time (AOT) compilation, which allows users to analyze the memory and FLOPS utilization of a training program without executing a single line of the program, including catching errors like OOMs that would otherwise result in wasted resources."

The mechanism. JAX programs can be lowered to XLA's intermediate representation (HLO) without executing them, by calling jax.jit(fn).lower(inputs). The lowered representation includes explicit memory allocation sizes for every tensor, communication collectives injected by GSPMD, and the full computation graph. AXLearn exposes this through the composer: after applying all mesh rules, config modifiers, and rematerialization policies, the composer generates the JAX program and optionally runs AOT compilation, reporting the estimated memory usage per device and total FLOP count.

The reliability guarantee. The paper emphasizes a critical property:

"Because the same codepath is used for AOT and actual training, users can be confident that a program that AOT-compiles will run at a larger scale."

This is not trivially true in all systems. In PyTorch, for instance, the computation graph is built dynamically during execution, so there is no "lowered" representation to inspect ahead of time. TorchTitan and other PyTorch-based systems must actually allocate memory on GPUs to detect OOMs. In AXLearn, the AOT path shares exactly the same config resolution, mesh rule application, sharding annotation, and remat policy logic as the actual training path — the only difference is that the lowered HLO is inspected rather than compiled to machine code. This is what makes the confidence guarantee hold: there is no separate "analysis mode" with different code paths that could diverge from actual execution.

Practical impact. The paper notes in Section 7.4 that AOT compilation was critical for scaling development with limited TPU capacity: "a significant portion of resources were consumed by jobs with low resource utilization or no progress due to preventable errors... We capitalized on JAX's compiler-first approach by deeply integrating AOT-compilation, which allowed users to debug training entirely on CPU." This is an operational benefit that emerges from the compilation-based architecture (JAX/XLA) but is made practical by AXLearn's integration of AOT into the composer workflow.

InvocationContext: State Management Under Functional Purity

This is the most architecturally subtle component of AXLearn, solving a fundamental tension between JAX's functional programming requirements and the imperative programming style that deep learning practitioners expect (Section 4.3).

The problem. JAX transformations like jax.jit (Just-In-Time compilation) and jax.grad (automatic differentiation) require that functions be purely functional — they must have no side effects, return all outputs explicitly, and not depend on mutable external state. However, neural network training is inherently stateful: model parameters must be stored and updated, pseudo-random number generator (PRNG) keys must be advanced to ensure stochasticity, training summaries (loss, accuracy, gradients) must be collected, and intermediate outputs must be aggregated across the module hierarchy.

In a naive functional implementation, the caller would need to manually pass all state (parameters, PRNG key, output collections) into every module invocation and manually collect all state updates on return. This would mean that adding a new layer to a model requires updating the state-passing code in every intermediate module — exactly the $O(N)$ complexity that AXLearn seeks to avoid.

The solution: InvocationContext stack. AXLearn introduces a data structure called InvocationContext that behaves analogously to a traditional call stack. The mechanism is:

1. When a parent module invokes a child module, an InvocationContext for the child is automatically pushed onto a stack. This context:

   • Retrieves the child's state (parameters) from the parent's state store.
   • Splits the parent's PRNG key to produce a fresh, deterministic child PRNG key (using JAX's functional PRNG splitting).
   • Creates a new, empty data store for the child's outputs (training summaries, metrics, etc.).

2. The child module executes, reading its parameters and PRNG key from the context, and writing any summaries or outputs to the context's output store.

3. When the child module returns, the context is popped from the stack:

   • The child's parameter updates (gradients) are collected into the parent's state store.
   • The child's output store is merged into the parent's output store.
   • The child's PRNG key is discarded (the parent retains its own key, which was split, not consumed).

This mechanism is depicted in Figure 3 of the paper.

Why this preserves encapsulation. The crucial design decision is that "InvocationContexts contain references to modules, but not vice-versa." Modules do not hold references to their own context, their parent's context, or any sibling contexts. Instead, the context stack is managed externally by the AXLearn runtime, and modules access their context through a global (thread-local) mechanism. This means:

• A module can access shared state (like tied weights) by traversing the context stack programmatically — looking for a context of a specific type — without knowing where in the module hierarchy that state lives. For example, a language model head that shares weights with the embedding layer can find the embedding layer's parameters by searching the context stack for an Embedding context, without the head module having a reference to the embedding module.

• A module never needs to pass state to its children explicitly. The child's InvocationContext handles state retrieval automatically, so the parent's __call__ method can simply invoke self.child(inputs) without manually unpacking and repacking parameters.

• The same module can be used in different positions in the hierarchy without modification, because it always retrieves its state from its current context rather than from a fixed parent reference.

Compatibility with third-party code. The paper highlights a subtle benefit: because the InvocationContext stack can be accessed from arbitrary function calls (not just module methods), it allows "deep integration with 3rd party libraries that are not natively aware of the AXLearn state system (e.g., optax optimizers)" and "compatibility with codepaths with unique execution behavior, such as JAX's custom vjp backward pass." In JAX's custom_vjp (custom vector-Jacobian product, used for defining custom backward passes), the function that computes the backward pass may not have access to the original module hierarchy. But because it can still access the InvocationContext stack, it can retrieve state (like saved activations for gradient computation) without the module architecture being explicitly threaded through the backward pass.

Interaction with JAX transformations. The InvocationContext is designed to be compatible with JAX's functional transformations. When jax.jit compiles a function, the context stack operations (push, pop, state retrieval, PRNG splitting) are traced as part of the computation graph and compiled into the XLA program. This means there is no runtime overhead for state management in the compiled execution — all context operations are resolved at compile time.

Comparison to prior approaches. The paper notes that earlier JAX-based systems (early Flax, for example) required users to manually thread state through the module hierarchy: "one manually specified parameters as inputs to each module, and bubbled up outputs through the call stack. This hurt modularity as it required implementations to be intricately aware of state structure, and hurt usability as users lost the comfort of the PyTorch imperative style" (Section 7.4). The InvocationContext was designed specifically to address this — it restores imperative-style programming (where state is accessed implicitly) while maintaining the functional purity that JAX requires.

Summary of Design Choices and Their Justifications

• Hierarchical configs over flat configs: Partially-specified configs with top-down resolution allow entire subtrees to be reused without knowledge of internal details. Flat configs force every module to be modified when new features are added because all parameters live in the same namespace.

• Config traversal over subclassing: Applying changes by visiting the config tree and replacing types (composition) rather than by subclassing modules and overriding methods (subtyping) decouples the transformation from the module hierarchy structure.

• Mesh rules over per-model hardware configs: Mapping hardware types to lists of config modifiers centralizes hardware-specific optimizations, avoiding the need for every model implementation to contain backend-conditional logic.

• GSPMD over manual communication: Compiler-injected communication separates parallelism specification from layer implementation, allowing the same layer code to work with FSDP, tensor parallelism, or data parallelism.

• Tagged remat points over hard-coded remat policies: Exposing named checkpoints in the module hierarchy that external policies target allows per-hardware rematerialization strategies without modifying layer code.

• Drop-in kernel replacement over backend-conditional dispatch: Encapsulating hardware-specific kernels behind a common layer interface with internal dispatch avoids backend conditionals in model code.

• AOT compilation sharing the exact training codepath: Guarantees that errors caught locally will also manifest at scale, avoiding false confidence from a separate analysis mode.

• InvocationContext stack over manual state threading: Preserves encapsulation by decoupling state management from the module hierarchy, allowing modules to be written in imperative style while satisfying JAX's functional constraints.

4. Key Insights and Innovations

Innovation 1: LoC-Complexity as a Formal Framework for Measuring System Modularity

The most intellectually distinctive contribution of this paper is not AXLearn itself but the LoC-complexity framework — a way to formally quantify the modularity of a deep learning training system by measuring how the integration cost of new features scales with system size. This is a diagnostic concept that did not exist before in the ML systems literature. Prior work discussed modularity in qualitative terms: systems were described as "flexible" or "extensible" based on design principles or rules of thumb. TorchTitan (Liang et al., 2024) argued that writing configs is easier than implementing subtyped modules, and Praxis used a composable config system, but neither provided a metric for evaluating how modular their designs actually were.

The paper's insight is that modularity can be measured asymptotically: given a system with $N$ modules and a feature with $M$ variants, what is the asymptotic lines-of-code change required to existing interfaces to integrate the feature? The answer — $O(1)$ for AXLearn versus $O(N)$ or $O(NM)$ for every other system surveyed (Table 2) — transforms modularity from a subjective design aesthetic into a rigorous, quantifiable property. This is analogous to how algorithmic complexity theory moved the analysis of algorithms from "this feels fast" to "this is $O(n \log n)$." The framework doesn't just describe AXLearn — it provides a lens for evaluating any future training system's design.

The framework is validated concretely: the paper traces how RoPE integration requires propagating positional embedding parameters through Megatron-LM's GPTModel → TransformerBlock → TransformerLayer → Attention hierarchy, how DeepSpeed's config flattening forces every model implementation to override positional_embedding_type, and how TorchTitan's per-model Attention implementations condition on RoPE configs with $O(NM)$ branching logic. These are not hypothetical failures — they are demonstrated through code-level analysis in Appendix B with specific LoC estimates (400 for Megatron-LM, 320 for DeepSpeed, 240 for TorchTitan, 0 for AXLearn). The framework's explanatory power is demonstrated by its ability to predict why certain architectures (like Praxis's template approach) achieve $O(M)$ for MoE but degrade to $O(NM)$ for RoPE — because config flattening leaks across the encapsulation boundary for positional embeddings but not for expert routing.

The significance of this contribution extends beyond the paper itself. The LoC-complexity framework provides a vocabulary for the field to discuss system design tradeoffs that were previously ineffable. When a new training system is proposed, the question is now not just "does it support MoE?" but "what is its LoC-complexity for MoE integration?" — a more precise and actionable criterion. This is a fundamental contribution, not an incremental refinement, because it introduces an entirely new axis of evaluation into the ML systems design space.

Innovation 2: The Composition-over-Subtyping Principle Applied Systematically to Deep Learning

The paper's second conceptual contribution is the observation that subtyping is the root cause of $O(N)$ complexity scaling in deep learning systems, and that systematic composition — applied not just to layers but to configs, parallelism, memory optimization, and state management — can recover constant complexity. This is not a new idea in software engineering: the "composition over inheritance" principle has been a staple of object-oriented design since the Gang of Four (Gamma et al., 1994). What is new is the paper's diagnosis that every existing large-model training system violates this principle systematically, and its demonstration that the violation can be eliminated entirely.

Prior work conflated modularity with configurability. DeepSpeed and TorchTitan provide configurable systems — you can set num_experts in a config file — but they achieve configurability through config flattening, which still requires every consumer of the config to be modified when new fields are added. Megatron-LM achieves modularity through subtyping (TransformerBlockSubmodules, ColumnParallelLinear), but subtyping couples the parent's implementation to the child's type, meaning replacing a child forces the parent to change. The paper's key diagnostic move is to identify that these patterns look modular at small scale but degrade to $O(N)$ complexity as the number of modules in the system grows. The "induction" argument in Section 2.1 — replacing a child forces the parent to change, which forces the grandparent to change, ad infinitum — crystallizes why subtyping fails in large codebases even when it works in simple examples.

AXLearn's composition achieves $O(1)$ complexity through three mutually reinforcing design choices: (1) modules declare their children by interface (config type) rather than concrete type, so child substitution doesn't affect the parent; (2) config modifiers operate on the config tree from the outside, so features can be applied universally without module awareness; and (3) the InvocationContext stack decouples state management from module hierarchy, so modules can access shared state without parent references. None of these ideas is individually novel, but their systematic combination — applied across layers, parallelism, rematerialization, and kernels — is what distinguishes AXLearn. The paper demonstrates this through the dramatic gap between the 10-line config modifier that applies MoE to 1,000+ experiments in AXLearn and the 4,000 LoC estimate for DeepSpeed (Table 2).

This is a fundamental contribution because it identifies the architectural invariant that separates $O(1)$ from $O(N)$ systems, not a small refinement of existing approaches. The invariant is that no module implementation may depend on the concrete type or internal structure of any other module, and the paper shows that prior systems violate this invariant in multiple ways (subtyping, config flattening, manual state threading, parallelism-specific layer variants). Establishing this invariant, and showing that it can be maintained without sacrificing performance (Table 3), is the paper's central architectural insight.

Innovation 3: Hardware-Agnostic Training Through Config-Based Optimization Injection

The third conceptual contribution is the "mesh rule" paradigm for hardware-agnostic training: the idea that all hardware-specific optimizations — parallelism strategies, rematerialization policies, kernel selection, quantization — can be expressed as config modifiers applied to a shared layer graph, rather than as code branches within layer implementations. Prior work treated hardware-specific optimization as fundamentally entangled with model code. Megatron-LM's tensor parallelism is implemented through ColumnParallelLinear and RowParallelLinear subclasses — the parallelism strategy is literally part of the layer's type system. PyTorch FSDP wraps modules at the torch.distributed level, but switching to tensor parallelism requires changing the layer code. Even JAX-based systems like MaxText and Flax, which benefit from GSPMD's compiler-injected communication, still embed hardware-specific logic in their model implementations: MaxText's Attention conditions on attention_type and rope_type configs, and Flax's Gemma model propagates RoPE parameters through the module hierarchy explicitly.

AXLearn's insight is that hardware agnosticism requires not just a portable compiler (XLA) but a separation of concerns between model specification and hardware optimization that is architecturally enforced. The mesh rule mechanism achieves this: a mapping from hardware type patterns to lists of config modifiers (MeshShapeModifier, RematSpecModifier, INT8ConfigModifier, FP8ConfigModifier) is applied to the config tree externally by the composer. The model code is completely unaware of whether it will be executed with FSDP within TPU v5e slices or 8-way tensor parallelism within H100 nodes — it only sees the resolved config after modifiers have been applied. This is fundamentally different from prior approaches, where the model code contains the hardware-specific logic (through conditional instantiation, parameter propagation, or subclass selection).

The paper validates this separation through a striking operational result: AWS Trainium2 "did not exist when we first began development," yet AXLearn was "one of the first deep learning systems that supports Trainium2 at scale" (Section 2.2) — without modifying any model code. This is possible because supporting a new hardware backend requires only adding a new mesh rule entry (specifying the appropriate mesh shape, remat policy, and kernel selection for Trainium2) and implementing backend-specific kernels behind AXLearn's existing layer interfaces. The $O(1)$ complexity of this process — one new mesh rule entry, zero model code changes — is a direct consequence of the architectural separation, and it demonstrates a level of hardware agility that no prior system achieved.

This contribution is conceptually fundamental because it redefines what "hardware support" means in a training system. In prior systems, supporting a new hardware backend meant the system developers implementing backend-specific optimizations in model code. In AXLearn, supporting a new backend means the hardware vendor or user implementing backend-specific kernels and writing a mesh rule — the training system itself is agnostic. This shifts the burden of hardware portability from the framework maintainer to the ecosystem, which is precisely the architectural property that has made Linux portable across architectures while maintaining a unified kernel interface. The paper is the first to articulate and demonstrate this property for large-scale deep learning training.

Innovation 4: State Encapsulation Enables $O(1)$ Complexity Under Functional Constraints

The fourth contribution is the InvocationContext abstraction — not as a mechanism (which Section 3.4 covers in detail) but as a diagnostic solution to a specific scaling pathology: manual state threading through functional module hierarchies produces $O(N)$ complexity when modules are added or removed. The paper identifies this as a previously unarticulated failure mode in functional deep learning systems.

Prior JAX-based systems (early Flax, Haiku to some extent) required users to manually pass parameters, PRNG keys, and output collections through the module call stack. This meant that adding a new layer to a model required updating the state-passing signature of every intermediate module between the new layer and the root — a direct violation of the modularity that JAX's functional purity was supposed to enable. The paper describes this operational pain in Section 7.4: "one manually specified parameters as inputs to each module, and bubbled up outputs through the call stack. This hurt modularity as it required implementations to be intricately aware of state structure." This $O(N)$ complexity arises from a subtle interaction between JAX's functional constraints and the hierarchical structure of neural networks, and prior work neither identified it as a systematic problem nor provided a general solution.

The InvocationContext solves this by inverting the responsibility for state management: instead of the module hierarchy managing state (which couples module implementations to the hierarchy structure), an external context stack manages state on behalf of the modules (which decouples implementations from hierarchy). The paper's design choice that "InvocationContexts contain references to modules, but not vice-versa" is the critical invariant: modules are unaware of their position in the hierarchy or the identities of their siblings, and all state access flows through the context stack rather than through instance attribute traversal. This means modules can be added, removed, or rearranged without any changes to other modules' state management code — the context stack adjusts automatically.

The paper provides a concrete illustration of why this matters beyond developer convenience: tied weights. In subtyping-based systems, implementing weight tying between an embedding layer and a language model head requires the head to reach through the module hierarchy to access the embedding's parameters (e.g., self.parent.embedding.weight). This creates a coupling between the head's implementation and the embedding's position in the hierarchy. In AXLearn, the head traverses the InvocationContext stack to find the embedding context programmatically, without knowing where the embedding lives. If the model architecture is restructured — say, the embedding is moved to a different parent — the head's weight-tying logic continues to work because it searches the context stack by type, not by position. This is a subtle property that directly emerges from the state encapsulation design, and it has no equivalent in subtyping-based systems.

This is a fundamental contribution to the design of functional ML systems because it identifies and solves a scaling pathology that was previously invisible. The InvocationContext is not merely a convenience — it is what prevents state management from being the bottleneck that makes functional systems less modular than imperative ones, which would be an ironic failure mode given that functional purity is theoretically more composable. The paper demonstrates that with proper state encapsulation, functional systems can achieve the same $O(1)$ integration complexity as their imperative counterparts, while retaining the compiler optimizability that makes hardware agnosticism possible.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper evaluates training performance across three model architectures: Llama2-7B, Llama2-70B (dense models), and Qwen-3 30B-A3B (a Mixture of Experts model). For modularity analysis, the evaluation is performed on hypothetical production codebases with 20 model variants and 10 attention variants, rather than on a specific dataset, since the metric is lines-of-code changes to existing interfaces. For inference benchmarking, the ShareGPT dataset is used with maximum input lengths of 1024 (7B) or 1,800 (70B) and maximum output length of 256 tokens.

• Base model(s). Training performance comparisons use Llama2-7B (with 4,096 context length), Llama2-70B (same context length), and Qwen-3 30B-A3B. The scaling study in Figure 4 uses two proprietary models: "Model A" (70B parameters, 4,096 context length) and "Model B" (150B parameters, 8,192 context length). The paper does not specify the architecture of these models beyond parameter count, noting they are production models from Apple.

• Metrics. The paper uses three distinct metric categories for its three evaluation axes:

  • LoC-Complexity (modularity): The asymptotic lines-of-code changes required to existing interfaces to integrate a new feature (RoPE or MoE), measured across existing modules in the system, not counting the implementation of the feature itself. This is formalized in Section 7.1 and quantified in Table 2 with concrete LoC estimates under a realistic production setting of 20 model variants and 10 attention variants.
  • Training throughput and MFU (performance): Iteration time (seconds per training step), Model FLOPs Utilization (MFU, the ratio of achieved FLOPs to theoretical peak FLOPs of the hardware), and throughput in tokens/second. All performance comparisons use a global batch size of 1024.
  • Inference latency and throughput: Time-To-First-Token (TTFT) in milliseconds, Time-Per-Output-Token (TPOT) in milliseconds, and overall throughput in tokens/second.

• Baselines. The paper compares against five training systems:

  • PyTorch FSDP (Zhao et al., 2023) or PyTorch XLA FSDP (PyTorch, 2025) on TPU (since standard FSDP does not run on TPU).
  • Megatron-LM (Narayanan et al., 2021) — GPU-only baseline.
  • MaxText (MaxText team, 2025) — runs on GPU and TPU.
  • DeepSpeed (DeepSpeed, 2025a,b), TorchTitan (Liang et al., 2024), Flax (Heek et al., 2023), Praxis (PAX team, 2022), and Haiku (Haiku, 2025) for modularity comparisons only (LoC analysis in Table 2).
  • vLLM (Kwon et al., 2023) for inference benchmarking on TPUs. None of the baselines support AWS Trainium2, so AXLearn is evaluated alone on that platform.

• Generation budget / compute accounting. Training performance is compared on identical hardware configurations: 256/512 H100 GPUs (32/64 AWS P5d instances, 8 GPUs each), 512 B200 GPUs (64 AWS P6 instances, 8 GPUs each), TPU v5p-512/1024 (64/128 GCP Cloud TPU hosts, 4 chips each), and 1,024 Trainium2 chips (64 AWS trn2 instances, 16 chips each). All runs use a global batch size of 1024. Systems within each hardware group in Table 3 are benchmarked on the exact same cluster with the same number of accelerators. The scaling study fixes per-device batch size. For LoC analysis, the budget is the number of existing modules that must be modified, measured asymptotically and concretely with estimates under a standardized production codebase assumption.

• Cross-validation / statistical protocol. The paper does not report statistical protocols such as multiple runs, confidence intervals, or cross-validation, with one exception: the inference benchmarks include specified parameters (maximum input/output lengths, VM instance types). The LoC estimates use conservative assumptions (20 model variants, 10 attention variants) but are derived from manual code inspection of public repositories, not from automated analysis. The paper does not report variance or statistical significance for throughput numbers, MFU values, or recovery latency measurements. For inference, the vLLM comparison is described as a single benchmark run per model per system.

Main Quantitative Results

The paper's evaluation spans three largely independent axes: modularity (LoC-complexity analysis), training performance (throughput and MFU across hardware), and inference performance. The scaling study and failure recovery measurements provide additional operational evidence.

Modularity: LoC-Complexity Analysis

The central claim about modularity is that AXLearn achieves $O(1)$ LoC-complexity while all other systems exhibit $O(N)$ or $O(NM)$ complexity. Table 2 presents the formal analysis.

RoPE integration complexity. The paper estimates that integrating a single RoPE variant into a production codebase with 20 model variants and 10 attention variants requires:

• AXLearn: 0 LoC changes to existing interfaces. The paper states "in AXLearn, 0 LoC changes to existing interfaces are necessary" (Table 2, LoC Estimate column). The same 10-line config modifier configures RoPE across over 1,000 experiments internally.
• Megatron-LM: 400 LoC (Table 2). Appendix B details the mechanism: RoPE parameters (rotary_percent, rotary_base, rotary_scaling, position_embedding_type) are flattened into each model's init signature and propagated through GPTModel → TransformerBlock → TransformerLayer → Attention. Each of the 20 model variants "conservatively" incurs at least 20 LoC changes.
• DeepSpeed: 320 LoC (Table 2). DeepSpeed's monolithic DeepSpeedInferenceConfig groups RoPE configs, and each model overrides positional_embedding_type() and positional_embedding_config(). With 10 attention variants (e.g., DSDenseBlockedAttention) each requiring ~20 LoC to handle the embedding type in init and forward, plus 120 LoC for 20 models at ~6 LoC each, the total reaches 320.
• TorchTitan: 240 LoC (Table 2). Each model's ModelArgs subclass must include RoPE fields, and each model's own Attention implementation conditions on RoPE configs. With 20 model variants averaging 2 LoC for config changes and 10 LoC for attention modifications, plus 10 attention variants, the total is 240.
• Flax: 600 LoC (Table 2). The Gemma model propagates RoPE parameters from TransformerConfig → Transformer → Block → Attention, requiring at least 30 LoC per model variant.
• Praxis: 300 LoC (Table 2). Praxis flattens RoPE configs like use_rotary_position_emb into each attention layer. With 10 attention variants (e.g., DotProductAttention, MultiQueryDotProductAttention) each requiring ~30 LoC, total is 300.
• MaxText: 200 LoC (Table 2). MaxText's Attention conditions on configs like attention_type and rope_type to choose the RoPE module, creating cross-product interactions between attention variants and RoPE variants. With ~10 LoC per variant, total is 200.

MoE integration complexity. The paper estimates:

• AXLearn: 0 LoC changes to existing interfaces. The same 10-line replace_config snippet applies MoE to over 1,000 experiments.
• DeepSpeed: 4,000 LoC (Table 2). This is the most dramatic finding. DeepSpeed requires subclassing each model from DSMoETransformerModelBase, which "requires in some cases a re-implementation of most methods." The QwenV2MoE implementation requires more than 200 LoC; across 20 model variants, this conservatively reaches 4,000 LoC.
• TorchTitan: 400 LoC (Table 2). Each model's TransformerBlock or DecoderLayer conditionally instantiates MoE or FFN based on moe_enabled. With 20 model variants, each requiring ~10 LoC for ModelArgs changes and ~10 LoC for layer changes, total is 400.
• Megatron-LM: 20 LoC (Table 2, but note this is the best non-AXLearn result for MoE). Megatron-LM uses composition via TransformerBlockSubmodules to specify MoE in place of MLP, but introduces an is_expert field that propagates through linear submodules. With 10 MLP variants each incurring 1 LoC for is_expert, plus 10 modules using linear submodules each incurring 1 LoC, total is 20.
• MaxText: 300 LoC (Table 2). MoE details are flattened into each model's DecoderLayer. With 10 LoC per decoder across 20 model variants (200 LoC), plus 5 LoC for each loss function that uses MoE configs for auxiliary losses across 20 variants (100 LoC), total is 300.
• Praxis: 5 LoC (Table 2). Praxis uses a "template" approach that achieves $O(M)$ complexity: each MoE variant incurs a 5 LoC change for flattened configs, but does not require per-model modifications.

The paper also reports the asymptotic analysis underlying these concrete estimates in Table 2:

• AXLearn: $O(1)$ for both RoPE and MoE. This is the headline result.
• Megatron-LM: $O(NM)$ for RoPE (parameters must propagate through the hierarchy for each variant), $O(N)$ for MoE (the is_expert field propagates through linear submodules).
• DeepSpeed: $O(NM)$ for both. Config flattening forces per-model overrides, and attention layers branch on embedding type.
• TorchTitan: $O(NM)$ for both. Per-model config classes and per-model attention implementations produce quadratic interactions.
• Flax: $O(NM)$ for RoPE (parameters propagate through manually threaded module signatures). No MoE implementation evaluated.
• Praxis: $O(NM)$ for RoPE (flattened configs interact with per-attention-layer implementations), $O(M)$ for MoE (template approach avoids per-model changes).
• MaxText: $O(NM)$ for both. Config-based branching in Attention and Decoder layers produces quadratic interactions.

The critical methodological detail is that these LoC estimates count only changes to existing interfaces, not the implementation of RoPE or MoE itself. The paper argues that "the same feature across systems tends to incur similar LoC up to some constant" (Section 7.1), making the interface-change count the differentiating factor.

Training Performance Across Hardware

Table 3 summarizes training throughput, iteration time, and MFU for Llama2-7B, Llama2-70B, and Qwen-3 30B-A3B across four hardware backends.

Llama2-7B results:

• On 32× H100-8 (256 GPUs): AXLearn achieves 1.4s iteration time, 54.2% MFU, and 3.0M tokens/s, matching MaxText (1.4s, 54.7%, 3.0M). Both outperform PyTorch FSDP (2.6s, 29.9%, 1.6M) and roughly match Megatron-LM (1.7s, 44.9%, 2.5M), though Megatron-LM has slightly lower throughput. The paper notes Megatron-LM has "stronger performance on H100 GPUs... because PyTorch currently has finer-grained scheduling capability over XLA," but frames this as a trade-off AXLearn accepts for hardware agnosticism.

• On TPU v5p-512: AXLearn achieves 2.5s iteration time, 66.2% MFU, and 1.7M tokens/s, outperforming MaxText (2.7s, 61.6%, 1.6M) and PyTorch XLA FSDP (3.5s, 46.7%, 1.2M). The paper attributes MaxText's lag "likely due to its choices of rematerialization." PyTorch XLA FSDP notably does not fail with OOM on 7B at this scale.

• On 64× Trainium2-16 (1,024 chips): AXLearn is the only system benchmarked, achieving 1.2s iteration time, 24.2% MFU, and 3.5M tokens/s. The paper notes that none of the baselines support Trainium.

Llama2-70B results:

• On 64× H100-8 (512 GPUs): AXLearn achieves 9.2s iteration time, 40.0% MFU, and 456K tokens/s, compared to PyTorch FSDP (10.6s, 34.7%, 396K), Megatron-LM (7.8s, 47.2%, 538K), and MaxText (9.4s, 39.1%, 446K). The pattern is consistent with 7B: AXLearn roughly matches MaxText, outperforms PyTorch FSDP, and trails Megatron-LM slightly.

• On TPU v5p-1024: AXLearn achieves 11.6s iteration time, 68.0% MFU, and 360K tokens/s, outperforming MaxText (12.3s, 64.4%, 341K). Critically, PyTorch XLA FSDP fails entirely with out-of-memory errors ("OOM" in Table 3), which the paper highlights as evidence that the PyTorch XLA FSDP pathway is not production-ready for large models on TPUs.

• On 64× Trainium2-16: AXLearn achieves 11.2s iteration time, 25.0% MFU, and 374K tokens/s. No baseline comparison available.

Qwen-3 30B-A3B (MoE) results:

• On TPU v5p-1024: AXLearn achieves 12.86s iteration time, 31.58% MFU, and 1.3M tokens/s, essentially matching MaxText (12.97s, 31.31%, 1.3M). The paper states these are "within measurement noise."

• On 64× B200-8 (512 GPUs): AXLearn achieves 4.31s iteration time, 19.22% MFU, and 3.9M tokens/s, compared to Megatron-LM (4.10s, 20.20%, 4.1M). Megatron-LM again shows a slight advantage on GPU.

The paper calculates the B200 MFU based on dense FLOPs, which explains the lower absolute MFU for the MoE model, but provides no further details on this computation. A notable pattern across all results: AXLearn's TPU MFU (66.2% for 7B, 68.0% for 70B) consistently exceeds its GPU MFU (54.2% for 7B, 40.0% for 70B), which the paper attributes implicitly to XLA compiler maturity on TPU versus GPU, though this is not stated explicitly.

Scaling Study

Figure 4 presents a weak-scaling study of two production models on TPUs, fixing per-device batch size.

• Model A (70B, 4,096 context length): Scaling from 256 to 4,096 TPU chips, iteration time remains nearly constant (approximately 8–9 seconds across the range, based on the y-axis scale showing a roughly flat line near the top of the plot), while MFU drops from 63.0% to 52.4%. The paper states these numbers "demonstrate that AXLearn achieves very close to linear scaling."

• Model B (150B, 8,192 context length): Scaling from 8,192 to 32,768 chips, MFU drops from 40.6% to 37.6%. The paper notes that MFU for the 150B model is lower "due to the need to limit the global batch size at 32,768 chip scale for good training convergence," and that the 150B experiments have 1/16 the per-chip sequence length compared to 70B experiments — a significant detail that explains the MFU discrepancy.

The absolute iteration time scale in Figure 4(a) appears to be ~8–10 seconds for Model A and ~1–2 seconds for Model B in Figure 4(b), though the paper does not quote these exact numbers in the text. Both figures show near-flat lines for iteration time and gradually declining MFU curves, consistent with the paper's "close to linear scaling" claim.

Inference Performance

Table 4 compares AXLearn inference performance against vLLM on TPUs for Llama2-7B and Llama2-70B using the ShareGPT dataset.

Llama2-7B (TPU v5p-8, max input 1024, max output 256):

• AXLearn achieves 40.1ms TTFT vs. vLLM's 538.6ms — a 13.4× speedup, not 500× as the paper claims. (The paper's "500x" likely refers to the 70B TTFT, not 7B, but the sentence structure is ambiguous.) The paper states "AXLearn achieves 500x and 6x speedup in TTFT and TPOT over vLLM, respectively," but the 7B numbers show 13.4× for TTFT and 2.5× for TPOT (9.1ms vs. 22.4ms). Throughput is 3,125 tokens/s vs. 1,117 tokens/s — a 2.8× improvement.

Llama2-70B (TPU v6e-8, max input 1,800, max output 256):

• AXLearn achieves 150.5ms TTFT vs. vLLM's 80 seconds — this is the ~530× speedup the paper highlights. TPOT is 28.1ms vs. 189.8ms (6.8×). Throughput is 1,139 tokens/s vs. 705 tokens/s (1.6×).

The paper explicitly caveats these results with two important notes. First, "TPU support for vLLM is still experimental, which likely contributes to the performance gap." Second, vLLM does not support the 70B model on v5p-8 at the time of benchmarking, requiring a different VM instance (v6e-8) for that comparison. The paper states its goal is "not to position AXLearn as a specialized inference engine, but to demonstrate that a modular training framework can achieve production-grade inference performance with minimal additional effort."

Failure Recovery Latency

Figure 5 shows a single recovery event during a production training job on 32,768 TPUs over a 1-hour window.

• Asynchronous checkpointing: The paper notes that "no training throughput reduction is observed during checkpoint creation" — the throughput line in Figure 5 shows a brief spike labeled "Checkpoint created" without visible throughput degradation.

• Failure detection: The paper states AXLearn "detects the training job has failed immediately" upon hardware failure and "initiates slice-level hot-swap."

• Hot-swap time: 4 minutes from failure detection to job recovery start.

• Checkpoint restoration: 9 minutes after hot-swap completes.

• Total training time lost: 21 minutes, which the paper explains includes "both the downtime during hot-swap and checkpoint restoration, as well as the lost progress from training steps completed after the most recent checkpoint before the failure occurred."

Figure 5 shows throughput dropping to zero at "Failure detected," remaining at zero through "Job recovery started" and "Checkpoint restored," and then returning to approximately the same throughput level (~4 steps/min) as before the failure. The absolute throughput before and after appears constant, suggesting no performance degradation from the recovery process.

Ablation Studies and Robustness Checks

The paper does not conduct traditional ablation studies — there are no experiments removing specific components of AXLearn and measuring the impact, no comparison of alternative design choices within AXLearn (e.g., hierarchical vs. flat configs holding other factors constant), and no sensitivity analysis of the LoC estimates to different codebase assumptions (e.g., what if there are 100 model variants instead of 20?).

Instead, the paper provides several forms of robustness evidence that serve a similar function:

Hardware diversity as a robustness check for the hardware-agnostic claim. AXLearn is evaluated on four distinct hardware backends (H100 GPU, B200 GPU, TPU v5p, AWS Trainium2) and achieves competitive performance on all of them. The fact that the same codebase produces 66.2% MFU on TPU v5p-512, 54.2% MFU on H100, and 24.2% MFU on Trainium2 (Table 3) without model code changes is itself the primary evidence for hardware agnosticism. The Trainium2 result is particularly informative because Trainium2 "did not exist when we first began development" (Section 2.2) — this is an unplanned robustness check demonstrating that the mesh rule architecture works for hardware that was not anticipated during initial design.

Multi-model evaluation as a robustness check for model-agnosticism. The paper benchmarks three architecturally distinct model families: Llama2-7B (dense, 7B parameters), Llama2-70B (dense, 70B parameters), and Qwen-3 30B-A3B (MoE, 30B total with 3B active parameters). AXLearn achieves competitive performance across all three at different scales (Table 3), providing evidence that the composable architecture works for both dense and MoE architectures. However, all three models are decoder-only Transformers, so the paper does not demonstrate agnosticism to fundamentally different architectures (e.g., encoder-decoder models like T5, mixture-of-experts with different routing mechanisms beyond what Qwen uses, or non-Transformer architectures like Mamba).

Weak scaling behavior as a robustness check for the scalability claim. Figure 4 demonstrates near-constant iteration time as chip count increases from 256 to 4,096 (70B model) and 8,192 to 32,768 (150B model), with MFU declining modestly (63.0% → 52.4% for the 70B, 40.6% → 37.6% for the 150B). The paper does not report the exact iteration times at each scale (only the MFU percentages), but the flat lines in Figure 4 visually support the linear scaling claim. The lower MFU for the 150B model is attributed to batch size constraints, not to a fundamental scaling pathology.

Inference performance as a robustness check for modular design generality. The paper describes the inference capability as "one surprising discovery" (Section 6), positioning it as evidence that the modular design's benefits extend beyond the intended training domain. The 2.8× throughput improvement over vLLM on TPUs for 7B (Table 4) demonstrates that encapsulation does not inherently limit performance in deployment scenarios. However, the vLLM comparison is weakened by the experimental status of vLLM's TPU support, and there is no comparison against optimized inference engines on GPU (e.g., TensorRT-LLM, vLLM on GPU), which would be a stronger test of generality.

Config modifier scalability. The paper states that the same 10-line config modifier snippet is "used to apply MoE to over 1,000 experiment configs, without any additional changes to other modules" (Section 4.1). This is a form of scale validation for the $O(1)$ claim: if config modifiers required changes when applied to many different configs, the effective complexity would not be constant. The paper provides this as an anecdotal statistic rather than a systematic measurement.

Recovery latency as a robustness check for production readiness. The single-event recovery measurement in Figure 5 demonstrates that AXLearn's fault tolerance mechanisms (watchdog, slice-level hot-swap, asynchronous checkpointing) function in a production environment at 32,768-chip scale. However, the paper reports only one recovery event, so there is no data on the distribution of recovery times, the frequency of failures requiring recovery, or the success rate of the hot-swap mechanism.

Negative results in cross-system comparison. The paper reports several negative results for competing systems that indirectly strengthen AXLearn's claims by demonstrating that alternative approaches fail at scale:

• PyTorch XLA FSDP fails with OOM on Llama2-70B at TPU v5p-1024 (Table 3: "OOM" entry). This is evidence that supporting large models on TPU requires more than just XLA compilation — the system-level memory optimizations (rematerialization configuration, optimizer offloading) that AXLearn provides through config modifiers are necessary.
• DeepSpeed requires 4,000 LoC for MoE integration (Table 2). This is the most dramatic negative result for a competing system's modularity, though the paper does not provide a detailed breakdown of the 4,000 LoC estimate beyond the heuristic described in Appendix B.

What is notably absent. The paper does not provide several ablations that would strengthen the claims:

• No comparison of AXLearn with itself under different configuration strategies. For example, what is the performance impact of using config traversal vs. manually specifying each replacement? Does the InvocationContext add measurable overhead compared to manual state threading?
• No sensitivity analysis of LoC estimates. The 20-model, 10-attention-variant assumption in Appendix B is not varied. What if a codebase has 100 model variants? Does the LoC complexity scaling hold empirically, or was the asymptotic analysis derived from a small number of public examples and extrapolated?
• No comparison of a single model implemented modularly vs. monolithically within AXLearn. This would isolate the modularity overhead from the JAX/XLA baseline performance.
• No data on the failure rate of AOT compilation predictions. The paper claims AOT compilation catches errors that would otherwise waste resources (Section 4.2, 7.4), but does not report how often AOT predictions match actual execution, how many OOMs are caught by AOT vs. at runtime, or what the false-positive rate is.
• No latency-to-solution comparisons. All training comparisons use iteration time and throughput. For practitioners making hardware decisions, time-to-convergence (e.g., hours to reach a target validation loss) would be more directly informative than per-step metrics, but is not reported.

Critical Assessment

The paper's experimental design reflects its nature as a systems paper rather than a machine learning research paper. The evaluation does not follow the hypothesis-driven experimental paradigm of ML papers (controlled ablation, statistical significance, generalization across datasets). Instead, it follows the systems paper paradigm: demonstrate that the system (1) solves a real problem better than existing alternatives, (2) does not sacrifice performance to achieve its design goals, and (3) functions at production scale. The evidence presented largely supports these goals, but with important caveats.

Claim 1: AXLearn achieves constant $O(1)$ complexity for feature integration, demonstrated by 0 LoC changes to existing interfaces versus hundreds to thousands in other systems.

This claim is supported by the LoC analysis in Table 2 and Appendix B, but the evidence has significant methodological limitations. The LoC estimates are based on manual inspection of public repositories and the paper's stated assumptions about a "realistic production setting" (20 model variants, 10 attention variants). The derivation in Appendix B walks through the logic for each system, but the actual LoC counts are estimates, not measurements from a systematic instrumentation of the codebases. For example, the 4,000 LoC estimate for DeepSpeed MoE integration is derived from the observation that DeepSpeed's QwenV2MoE requires "more than 200 LoC" and "conservatively" multiplying by 20 model variants — but this assumes all 20 model variants would require comparable changes, which may not be the case if some variants share infrastructure. The 200 LoC for QwenV2MoE is itself presented as an example without a detailed line-by-line accounting.

More fundamentally, the paper evaluates only two features (RoPE and MoE) and only the integration of single variants (not the full $M$-variant case that the asymptotic analysis predicts). For the $O(NM)$ claim to be validated, one would need to measure LoC changes when adding multiple variants of each feature and observe quadratic scaling. The paper does not do this — Table 2 reports only single-variant LoC estimates and then states the asymptotic complexity separately. The gap between the asymptotic claim and the empirical evidence is bridged by reasoning in Appendix B (tracing how parameters propagate through module hierarchies), not by measurement. This is acceptable for a systems paper making an architectural argument, but the complexity claims should be understood as qualitative characterizations of design patterns rather than empirically validated scaling laws.

The fact that the paper can provide a 10-line snippet that configures MoE across 1,000 experiments is a genuine demonstration of the design's practical benefits, but it is a point demonstration rather than a systematic evaluation. Would a 100th feature also require 10 lines, or does the approach break down for features that don't fit the "drop-in replacement" pattern? The paper does not explore the boundaries of what kinds of features can be expressed as config modifiers.

Claim 2: AXLearn matches or exceeds the training performance of hardware-optimized alternatives.

Table 3 supports this claim with qualifications. On TPUs, AXLearn consistently achieves the highest MFU (66.2% vs. 61.6% for MaxText on 7B, 68.0% vs. 64.4% on 70B). On H100 GPUs, AXLearn roughly matches MaxText but trails Megatron-LM (54.2% vs. 44.9% MFU for 7B, but note AXLearn's higher throughput: 3.0M vs. 2.5M tokens/s). The paper acknowledges the Megatron-LM performance gap honestly, attributing it to PyTorch's finer-grained scheduling. On Trainium2, AXLearn has no competitors to compare against, so the 24.2% MFU for 7B cannot be assessed as competitive or subpar — it is simply the number achieved.

A critical limitation is that all comparisons are on a single global batch size (1024). Different systems may have different optimal batch sizes, and the optimal batch size may vary with model scale, hardware, and parallelism strategy. The paper does not report whether 1024 is the optimal batch size for any system or model, or whether results are sensitive to this choice. The fact that the 150B model in the scaling study uses 1/16 the per-chip sequence length of the 70B model "for good training convergence" suggests batch size tuning matters, but this is not explored systematically.

The Llama2 and Qwen models are well-known architectures with accessible reference implementations, making the performance comparisons relatively fair — each system's implementers have had the opportunity to optimize for these architectures. However, this also means the benchmarks measure maturity of optimization for specific architectures as much as they measure system quality. A newly released system like AXLearn achieving parity with Megatron-LM (developed over many years with specific focus on GPU optimization) is impressive, but the comparison favors AXLearn on TPU (where Megatron-LM does not run) and Trainium (where no competitor runs).

Claim 3: AXLearn scales to very large model sizes and chip counts.

Figure 4 demonstrates weak scaling from 256 to 4,096 chips for a 70B model and 8,192 to 32,768 chips for a 150B model with near-constant iteration time. The MFU decline (63.0% to 52.4% for 70B, 40.6% to 37.6% for 150B) is modest. However, "near-linear scaling" is an informal claim — the paper does not report the scaling efficiency (the ratio of achieved throughput to ideal linear scaling) as a formal metric. The visual evidence in Figure 4 is convincing for the ranges tested, but the paper does not push beyond 32,768 chips (the 150B model at 4× scaling from 8,192), so the upper bound of scalability is not probed. Megatron-LM and MegaScale have demonstrated scaling to 10,000+ GPUs, so AXLearn's demonstration is competitive but not record-setting.

A more subtle limitation: the scaling study uses proprietary models (Model A and Model B) with unspecified architectures. Without knowing what these models are, the scaling results are difficult to contextualize. Are these dense Transformers? MoE? Do they have unusually large or small layers that affect communication patterns? The paper provides only parameter counts and context lengths, which is insufficient to assess whether the scaling behavior is genuinely representative or specific to these particular architectures.

Claim 4: AXLearn's modular design extends to inference, achieving production-grade performance with minimal effort.

Table 4 shows AXLearn outperforming vLLM on TPUs, but the paper's own caveat — "TPU support for vLLM is still experimental" — substantially weakens this claim. The comparison is not against a mature, optimized TPU inference engine (which may not exist), but against a system whose TPU backend is described as experimental. The 500× TTFT improvement for 70B is dramatic but may primarily reflect vLLM's immaturity on TPU rather than AXLearn's inference efficiency. A fair comparison would require benchmarking against the best available inference engine on each platform (e.g., vLLM on GPU vs. AXLearn on GPU), which the paper does not report. The inference results are better understood as existence proof that modular training code can serve inference rather than as evidence of inference performance competitiveness.

Claim 5: AXLearn's modularity enables rapid experimentation at scale (1,000+ experiments, hundreds of engineers).

This claim is supported by operational statistics (Section 7.4: "over 10,000 experiments under development at a given time, running across tens of different hardware clusters") but these are self-reported without independent verification. The paper's narrative that AXLearn's adoption "largely owes to its modularity" is an interpretation of the adoption data, not a causal demonstration. Alternative explanations — organizational mandate, integration with Apple's infrastructure, lack of alternatives supporting Apple's hardware mix — are not ruled out.

What would strengthen the evaluation. Several additional experiments would substantially increase confidence in the paper's claims, though their absence is understandable for a production system paper:

1. Benchmarking on a non-Transformer architecture (e.g., a CNN, an SSM like Mamba, or a retrieval-augmented model) would test whether the composition-over-subtyping approach generalizes beyond the Transformer family, where the dominant pattern of "attention layer + FFN layer + normalization" makes drop-in replacement particularly natural.

2. A controlled experiment measuring development time, not just LoC. If AXLearn truly accelerates experimentation, a user study comparing the time to implement an MoE variant in AXLearn vs. DeepSpeed would be more directly interpretable than static LoC analysis.

3. Sensitivity analysis of mesh rules. The paper demonstrates one TPU rule and one GPU rule in Appendix A. How many such rules are needed in practice to cover Apple's hardware fleet? Does the rule-matching mechanism (regex on instance types) scale to tens or hundreds of hardware configurations, or does it become unwieldy?

4. AOT compilation accuracy data. How many training jobs were caught by AOT compilation that would have failed at scale? What is the false positive rate? Without these numbers, the practical value of AOT compilation is asserted rather than demonstrated.

5. Statistical characterization of the failure recovery data. A single recovery event (Figure 5) shows 21 minutes of lost time, but the distribution of recovery times across many failures would reveal whether this is typical or an outlier. The paper mentions "opaque failures often out of our control" (Section 7.4) but does not characterize their frequency or the recovery mechanism's success rate.

Overall assessment. The paper's experimental evidence is appropriate for a systems paper introducing a production training framework. The performance benchmarks demonstrate that AXLearn achieves its design goals (hardware agnosticism, modularity) without sacrificing training throughput — a necessary condition for adoption. The LoC analysis provides a compelling, if not rigorous, demonstration of the encapsulation design's benefits. The scaling study and failure recovery data establish production credibility. However, the evaluation is narrower than the paper's ambitious framing suggests. The $O(1)$ complexity claim is validated for exactly two features (RoPE, MoE) on Transformer architectures; the hardware agnosticism claim is validated on three backends; and the scalability claim is validated up to 32,768 chips. These are substantial achievements, but the paper's rhetoric sometimes outstrips its evidence, particularly in claiming that the LoC-complexity framework is a general metric rather than a diagnostic tool developed for the specific features and systems analyzed.

6. Limitations and Trade-offs

Encapsulation Benefits Demonstrated Only for Drop-In Replacement Patterns

The paper's central claim — that AXLearn achieves $O(1)$ complexity for feature integration — is validated exclusively through two features that fit naturally into a "drop-in replacement" pattern: Rotary Position Embeddings (RoPE) and Mixture of Experts (MoE). Both features replace an existing component (positional encoding, feed-forward network) with a variant that shares the same dimensional interface. The paper does not demonstrate $O(1)$ complexity for features that fundamentally change the model architecture in ways that cannot be expressed as config traversal and type substitution — for example, adding cross-attention layers to a decoder-only model, introducing a novel training objective that requires new state management, or implementing a feature like retrieval-augmented generation that requires new data flow patterns between non-adjacent modules.

The consequence is that the LoC-complexity framework, while rigorous for the evaluated cases, may not generalize. A practitioner attempting to add a feature that restructures the module hierarchy — interleaving new layers between existing ones, introducing skip connections that span multiple levels, or changing the forward/backward computation graph in ways that require new InvocationContext interactions — may find that AXLearn's encapsulation boundaries become obstacles rather than enablers. The paper acknowledges no such boundary, implying that all features can be expressed as drop-in replacements, which is an architectural claim about neural network design space (that all architectural innovations are compositions of interface-compatible modules) rather than a property of AXLearn. If a feature requires changing the interface between modules (e.g., adding a new tensor input to the attention layer that wasn't anticipated in AttentionLayer.Config), the config traversal approach would need to modify the parent's config class definition — breaking the $O(1)$ guarantee because the parent's Config must be updated to declare the new child interface field.

The paper provides no evidence on how LoC-complexity scales for a broader set of feature types. The two features tested (RoPE and MoE) were chosen because they are well-known and widely adopted, not because they stress-test the encapsulation architecture. The appendix provides detailed LoC analysis for these two features across seven competing systems (Appendix B), but the paper never discusses what kinds of features cannot be expressed as config modifiers, what the escape hatches are when encapsulation must be violated, or whether the system has been tested with features that required such violations. The mitigation status is that this limitation is unaddressed and unacknowledged — the paper presents the $O(1)$ claim as a general property of the architecture, not as an observation about two conveniently structured features.

Hardware-Agnostic Performance Relies on Backend-Specific Kernel Engineering

The paper is transparent that achieving peak hardware utilization requires backend-specific kernel implementations: Section 4.2 notes that the FlashAttention layer "transparently dispatches kernels based on the backend," using cuDNN on GPU, Nki kernels on Trainium, and SplashAttention on TPU. Section 7.4 further acknowledges that "achieving peak hardware utilization on each backend requires backend-specific kernel implementations, which demands specialized expertise and ongoing maintenance as hardware evolves." However, the paper understates the scope of this dependency.

The reported performance numbers (Table 3) depend critically on these kernels. On H100 GPUs, AXLearn achieves 54.2% MFU for Llama2-7B, trailing Megatron-LM's 44.9% MFU (with throughput as the primary metric), but the paper's own explanation — "PyTorch currently has finer-grained scheduling capability over XLA" — suggests that the XLA compilation path, not just kernel quality, imposes a performance ceiling. On Trainium2, AXLearn achieves only 24.2% MFU for Llama2-7B, dramatically lower than on other backends. The paper does not explain this gap. Is it due to immature kernels? Compiler limitations? Fundamental architectural properties of Trainium2? A practitioner evaluating Trainium2 for large-scale training cannot determine from the paper whether the 24.2% MFU is the best achievable or whether significant optimization headroom remains.

The consequence is that hardware agnosticism in AXLearn means API-level agnosticism but not performance portability. The same model code runs on all backends, but achieving competitive performance requires backend-specific kernel development that the paper describes as requiring "specialized expertise." This shifts the burden of hardware support from model developers (who would otherwise write backend-conditional code) to kernel developers (who must implement and maintain backend-specific attention, normalization, and quantization kernels). For organizations with the resources to maintain a kernel engineering team (like Apple), this is a manageable trade-off. For smaller teams or individual researchers who want to run their models on a new accelerator, the "hardware-agnostic" claim is misleading — they would need to either accept suboptimal performance or write custom kernels, which is precisely the expertise the paper aims to abstract away.

The paper provides no data on how much performance degrades when using the default (non-custom) kernel implementations. This is the critical missing ablation: what MFU does AXLearn achieve on H100 without the cuDNN/Pallas attention kernels? On Trainium2 without the Nki kernel? Without these numbers, a practitioner cannot assess the performance cost of the "hardware-agnostic" path versus the "hardware-optimized" path.

Inference Performance Claims Are Not Benchmarking Against Mature Inference Engines

The paper presents AXLearn's inference performance as evidence that "a modular training framework can achieve production-grade inference performance with minimal additional effort" (Section 7.2). The comparison in Table 4 shows AXLearn dramatically outperforming vLLM on TPUs — 500× faster TTFT for Llama2-70B, 2.8× higher throughput for Llama2-7B. However, the paper explicitly acknowledges that "TPU support for vLLM is still experimental, which likely contributes to the performance gap." This caveat is more consequential than the paper's framing suggests.

The comparison is between a training system that has been optimized for TPU inference (AXLearn on its native platform) and an inference engine whose TPU backend is, by the paper's own characterization, experimental. This is not a test of AXLearn's inference capabilities versus state-of-the-art inference systems — it is a test of AXLearn versus a system that is not production-ready on the target hardware. The fair comparison would be AXLearn versus the best inference engine available on each hardware platform: vLLM or TensorRT-LLM on GPU, a mature TPU inference engine (if one exists) on TPU. The paper reports no inference benchmarks on GPU or Trainium2.

The consequence is that the inference performance claims are uninterpretable as evidence of competitiveness. A practitioner considering AXLearn for inference cannot determine from Table 4 whether AXLearn's inference performance is genuinely state-of-the-art or merely better than an experimental backend. The 500× TTFT improvement for 70B is likely dominated by vLLM's TPU immaturity (an 80-second TTFT is not characteristic of production inference engines on any platform) rather than by AXLearn's architectural advantages. The paper's framing — presenting this as a "surprising discovery" that modular training code achieves high inference performance — is undercut by the choice of a baseline that does not establish the performance frontier.

The mitigation status is partial: the paper acknowledges the vLLM TPU limitation but does not address it with additional benchmarks on other platforms or against more mature baselines. Section 6 states that "with additional effort we can support unified training and inference on other backends," but provides no evidence or timeline.

Difficulty Estimation Cost for Config Modifier Selection Is Unaccounted For

A subtle but practically significant limitation: the mesh rule system and config modifiers that enable $O(1)$ complexity require that someone — either the user or a platform engineer — write the correct mesh rules for each hardware target. The paper demonstrates this with a roughly 10-line example in Appendix A that configures FSDP within TPU v5e slices, offloads activations, and enables INT8 training. But the paper does not address how a user determines what the correct mesh shape, rematerialization policy, or quantization strategy should be for a given model on a given hardware platform.

This is the AXLearn analog of the "difficulty estimation" problem in the compute-optimal test-time scaling paper — the config modifiers are powerful once written, but writing them requires deep expertise in the interaction between model architecture, hardware characteristics, parallelism strategies, and compiler behavior. The paper acknowledges implicitly that this expertise is scarce, noting that achieving peak utilization "requires backend-specific kernel implementations, which demands specialized expertise and ongoing maintenance" (Section 7.4), but frames this as a kernel development concern rather than a configuration concern. In practice, the mesh rules embed similar expertise: the decision to use mesh(data=-1, fsdp=256) on TPU v5e versus mesh(fsdp=-1, model=8) on H100 reflects non-obvious hardware tradeoffs (TPU v5e's interconnect characteristics, H100's NVLink bandwidth, memory capacity constraints) that a typical ML researcher would not know.

The consequence is that AXLearn shifts the bottleneck from code changes (modifying model implementations for each hardware target) to configuration expertise (writing correct mesh rules). For an organization like Apple with dedicated platform engineering teams, this is a net improvement: the platform team writes the mesh rules once, and hundreds of model engineers benefit. For smaller organizations or individual researchers working on a new hardware platform where no mesh rules exist, the situation is arguably worse than with a system like Megatron-LM — at least Megatron-LM's GPU-specific optimizations are transparent in the code and can be studied and adapted. AXLearn's mesh rules are external configuration that assumes the user knows the correct parallelism strategy, remat policy, and kernel selection for their model-hardware combination.

The paper provides no guidance on how to derive mesh rules for a new hardware platform, no tooling for automatic mesh rule generation (e.g., auto-tuning parallelism strategies), and no evaluation of how performance degrades with suboptimal mesh rules. The AOT compilation feature (Section 4.2) can catch OOMs from poor mesh choices, but cannot suggest better ones. This limitation is unacknowledged and unaddressed.

LoC-Complexity Estimates Are Expert Judgments, Not Systematically Measured

The paper's headline metric — the LoC-complexity framework and the concrete estimates in Table 2 — is presented as a rigorous, quantifiable evaluation of system modularity. However, the estimates are derived through manual code inspection and heuristic reasoning (Appendix B), not through systematic measurement. The paper states the assumptions: a "realistic production setting" with 20 model variants and 10 attention variants. But the LoC numbers are not generated by instrumenting the codebases, counting changed lines across actual commits that added RoPE or MoE support, or surveying system users about their integration experiences. They are the authors' estimates, based on reading public code and reasoning about what changes would be necessary.

This matters because the specific numbers drive the paper's strongest quantitative claims: AXLearn needs 0 LoC versus 320–600 for RoPE, 0 versus 20–4,000 for MoE. If the DeepSpeed MoE estimate (4,000 LoC) is off by a factor of 2×, the qualitative story (AXLearn is better) remains true but the quantitative gap is mischaracterized. More importantly, the estimates assume that every model variant in a codebase requires changes proportional to the public examples the authors inspected. In practice, a codebase might share infrastructure across variants (reducing the effective $N$) or might have idiosyncratic implementations that require more changes (increasing it). The paper's Appendix B reasoning for DeepSpeed — "for 20 model variants, each can incur 100s of LoC... which conservatively incurs 4,000 LoC changes" — assumes a linear scaling from one example (QwenV2MoE at 200+ LoC) without evidence that this linear extrapolation holds.

The consequence is that the LoC-complexity framework, while conceptually valuable, should be understood as a qualitative design analysis tool rather than a quantitatively validated metric. The asymptotic claims ($O(1)$ vs. $O(N)$ vs. $O(NM)$) are logically derived from architectural properties (encapsulation vs. subtyping vs. config flattening), but the concrete LoC estimates are back-of-the-envelope calculations that have not been validated against ground truth. A practitioner deciding between systems cannot take the Table 2 numbers as precise predictions of their own integration costs — their codebase structure, model variant count, and feature requirements will differ from the paper's idealized production setting.

The mitigation status: the paper is transparent about its methodology (Appendix B describes the derivation logic), but does not acknowledge the gap between expert estimates and systematic measurement. No sensitivity analysis explores how the estimates change under different assumptions about codebase structure.

Scaling Study Limited to Two Proprietary Models and Four-Fold Chip Scaling

The paper demonstrates AXLearn's scalability through a weak-scaling study (Figure 4) on two proprietary models (Model A at 70B, Model B at 150B) with scaling ranges of 256→4,096 chips (16×) for Model A and 8,192→32,768 chips (4×) for Model B. While the results show "near-linear scaling" with modest MFU decline, several aspects of this evaluation limit its generalizability.

First, the models are proprietary with unspecified architectures. Without knowing whether they are dense Transformers, MoE models, or use custom attention mechanisms, the observed scaling behavior cannot be attributed to AXLearn's design versus the models' characteristics. If both models happen to have communication patterns that map naturally to the parallelism strategies AXLearn employs, the scaling results may not generalize to architectures with different communication requirements (e.g., models with unusual tensor shapes, non-standard attention patterns, or dynamic computation graphs).

Second, the scaling range for Model B is only 4× (three data points: 8,192, 16,384, 32,768 chips). While 32,768 chips is a substantial absolute scale, the narrow range means the paper does not probe the regime where communication overhead typically becomes the dominant scaling bottleneck. Prior work (MegaScale, Jiang et al., 2024) has demonstrated scaling to 10,000+ GPUs with detailed characterization of the failure modes at extreme scale (stragglers, communication contention, failure frequency). AXLearn's scaling study does not characterize these effects or push into the regime where they become visible.

Third, the paper notes that the 150B model uses 1/16 the per-chip sequence length of the 70B model "for good training convergence" — a significant detail that complicates the weak-scaling comparison. Weak scaling typically fixes per-device work, but if the per-device sequence length changes by 16× between the two models, the computational intensity (FLOPs per byte of communication) is not comparable. The paper does not report the weak-scaling behavior of a single model across the full chip range, which would be the cleanest demonstration of scalability.

The consequence is that the scaling results, while promising, do not constitute a thorough characterization of AXLearn's scalability limits. A practitioner training a model with different architectural properties, or at scales beyond 32,768 chips, cannot confidently extrapolate from Figure 4. The mitigation status is that the paper does not claim to have fully characterized scalability limits, but also does not acknowledge the narrowness of the scaling evaluation relative to the claims of production-scale operation ("models at billion-to-trillion parameter scales," Section 7.4).

7. Implications and Future Directions

How This Work Changes the Landscape

This paper does not introduce a new training algorithm, a novel parallelism strategy, or a breakthrough in hardware utilization. Its contribution is architectural: it demonstrates that strict encapsulation — applied systematically across every component of a training system — can reduce feature integration complexity from $O(N)$ to $O(1)$ while matching state-of-the-art performance across multiple hardware backends. This is less a paradigm shift and more a correction to a widespread architectural error that the paper diagnoses across seven major training frameworks: the conflation of configurability with modularity, and the routine violation of encapsulation through subtyping, config flattening, and manual state threading.

The conceptual shift this work enables is the recognition that modularity in ML training systems can and should be measured quantitatively. Prior to this paper, discussions of system extensibility were qualitative — "this framework is flexible," "this library is easy to extend" — based on developer intuition or marketing claims. The LoC-complexity framework introduced in Section 7.1 provides a rigorous vocabulary: for a system with $N$ modules and a feature with $M$ variants, what is the asymptotic integration cost to existing interfaces? The paper's analysis shows that existing systems cluster at $O(N)$ or $O(NM)$, while AXLearn achieves $O(1)$. This transforms system design from a matter of taste into a matter of provable scaling properties. A follow-up system can now be evaluated not by whether it "feels modular" but by what its LoC-complexity is for a standard feature like RoPE or MoE.

The paper also reframes the conversation around hardware agnosticism. Prior work treated hardware support as an accumulation of backend-specific optimizations — add cuDNN for GPU, add SplashAttention for TPU, and so on. AXLearn's mesh rule architecture (Section 4.2, Appendix A) demonstrates that hardware agnosticism is fundamentally a separation-of-concerns property: all hardware-specific decisions (parallelism strategy, rematerialization policy, kernel selection, quantization) can be expressed as config modifiers applied externally to a shared layer graph, so long as the layer graph itself is sharding- and remat-aware. This shifts the burden of hardware support from model code (where it scales with $O(N)$) to configuration rules and backend kernel implementations (where it scales with $O(1)$ per hardware target). The paper validates this through the unplanned arrival of AWS Trainium2, which was supported without model code changes (Section 2.2). This is a different model of hardware portability than the field has previously articulated, and it makes the case that compiler-injected parallelism (GSPMD/XLA) is not just a performance optimization but an architectural enabler of encapsulation.

The paper resolves a latent contradiction that practitioners have lived with: why do ML frameworks feel harder to extend as they grow, even though they provide configuration systems? The answer, per the paper's analysis, is that configuration systems without encapsulation (config flattening in TorchTitan and DeepSpeed) still propagate changes through the module hierarchy because every consumer of the config must be modified when new fields are added. The induction argument in Section 2.1 — replacing a child forces the parent to change, which forces the grandparent — explains why a 4-line change at the FFN layer becomes a 200-line change across the model in DeepSpeed's QwenV2MoE transition. This is not a failure of developer discipline; it is a failure of the abstraction. AXLearn's correction — hierarchical, partially-specified configs where parents only know child interfaces, not child internals — provides the design pattern that breaks the induction.

Certain research directions become more attractive in light of this work. Compiler-first training architectures (GSPMD/XLA) receive a strong endorsement because they separate communication from model code, which is what makes sharding-aware layers possible without coupling to parallelism strategy. Automated parallelism optimization (auto-tuning mesh shapes, remat policies, and sharding annotations) becomes more important because the paper shows these can be cleanly separated from model implementation — but the paper does not solve the problem of determining the optimal configuration, only the problem of applying it without code changes. Standardized layer interfaces for common architectural patterns (attention layers, FFN layers, normalization) become more valuable because drop-in replacement is the mechanism through which $O(1)$ integration is achieved.

Conversely, some directions become less attractive. Fork-and-modify model implementations (the MaxText approach, where each model variant is implemented by copying and customizing) appear architecturally fragile after this analysis — the paper estimates 200–300 LoC for feature integration in MaxText (Table 2) despite its JAX foundation. Monolithic config classes (TorchTitan's ModelArgs, DeepSpeed's DeepSpeedInferenceConfig) are directly implicated as the mechanism behind $O(NM)$ scaling and should be deprecated in favor of hierarchical, composed configurations. Manual state threading in functional frameworks (the pattern the paper describes in early Flax, Section 7.4) is shown to produce the same $O(N)$ complexity as subtyping, and the InvocationContext abstraction provides a proven alternative.

Follow-Up Research This Work Enables

LoC-complexity benchmarking as a standard evaluation for training systems. The paper proposes LoC-complexity as a metric but validates it on only two features (RoPE, MoE) with expert estimates rather than systematic measurement. A strong follow-up would develop a benchmark suite of standard features — say, ten common architectural modifications (grouped-query attention, sliding window attention, different normalization schemes, different activation functions, weight tying, sparse attention patterns, adapter layers, retrieval augmentation hooks, quantization-aware training wrappers, and novel loss functions) — and measure the actual LoC changes required to integrate each feature into each major training framework using a standardized set of model architectures. This would validate whether AXLearn's $O(1)$ property holds for features beyond drop-in replacements, and would measure the true LoC overhead in each system through instrumentation rather than expert judgment. The benchmark could be maintained as a community resource, similar to MLPerf for performance, and would provide a rigorous foundation for the claims this paper makes qualitatively.

Boundary testing of the composition-over-subtyping paradigm. The paper demonstrates $O(1)$ complexity for features that replace existing modules with interface-compatible variants. A critical stress test would be to implement features that change the interface between modules — for instance, adding cross-attention to a decoder-only model, introducing per-layer learning rates, or implementing a training objective that requires auxiliary outputs from intermediate layers. These features cannot be expressed as replace_config traversals because they require the parent module's config class to declare new child interfaces. Do they force a reversion to $O(N)$ complexity in AXLearn? Or does AXLearn provide mechanisms (e.g., generic extension fields in Config, dynamic child registration) that preserve $O(1)$ even for interface-changing features? Answering this question would define the true boundaries of the encapsulation approach and identify which kinds of architectural innovation are naturally supported versus which require framework modification.

Automated mesh rule generation through performance auto-tuning. The paper identifies but does not solve the problem of determining correct mesh rules for a given model-hardware combination. A valuable follow-up would develop a system that automatically searches the space of mesh shapes, rematerialization policies, and sharding annotations for a given model and hardware target, using a cost model (estimated via AOT compilation FLOP counts and memory usage) or a small number of profiling runs. The AXLearn AOT compilation infrastructure (Section 4.2) is uniquely suited for this because it can evaluate configurations locally without consuming accelerator time. A system that takes a model config, a hardware target, and returns a mesh rule achieving ≥90% of optimal MFU without human intervention would convert AXLearn's expert-dependent configuration system into a self-service capability. The paper's existing mesh rule examples (Appendix A) provide the template; the missing piece is the search algorithm and cost model.

Quantifying the developer velocity impact of encapsulation with a controlled user study. The paper claims AXLearn's modularity accelerates experimentation (Section 7.4: "thousands of models involving hundreds of engineers... rapid adoption largely owes to its modularity"), but provides only adoption statistics as evidence. A controlled study could measure time-to-implementation for a set of standardized ML engineering tasks — add RoPE to an unfamiliar model, swap FFN for MoE, port a model from GPU to TPU, integrate a new attention kernel — across AXLearn and one or two competing systems, with participants of comparable experience. The outcome would be not LoC counts but wall-clock time, error rates, and participant-reported frustration. This would test whether the $O(1)$ complexity translates to real productivity gains or whether the config traversal approach introduces new failure modes (silent misconfigurations, difficulty debugging config modifiers) that offset the code reduction.

Inference performance parity study across hardware backends. The paper's inference benchmarking (Table 4) compares AXLearn against vLLM on TPU only, where vLLM's support is experimental. A rigorous follow-up would measure AXLearn inference performance against the best available inference engine on each hardware platform: vLLM or TensorRT-LLM on GPU, a mature TPU serving solution on TPU, and AWS Neuron-based serving on Trainium2. The goal would be to determine whether AXLearn's modular training code can achieve inference performance within, say, 20% of specialized inference engines across all platforms — which would make the "unified training and inference" claim practically compelling — or whether the 2.8× advantage in Table 4 is entirely attributable to the baseline's immaturity rather than AXLearn's architectural advantages. The paper acknowledges this limitation (Section 7.2: "Our goal is not to position AXLearn as a specialized inference engine"), making a rigorous multi-backend comparison the natural next step.

Practical Applications and Downstream Use Cases

Multi-cloud, multi-vendor training for organizations hedging hardware risk. The paper's most directly actionable contribution is the mesh rule architecture that enables the same model code to train efficiently on GPU, TPU, and Trainium without modification. For any organization that cannot or will not commit to a single hardware vendor — because of supply constraints, pricing dynamics, or strategic risk management — AXLearn's approach provides a concrete template. The marginal cost of adding a new hardware backend is implementing backend-specific kernels behind existing layer interfaces and writing a mesh rule entry (Appendix A), rather than modifying every model implementation. The paper's unplanned Trainium2 support (Section 2.2) demonstrates this works in practice. Organizations can adopt this architecture to maintain a single model codebase while deploying training to whichever hardware is most available or cost-effective at a given time, with the mesh rules encoding per-platform optimization expertise written once by platform engineers and reused across all model teams.

Large-team ML research environments with high model variant diversity. The paper's operational context — "thousands of models involving hundreds of engineers," "over 10,000 experiments under development at a given time," across "tens of different hardware clusters" (Section 7.4) — is unusually large, but the underlying dynamic applies at smaller scale. In any organization where multiple teams develop model variants that share architectural components, the $O(1)$ versus $O(N)$ complexity gap translates to reduced merge conflicts, fewer bugs from incomplete propagation of config changes, and faster iteration when a new architectural idea (like RoPE or MoE) needs to be tested across all variants. The paper's concrete demonstration — a 10-line config modifier that applies MoE to over 1,000 experiments without additional changes — is directly transferable. Organizations can adopt the hierarchical config + config modifier pattern even without adopting the full AXLearn stack, by restructuring their existing configuration systems to use partial specification with top-down resolution and external traversal-based modification rather than monolithic config classes with per-model overrides.

AOT-based resource-efficient development for teams with limited accelerator access. The paper describes AOT compilation (Section 4.2) as a mechanism that "allowed users to debug training entirely on CPU" and was critical for scaling development when TPU capacity was limited (Section 7.4). For academic groups, startups, or teams in developing regions that have limited access to expensive accelerators, AXLearn's AOT workflow provides a template for catching memory errors, sharding mistakes, and OOM conditions before consuming scarce GPU/TPU hours. The reliability guarantee — "because the same codepath is used for AOT and actual training, users can be confident that a program that AOT-compiles will run at a larger scale" — means that AOT is not a heuristic but a precise predictor of distributed execution feasibility. Integrating AOT compilation into CI/CD pipelines for model changes (a "golden configuration" testing suite, as the paper describes in Section 7.4) could dramatically reduce wasted compute from preventable configuration errors.

When to Prefer This Method

The paper articulates an explicit tradeoff: AXLearn's encapsulation-first design versus Megatron-LM's hardware-specific optimization on GPUs. Section 7.2 states directly that Megatron-LM "has stronger performance on H100 GPUs... because PyTorch currently has finer-grained scheduling capability over XLA. However, using XLA allows AXLearn to be hardware-agnostic. This is a trade-off we are willing to take." This is not a universal prescription but a context-dependent choice. The conditions are:

• Prefer AXLearn's encapsulation architecture when: (a) your organization trains models across multiple hardware backends (GPU, TPU, custom accelerators), or anticipates needing to switch between them based on availability or cost; (b) you maintain a codebase with many model variants (tens to hundreds) that share architectural components, and changes must propagate across them frequently; (c) your team includes model researchers who should not need to understand parallelism strategies or hardware-specific optimizations; (d) you can invest in backend-specific kernel development to close the performance gap on each platform; (e) your training scale reaches the regime where AOT compilation catching errors locally saves significant wasted compute.

• Prefer Megatron-LM or other hardware-tuned frameworks when: (a) you train exclusively on NVIDIA GPUs and have no plans to diversify hardware; (b) your organization has optimized your entire training pipeline around PyTorch and cannot absorb the migration cost to JAX/XLA, even if the paper shows the migration can be incremental (Section 7.4); (c) you need the absolute maximum performance on a single hardware platform and the ~10% throughput advantage Megatron-LM shows on H100 (Table 3: 2.5M vs. 3.0M tokens/s for Llama2-7B, though note the throughput metric favors AXLearn while MFU favors Megatron-LM) is material to your training budget — understanding that this advantage may narrow as XLA GPU compilation matures.

• The paper does not position AXLearn against DeepSpeed, TorchTitan, or MaxText as performance competitors — the modularity gap (Table 2) is the primary differentiator, not throughput. For teams where modularity is not a bottleneck (small codebase, few model variants, stable architecture), the performance differences in Table 3 are within the range where other factors (team familiarity, ecosystem integration, deployment tooling) would dominate the framework choice.