TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling

1

TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling

arXiv:2608.15211v1 [cs.CV] 15 Aug 2026

Ruohan Wu, Ziqi Zhu, Yang Zhao, Jiarui Tang, Yingzhe Cui, Junshi Chen, Zhao Jing, Jun Shi, and Hong An

Abstract—Training high-resolution AI-based Earth forecasting models is memory-intensive. Window-based Swin Transformers reduce the quadratic cost of global attention, but existing distributed systems such as AERIS primarily target pixel-level models and do not jointly support convolutional sampling modules and shifted-window execution. Long-lead rollout finetuning further increases activation memory. To address these challenges, we present TERRA, a hierarchical parallel training framework for high-resolution Earth forecasting. TERRA introduces SamplingAware Window, Sequence, and Tensor Parallelism (SAWSTP), which preserves spatially contiguous layouts for sampling modules and routes tokens into topology-aware ragged window layouts for Transformer execution. For long-lead finetuning, Memory Orchestration (MO) provides rollout-aware checkpoint planning and combines input buffering with budget-constrained activation offloading. Experiments on the 1/12◦ GLORYS-based Wenhai workload show that TERRA supports models with up to 11.4B parameters on 96 H200 GPUs and sustains up to 39.76 PFLOPS, achieving 65.0% strong-scaling and 94.1% weak-scaling efficiency. Compared with checkpoint-only policies, MO further reduces peak allocated GPU memory by 32.2%–51.8% with at most 20.0% step-time overhead, which makes finetuning with smaller patch sizes and longer rollouts feasible for improved forecasting accuracy. Index Terms—Deep learning, weather and ocean forecasting, distributed training, Swin Transformer, memory management.

I. I NTRODUCTION

A

I-based Earth forecasting has emerged as a promising approach for advancing high-resolution weather and ocean

This work was supported in part by Laoshan National Laboratory under Grant LSKJ202300305 and in part by the Anhui Provincial Natural Science Foundation under Grant 2508085QF210. (Corresponding authors: Jun Shi and Hong An.) Ruohan Wu and Jun Shi are with the School of Computer Science and Technology, University of Science and Technology of China, Hefei 230026, China (e-mail: ruohanwu@mail.ustc.edu.cn; shijun18@ustc.edu.cn). Ziqi Zhu is with the School of Artificial Intelligence and Data Science, University of Science and Technology of China, Hefei 230026, China (e-mail: ta1ly@mail.ustc.edu.cn). Yang Zhao, Jiarui Tang, Junshi Chen, and Hong An are with the School of Computer Science and Technology, University of Science and Technology of China, Hefei 230026, China, and also with the Department of Ocean Big Data and Artificial Intelligence, Laoshan Laboratory, Qingdao, China (email: yang zhao@mail.ustc.edu.cn; tjr@mail.ustc.edu.cn; cjuns@ustc.edu.cn; han@ustc.edu.cn). Yingzhe Cui is with the Department of Ocean Big Data and Artificial Intelligence, Laoshan Laboratory, Qingdao, China (e-mail: cuiyingzhe@lsnl.cn). Zhao Jing is with the Frontiers Science Center for Deep Ocean Multispheres and Earth System and the Key Laboratory of Physical Oceanography/Academy of Future Ocean, Ocean University of China, Qingdao, China, and also with the Department of Ocean Big Data and Artificial Intelligence, Laoshan Laboratory, Qingdao, China (e-mail: jingzhao@ouc.edu.cn).

prediction [1], [2], [3], [4], [5], [6], [7]. Recent models increasingly adopt finer spatial resolutions to represent complex Earth-system dynamics. However, training these models at scale remains challenging because high-resolution, multivariate training data increases the number of spatial tokens and creates large activation tensors. Recent systems have demonstrated that large-scale parallel training can support Earth forecasting models with greater parameter capacity. ORBIT [8] combines Fully Sharded Data Parallelism (FSDP) [9] with tensor parallelism (TP) [10] to train Vision Transformer (ViT)-based Earth models [11] with up to 113B parameters on 1.4◦ data. At higher spatial resolutions, however, activation memory rather than parameterrelated memory can become the dominant training bottleneck. Window-based Swin Transformers [12] are therefore particularly attractive because they avoid the quadratic cost of global attention. AERIS [13] scales pixel-level Swin Transformers with up to 80B parameters on 0.25◦ data through SequenceWindow Parallelism (SWiPe), which combines window parallelism (WP), sequence parallelism (SP) [14], and pipeline parallelism (PP) [15]. Despite this progress, the SWiPe method for pixel-level Swin Transformers still faces several challenges. First, SWiPe does not support convolutional sampling modules, which reduce the token resolution and activation memory of Transformer blocks and can also improve forecasting accuracy. Second, SWiPe relies on PP, whose bubbles cannot be effectively amortized under small global batch sizes. Third, AERIS considers only fixed window assignments. When the window size is small and the number of windows per rank is sufficient, complete-window WP provides sufficient parallelism, making SP communication unnecessary. When the window size is large, fixed window assignment requires substantial padding to maintain uniform window ownership. Its roundrobin ownership also increases the communication required by shifted-window execution. Moreover, using smaller patch sizes [16] and finetuning with longer lead times can improve long-lead forecasting accuracy. However, these methods also significantly increase activation memory, especially for the 1/12◦ GLORYS-based Wenhai workload [6], [17]. To address these challenges, we present TERRA, a hierarchical parallel training and memory orchestration framework for high-resolution Earth forecasting. TERRA introduces Sampling-Aware Window, Sequence, and Tensor Parallelism (SAWSTP), which defines an I/O topology for sampling modules and a Transformer topology for Swin Transformer

2

blocks. The I/O topology applies Domain Parallelism (DMP) to partition high-resolution inputs into spatially contiguous regions and parallelizes convolutional sampling modules through halo exchange [18]. The Transformer topology integrates WP, SP, and TP with FSDP, avoiding PP and its pipeline bubbles at small global batch sizes. Differentiable token routes connect the two topologies by converting spatial DMP partitions into the window-oriented layouts required by Transformer blocks and support distributed shifted-window execution across different WP, SP, and TP configurations. Building on these token routes, TERRA introduces costaware window assignment to adapt the Transformer topology to different window sizes and parallel configurations. It uses ragged window ownership and selects among WP/SP topologies to reduce padding and shifted-window communication while balancing workloads. To reduce the activation-memory pressure of long-lead finetuning, TERRA further introduces Memory Orchestration (MO). Building on activation checkpointing [19], [20], [21], [22], MO separately profiles the retained activation memory and backward recomputation peaks of sampling and Transformer checkpoint units. It uses these measurements to construct a lead-dependent checkpoint plan that balances rollout-wide activation accumulation against local recomputation peaks. Input buffering bounds the GPU lifetime of rollout inputs and labels. If the predicted peak still exceeds the target memory budget, MO selectively offloads activation boundaries to host memory [23], [24], [25], [26] until the predicted peak satisfies the budget. We implement TERRA as an open-source training framework available at https://github.com/ruohan12345/TERRA and evaluate it through pretraining and rollout finetuning of hierarchical GLORYS-based Wenhai models. The contributions of this work are summarized as follows: • We design SAWSTP, a parallelism strategy that applies halo-exchange-based DMP to convolutional sampling modules and combines WP, SP, TP, and FSDP for Transformer blocks. It scales Wenhai model training to 96 H200 GPUs, achieving up to 65.0% strong-scaling and 94.1% weak-scaling efficiency, with sustained performance of up to 39.76 PFLOPS. • We introduce cost-aware ragged window assignment that adapts the WP/SP topology to the window size and parallel configuration. It reduces padding and shiftedwindow communication while balancing workloads, providing a 1.03×–1.18× training-step speedup over the fixed AERIS-style baseline. • We develop MO for long-lead finetuning by combining profile-guided, lead-dependent checkpoint planning with input buffering and budget-constrained activation offloading. MO reduces peak allocated GPU memory by 32.2%– 51.8% relative to checkpoint-only policies, with at most 20.0% step-time overhead. • We pretrain, finetune, and evaluate hierarchical Wenhai models on the GLORYS-based workload. The results demonstrate that hierarchical sampling, smaller patch sizes, and longer rollout finetuning enabled by TERRA improve long-lead forecasting accuracy.

TABLE I N OTATION FOR THE W ENHAI MODEL . Notation H, W Ĥ, Ŵ C p α d Nh w L

Description Spatial resolution of the input field Padded spatial resolution of the input field Number of input physical variables Patch size Stride of each down- or up-sampling module Embedding dimension Number of attention heads Attention-window size Number of Swin Transformer blocks

II. BACKGROUND AND M OTIVATION A. AI-based Earth Forecasting Models AI-based models have become an important approach for global weather and ocean forecasting [1], [2], [3], [4], [5], [6], [7]. However, the increasing spatial resolution substantially raises the computational and memory requirements of model training. For example, Pangu-Weather is trained on 1/4◦ ERA5 data with 721 × 1440 grid points [2], [27], whereas Wenhai uses 1/12◦ GLORYS data with 2041 × 4320 grid points [6], [17]. TERRA targets the Wenhai workload, whose high-resolution multivariate ocean inputs impose substantial memory pressure during training. In Wenhai, each input is a dense multivariate ocean field on a regular latitude–longitude grid, represented as Xt ∈ RB×H×W ×C , where B is the micro-batch size, H and W are the spatial dimensions, and C = 93 is the number of ocean variables. As shown in Fig. 1(a), Wenhai first maps the input field to patch tokens, which are then processed by a hierarchical encoder–decoder architecture that combines convolutional down- and up-sampling modules with windowbased Swin Transformer blocks [12]. This design reduces the number of tokens processed by the Transformer blocks while preserving multiscale spatial representations. Table I summarizes the main notation of the Wenhai model. As shown in Fig. 1(b), AI-based Earth forecasting is autoregressive. Given an initial state Xt at time t, the predicted state at lead time k∆t is fed back as the input to the next forecast step, producing the state at lead time (k + 1)∆t. To improve long-lead forecast accuracy, models are typically pretrained for single-step forecasting and then finetuned over multiple autoregressive steps [1], [3], [4], [6]. During pretraining, the model learns a single-step forecasting operator fθ that maps Xt to Xt+∆t . During finetuning, the same model parameters are reused over a K-step rollout:   X̂t+k∆t = fθ X̂t+(k−1)∆t , 1 ≤ k ≤ K, (1) where X̂t = Xt and K is the rollout length. The finetuning objective sums the prediction losses over the K forecast steps, which directly optimizes long-lead forecast accuracy. However, the finetuning also retains intermediate activations across the rollout until backpropagation, which substantially increase the memory footprint.

3

Lead t

Pred Lead t+1

Finetune Model

Pretrain Model

Pred Lead t+1

Model

Pred Lead t+2

Lead t+1

Loss Function

Lead t+2

Loss Function

Lead t Pad

Pad Back

Patch Embedding

Output Embedding

Loss 1 Down Blocks

Swin Transformer Blocks

Pre-Processing

Up Blocks

…

Loss 2

Final Loss

Post-Processing

(a) Hierarchical Swin-based model.

(b) One-step pretraining and multi-step rollout finetuning.

Training high-resolution forecasting models presents both substantial I/O overhead and memory pressure. A single training sample can contain a large multivariate spatial field, making the micro-batch size per GPU typically small (e.g. B = 1). Recent systems therefore distribute one sample across multiple devices to reduce the per-rank input and activation footprint [8], [13], [28], [29], [30]. The GPU memory footprint of DNN training consists primarily of model parameters, gradients, optimizer states, and activations [31]. Different parallel and memory-management techniques target different components. Data parallelism (DP) replicates the model and distributes samples across ranks, but it cannot reduce the memory required by an individual high-resolution sample. FSDP and ZeRO-style methods shard parameters, gradients, and optimizer states [9], [32], [33], [34], [35]. TP [10], [36] partitions Transformer parameters and intermediate activations, SP [14] partitions token sequences, and PP partitions model layers into stages [15]. Activation checkpointing trades additional recomputation for a smaller activation footprint [19], [20], [21], [22], while activation offloading transfers selected activations outside GPU memory for later reuse [23], [24], [25], [26]. To improve the efficiency of large-scale Earth modeling, a series of memory-efficient parallel training systems has been proposed. For example, Jigsaw focuses on spatially partitioned execution for high-resolution inputs through distributed patch embedding and matrix operations [28]. ORBIT combines FSDP and tensor parallelism to scale the parameter capacity of ViT-based Earth foundation models [8], while ORBIT-2 further targets hyper-resolution climate downscaling with spatially partitioned execution of extremely long token sequences [29]. AERIS addresses distributed training of pixellevel Swin Transformers through SWiPe, which partitions window activations and distributes the computation of shiftedwindow attention blocks across GPUs [13].

C. Motivation The following observations motivate the design of TERRA for high-resolution hierarchical Earth forecasting models.

No sampling With sampling OOM

100 80 60 40 20 0 p=2

p=3

p=4

p=6

p=8

(a) Patch-size memory. Peak allocated GPU memory (GB)

B. Memory-Efficient Parallel Training

Peak allocated GPU memory (GB)

Fig. 1. High-resolution Earth forecasting model and training stages.

p=4 p=6 p=8 OOM

100 80 60 40 20 0 1

2

4

6

Finetune lead time

(b) Rollout memory. Fig. 2. Memory pressure in training high-resolution Earth forecasting models.

1) Limitations of Existing Swin Parallelism for HighResolution Earth Modeling and Parameter Scaling: Prior work has shown that larger models and smaller patch sizes can improve weather-forecasting accuracy [16]. Existing distributed Swin Transformer systems such as AERIS use SWiPe to train pixel-level Swin Transformers with a patch size of p = 1 and scale models to 80B parameters [13]. However, extending this design to higher-resolution training data and more parameters presents additional challenges. Finer spatial resolutions and smaller patch sizes both increase the number of patch tokens and the associated activation memory. AERIS is evaluated on 0.25◦ data, whereas Wenhai uses 1/12◦ data. At a fixed patch size, this resolution difference produces roughly nine times more tokens and activations. As shown in Fig. 2(a), reducing the patch size substantially

4

increases peak memory, while sampling modules alleviate this pressure by reducing the token resolution processed by the Transformer blocks and can improve forecasting accuracy through multi-scale feature learning. However, combining parallel Transformer execution with parallel sampling modules remains challenging because they favor different tensor layouts. Sampling modules require spatially contiguous partitions for convolution, whereas Swin Transformer blocks require window-major token layouts for window attention. For parameter scaling, Table II summarizes the model settings and supported parallelism of ORBIT, AERIS, and TERRA. AERIS relies on PP to partition Transformer layers across devices. Its WP and SP distribute activation across multiple GPUs but do not reduce parameter-related memory [13]. PP bubbles are difficult to amortize under small global batch sizes. For example, Wenhai uses a global batch size of only 16. ORBIT instead combines TP with FSDP to scale Earth foundation models [8]. However, complex distributed windowattention operations with shifted windows make TP integration with SWiPe challenging. 2) Efficient Window Assignment Requires Topology-Aware Optimization: To support distributed shifted-window execution, SWiPe adopts a fixed window-assignment strategy that maps attention windows to ranks through inter-window WP and intra-window SP. However, this fixed mapping does not adapt to the number of complete windows available under a given parallel topology. When many complete windows are available per rank, the fixed WP mapping can incur avoidable WP communication during distributed window shifting, as well as padding overhead from its regular layout. TERRA mitigates both sources of overhead through a row-contiguous window-assignment method. When only a few complete windows are available per rank, WP may not divide the window grid evenly. Padding the window grid to equalize window counts increases computation and memory consumption, whereas avoiding padding can leave ranks with imbalanced workloads. Although SP can mitigate this imbalance by further partitioning individual windows across ranks, it introduces additional all-to-all communication. Therefore, window assignment should jointly consider padding overhead, window-shift communication, workload balance, and the required topology of WP and SP. 3) Long-Lead Finetuning Requires Activation-Aware Memory Management: Long-lead finetuning fundamentally changes the memory characteristics of high-resolution Earthmodel training. During one-step pretraining, activations are bounded within a single forward–backward pass. In contrast, rollout finetuning repeatedly reuses the same model parameters across multiple forecast steps and retains intermediate rollout states and activations until backpropagation, making long-lead finetuning increasingly dominated by activation memory. As shown in Fig. 2(b), peak memory increases rapidly with rollout lead time and can eventually cause out-of-memory (OOM) failures. Although activation checkpointing and activation offloading have been extensively studied for DNN training, existing policies do not adequately capture the heterogeneous memory behavior of sampling and Transformer checkpoint units in

TERRA Hierarchical Swin Transformer

Task

Parallel Framework SAWSTP (DP, DMP, WP, SP, TP, FSDP)

Cost-aware Window Assignment

Pretraining

Input Resolution

Memory Orchestration

Patch size

Window Size

Checkpoint Planning

Budget Offloading

Input Buffer

Finetuning

Fig. 3. Overview of TERRA.

hierarchical Earth models. High-resolution sampling blocks retain individually large activation tensors, which are difficult to mitigate through activation checkpointing alone. Moreover, long-lead finetuning requires high-resolution input and label tensors from multiple lead times to reside on the GPU. These tensors alone impose a substantial memory footprint. Consequently, activation-memory optimization should employ separate strategies for different activation units while accounting for the rollout process and remaining compatible with distributed execution, rather than relying on a single global checkpointing or offloading strategy. III. TERRA OVERVIEW To address the aforementioned challenges, we propose TERRA, a hierarchical parallel training and memoryorchestration framework for high-resolution AI-based Earth forecasting models. As shown in Fig. 3, TERRA targets hierarchical Swin Transformer models and supports configurations with varying input resolutions, patch sizes, window sizes, and model scales. It integrates SAWSTP with MO to support distributed pretraining and long-lead finetuning under these configurations. SAWSTP partitions each high-resolution training sample into spatially contiguous DMP shards for patch embedding and sampling modules, and then routes the resulting tokens to Swin Transformer blocks executed with WP, SP, TP, and FSDP. Building on SAWSTP, TERRA employs cost-aware window assignment to select a ragged window assignment based on the input and parallel configuration. During long-lead finetuning, MO manages the per-rank memory pressure arising from rollout activations and high-resolution input and label data while preserving the execution semantics of SAWSTP. IV. PARALLEL T RAINING F RAMEWORK A. Hierarchical Swin Parallelism As shown in Fig. 4, TERRA realizes distributed execution of hierarchical Swin Transformer models through SAWSTP. Hierarchical Swin Transformer models combine convolutional down- and up-sampling modules with Swin Transformer blocks that contain window-based multi-head self-attention (W-MSA) and shifted window-based multi-head self-attention (SW-MSA) [12]. Sampling modules operate on spatially contiguous feature maps to preserve local convolutional dependencies, whereas Transformer blocks require window-major token layouts for distributed attention and shifted-window execution.

5

TABLE II C OMPARISON OF MODEL SETTINGS AND PARALLELISM SUPPORT IN ORBIT [8], AERIS [13], AND TERRA. Model setting Model architecture Vision Transformer Pixel-level Swin Transformer Hierarchical Swin Transformer

Framework ORBIT AERIS TERRA

DP

Parallel sampling × × ✓

WP × ✓ ✓

Parallelism support SP TP × ✓ ✓ × ✓ ✓

I/O Topology

PP × ✓ ×

DP ✓ ✓ ✓

FSDP ✓ × ✓

I/O Topology

Transformer Topology Rank 0

Patch Embedding

Rank 1 DMP

Resolution 1.4◦ 0.25◦ 1/12◦

Rank 2

Swin Transformer Blocks

DMP Token Route

Down Blocks

Rank 3

Top Halo

DMP Token Route

Patch Recovery Up Blocks

Parallel Sampling Blocks

Parallel Swin Transformer Blocks ×2

DMP Conv / TConv Stride=2

DMP Conv Stride=1

DMP GN

Block 0 W-MSA

Block 1 SW-MSA

Block 0 MLP

Block 1 MLP

Rank 0 Rank 1 Rank 2 Rank 3

Layernorm Shift Token Route

SiLU

Bottom Halo

Window Attention Shift Token Route

Fig. 4. SAWSTP Training Framework for the hierarchical Swin Transformer.

To overcome this mismatch, SAWSTP employs two stagespecific logical topologies, namely a sampling-aware I/O topology and a multi-level Transformer topology. The two topologies use the same participating ranks but organize them into different communication groups. The I/O topology is built on DMP, which partitions each high-resolution training sample and its corresponding labels into spatially contiguous shards. Each rank therefore loads and processes only the input and label tensors that it owns. This topology preserves the DMP layout through patch embedding, sampling, patch recovery, and loss computation. The Transformer topology remaps the sampled feature maps into window-major layouts for Swin Transformer execution. It factorizes Transformer computation with WP, SP, and TP, which determine complete-window ownership, intra-window sequence partitioning, and operator-level sharding, respectively. FSDP is orthogonal to these parallel dimensions and can shard parameters, gradients, and optimizer states over the combined DP × WP × SP process group. The differentiable DMP token route connects the two topologies by converting between DMP-partitioned and window-major layouts. Within the Transformer topology, the shift token route redistributes tokens for shifted-window attention. During backpropagation, the inverse routes ensure correct gradient propagation.

B. Sampling-Aware I/O topology The I/O topology first pads the input field Xt ∈ RB×H×W ×C to a padded shape of B × Ĥ × Ŵ × C, where Ĥ and Ŵ satisfy the divisibility requirements of patch embedding, sampling modules, and subsequent window formation. It then partitions the padded input into Pg contiguous spatial stripes along the height dimension, where Pg is the number

of ranks in the DMP group. Each rank owns a local tile and directly loads only its corresponding input and label tensors. Given a patch size p, each rank applies patch embedding to its local tile, producing a feature map of shape B× PĤg p × Ŵ p ×d, where d is the embedding dimension. This feature map serves as the input to the hierarchical sampling modules. Downsampling and up-sampling modules use strided convolution (Conv) and transposed convolution (TConv), respectively, with stride α = 2 in this work. Each sampling module also contains residual blocks composed of Conv, Group Normalization (GN), and SiLU operators. Consequently, a downsampling module reduces the local feature-map shape to Ŵ B × αPĤg p × αp × d, while the corresponding up-sampling module restores the spatial resolution. The down-sampled features are subsequently routed to the Transformer blocks, substantially reducing Transformer-stage activation memory. This encoder–decoder sampling design also supports multiscale feature learning, and the evaluation in Fig. 15(a) confirms its benefit to forecasting accuracy. Parallel convolution in sampling modules has been extensively studied in prior training systems [18], [30], [37], [38], [39]. SAWSTP adopts a halo-exchange-based implementation [39] that partitions activations along the height dimension while keeping the Conv, TConv, and GN parameters replicated across DMP ranks. Before each local Conv or TConv operation, neighboring ranks exchange only the boundary rows required by the operator. The communication volume scales with the halo width rather than the full local feature-map size. The halo exchange is differentiable, allowing boundary gradients to be returned to their original ranks during backpropagation.

6

WP Topo(2, 2, 1, 1) (0, 0) (0, 1) (0, 0) (0, 1) WP-x

P2P

Shift Token Route (1, 0) (1, 1) (1, 0) (1, 1)

WP-y

WP-SP Topo(1, 2, 2, 1)

I/O Topology: DMP Stripe Layout

DMP Token Route

(0, 1, 0)

(0, 0, 0) (0, 0)

(0, 1)

(0, 0)

(0, 1)

Shift Token Route

SP: All-to-All

All-to-All

(0, 1, 1)

(0, 0, 1)

(0, 0, 0)

(0, 1, 0)

(0, 0, 1)

(0, 1, 1)

WP-SP-TP Topo(1, 1, 2, 2) TP: All-gather+Reduce-scatter

Color for a complete attention window

(0, 0, 0) (0, 0)

(0, 0, 0, 0)

(0, 0)

(0, 0, 0, 1)

Shift Token Route

SP: All-to-All (0, 0)

(0, 0)

(0, 0, 1)

(0, 0, 1, 0)

(0, 0, 1, 1)

Fig. 5. Multi-level Transformer topology with DMP and shift token routes in SAWSTP.

C. Multi-Level Transformer topology After the sampling modules produce lower-resolution feature maps under the I/O topology, the spatially contiguous DMP-stripe layout must be converted into the window-major layout required by Swin Transformer execution. SAWSTP performs this conversion through the DMP token route. To support different Transformer execution strategies, SAWSTP represents the Transformer stage using a factorized topology Topo(m, n, s, t), where (m, n) defines a two-dimensional WP grid for complete-window ownership, s denotes the SP degree for intra-window sequence partitioning, and t denotes the TP degree for sharding Transformer operators and intermediate activations. Under a selected topology, each Transformer rank is identified by a factorized coordinate (i, j, k, l), where (i, j) specifies the WP coordinate, k identifies the intra-window SP coordinate, and l identifies the TP coordinate. For clarity, Topo(m, n), Topo(m, n, s), and Topo(m, n, s, t) denote pure WP, WP+SP, and WP+SP+TP, respectively. Fig. 5 illustrates how the same DMP-stripe layout is transformed into the window-major layouts required by these Transformer topologies. Each complete attention window is represented by a 2 × 2 block with the same color, while the coordinate labels indicate the ranks responsible for the corresponding distributed window data. The gray boxes show the layouts directly produced by the DMP token route, and the layouts before the gray boxes conceptually show how complete windows are split into SP and TP sub-tiles. Under pure WP, each complete window is assigned to one WP rank, while the shift token route uses point-to-point (P2P) communication to redistribute tokens across WP ranks for shifted-window execution. When SP is enabled, each window is further partitioned along the intra-window sequence dimension, and Ulysses-based all-to-all communication reorganizes these sequence shards around window attention. The Topo(m, n, s, t) extends the WP+SP design of AERIS with FSDP and TP. AERIS-style WP+SP partitions window

activations but does not shard Transformer parameters. AERIS primarily uses PP to reduce parameter-related memory. However, pipeline bubbles are difficult to amortize under the limited global batch sizes of high-resolution Earth forecasting models such as Wenhai. SAWSTP instead applies FSDP over the DP × WP × SP process group. For Transformer blocks with a large hidden dimension d, FSDP must all-gather the full parameters of each wrapped unit before its execution, which can cause OOM failures. SAWSTP therefore introduces Megatron-style TP [10], [36] within each WP–SP group. As shown in Fig. 5, TP further partitions the computation assigned to each SP rank. It uses all-gather to materialize the activations required for a SP rank and reduce-scatter to restore the TP-sharded activation states afterward. Therefore, for a fixed world size, TP further shards Transformer parameters while preserving the per-rank activation layout of AERIS-style WP+SP partitioning. D. Differentiable Token Routes SAWSTP couples a DMP-based I/O topology with a multilevel Transformer topology that integrates WP, SP, and TP. These topologies assign tensor ownership, layouts, and communication groups differently, making it nontrivial to ensure correct token routing in the forward pass and gradient propagation in the backward pass. To address this issue, TERRA indexes both the DMP-stripe and window-major layouts using the same global two-dimensional token grid. The corresponding conversion from DMP-stripe layout to window-major layout is implemented as the DMP token route. After down-sampling, the Transformer-stage token grid has spatial size HT × WT , where HT =

Ĥ , pα

WT =

Ŵ . pα

(2)

The corresponding attention-window grid contains Mh × Mw windows, where Mh = HT /w and Mw = WT /w. Each

7

Fixed

Ragged

Density-aware

SP=2

Rank 0

Rank 1

Rank 2

Rank 3

Padding

Fig. 6. Fixed, ragged, and density-aware window assignment.

Transformer-stage token on the token grid is identified by a global (u, v). Its attention-window coordinate is   u  coordinate , wv ), and its within-window offset is determined by ( w (u mod w, v mod w). Under the DMP-stripe layout, (u, v) identifies the source rank and local offset of a token in its spatial stripe. Given the selected Topo(m, n, s, t) topology and window-assignment policy, (u, v) determines the owner of its complete attention window, its destination rank, and its local offset in the windowmajor tensor. Tokens with the same destination rank are gathered using torch.index_select in PyTorch [40] and packed into contiguous send buffers. TERRA then exchanges the peer-wise buffers through batched P2P communication. Each destination rank restores received tokens at their destination positions to construct the required window-major tensor. The shift token route uses the same coordinate-based mechanism. Window shifting changes the global coordinate (u, v) of each token, from which TERRA determines its new destination rank and local offset. The corresponding tokens are gathered, communicated, and restored using the same indexed packing procedure. Both routes are implemented as autograd-aware communication primitives. Their backward passes apply the inverse mappings to return gradients to the corresponding source ranks and local offsets. This design allows TERRA to support different Transformer topologies and window-assignment policies and provides the basis for the window-assignment optimizations presented in Section V. V. C OST-AWARE W INDOW A SSIGNMENT A. Ragged Window Assignment The multi-level Transformer topology maps complete attention windows to ranks along the WP dimension. As shown in the left part of Fig. 6, the AERIS-style fixed assignment uses two-dimensional round-robin mapping to distribute the windows of an Mh ×Mw grid across an m×n WP mesh. The illustrated example uses Mh = 9, Mw = 5, and m = n = 2. Fixed assignment provides uniform window ownership only when Mh and Mw are divisible by m and n, respectively. Otherwise, the input spatial dimensions are padded to Ĥ × Ŵ , where     H W Ĥ = pαwm, Ŵ = pαwn. (3) pαwm pαwn

This ensures that every WP rank receives the same number of complete windows. However, the additional padded tokens increase activation memory and computation. Moreover, twodimensional round-robin ownership requires distributing attention windows to adjacent WP ranks, which brings substantial P2P communication during shifted-window execution. To address this issue, TERRA instead adopts a ragged Topo(mn, 1) assignment. As shown in the middle part of Fig. 6, it avoids padding introduced for equal window ownership and preserves row-major contiguous ownership. Let M = Mh Mw be the total number of complete windows. After flattening the window grid in row-major order, the M complete windows are indexed from 0 to M − 1. WP rank r, where 0 ≤ r < mn, owns the contiguous window-index range [Lr , Rr ), where     (r + 1)M rM , Rr = . (4) Lr = mn mn The ragged assignment is enabled by the DMP token route, which allows each WP rank to receive tokens from the DMP ranks for a different number of complete windows. Thus, every complete window is assigned exactly once, and the window counts of any two ranks differ by at most one. Moreover, only boundary windows whose shifted-window token fragments cross a WP-rank boundary require P2P communication, whereas windows in the interior of each rank’s contiguous range remain local. This substantially reduces the communication volume of shifted-window execution. B. Density-Aware Window Assignment The ragged Topo(mn, 1) assignment is effective when the window grid is dense enough to provide sufficient assignment granularity. As the window size w increases, the number of complete windows decreases, and each complete window becomes a coarse assignment unit. The resulting ragged WP assignment can therefore become imbalanced. As shown in the right part of Fig. 6, assigning ten windows to four ranks with ragged Topo(4, 1) gives two ranks three windows each and the other two ranks two windows each. When the WP group size exceeds the number of complete windows, pureWP assignment can even leave some ranks with no assigned window. SP can mitigate this load imbalance by partitioning each complete attention window across SP ranks. For example, setting s = 2 changes the topology to Topo(2, 1, 2), and the resulting twenty window sub-tiles allow each of the four ranks to process five sub-tiles. However, SP introduces allto-all communication around window attention. This tradeoff motivates TERRA’s density-aware assignment, which extends the ragged WP assignment with SP when the performance benefit of improved assignment granularity covers the additional communication cost. Table III summarizes the communication characteristics of the fixed baseline and the two TERRA assignments. AERISstyle fixed assignment incurs P2P communication during window shifting. Ragged assignment limits this overhead to rankboundary windows. Density-aware assignment preserves this

8

TABLE III C OMMUNICATION CHARACTERISTICS OF WINDOW ASSIGNMENTS . Strategy Communication Fixed Topo(m, n) Two-dimensional WP P2P Ragged Topo(mn, 1) Boundary-only P2P Density-aware Topo(mn/s, 1, s) Boundary-only P2P + SP all-to-all

Algorithm 1 Density-aware window assignment 1: Input: rank count mn for WP and SP 2: Output: selected topology 3: bestTime ← +∞ 4: for each valid divisor s of mn do 5: candidateTopology ← Topo(mn/s, 1, s) 6: candidateTime ← Profile(candidateTopology) 7: if candidateTime < bestTime then 8: bestTime ← candidateTime 9: bestTopology ← candidateTopology 10: end if 11: end for 12: return bestTopology

reduced P2P communication while additionally introducing SP all-to-all communication. Algorithm 1 describes an offline profiling-based strategy for density-aware window assignment. Given mn ranks assigned to the WP and SP dimensions, TERRA enumerates the valid divisors s of mn. For each s, it constructs the corresponding Topo(mn/s, 1, s) candidate, profiles the candidate for several iterations, and selects the topology with the lowest average training-step time. The selected topology is then used throughout the subsequent training steps. For high window density, s = 1 is typically selected because ragged WP already provides sufficient load balance without introducing all-to-all communication. For low window density, a larger s is selected only when its load-balancing benefit outweighs the additional SP communication overhead. VI. M EMORY O RCHESTRATION A. Activation Bottleneck in Finetuning SAWSTP reduces the per-rank memory footprint of hierarchical Swin Transformer training. However, long-lead finetuning remains dominated by activation memory, particularly with small patch sizes. Fig. 7 shows the allocated GPU-memory timelines during the forward (FWD) and backward (BWD) passes of ten-lead rollout finetuning. During FWD, checkpoint boundaries generated at each lead remain resident until their corresponding BWD computations and therefore accumulate across the rollout. During BWD, recomputation of a sampling or Transformer checkpoint unit produces a short-lived local peak. Although numerous memory-optimization techniques based on activation recomputation [19], [20], [21], [22] and activation offloading [23], [24], [25], [26] have been proposed, it remains nontrivial to jointly determine recomputation granularity and offloading decisions for hierarchical rollout finetuning under SAWSTP. Sampling and Transformer checkpoint

units exhibit different checkpoint-boundary tensor sizes and recomputation peaks, while high-resolution input and label tensors further introduce a non-negligible baseline memory footprint. These characteristics require an architecture-aware policy that models rollout-wide tensor lifetimes and satisfies a target GPU-memory budget. TERRA addresses this challenge through Memory Orchestration (MO), a profile-based two-stage policy. MO first constructs an architecture-aware checkpoint-only plan by modeling rollout-wide boundary accumulation and local recomputation peaks. When checkpointing alone cannot satisfy a target GPU-memory budget, MO combines input buffering with selective boundary offloading to meet the budget.

B. Profile-Based Checkpoint Planning The first stage of MO constructs a checkpoint-only policy from a block-wise runtime memory profile under the selected SAWSTP configuration. Inspired by the runtime memory tracer in PatrickStar [31], TERRA records the profile at runtime and visualizes the allocated GPU-memory timeline. It uses torch.cuda.memory_allocated for the allocated-memory timeline and torch.cuda.max_memory_allocated for peak memory. As shown in Fig. 7, the p = 2, DMP = 4, and DP = 2 configuration exhibits distinct forward and backward memory patterns under different checkpointing policies. MO distinguishes sampling and Transformer checkpoint units because they exhibit different boundary activation sizes and recomputation peaks. As shown in the lower-left part of Fig. 4, each sampling module contains several Conv-based residual blocks. For sampling modules, MO considers two checkpoint granularities. At the coarse granularity, D0 treats patch embedding and the subsequent down-sampling module as one checkpoint unit, whereas U0 combines the up-sampling module with patch recovery. At the fine granularity, D1 combines patch embedding with the first stride-two convolution and checkpoints each subsequent residual block independently. U1 applies the analogous fine-grained partitioning to the upsampling module and patch recovery. The Transformer stage contains L = 10 Swin Transformer blocks. MO profiles checkpoint units containing tseg ∈ {1, 2, 3, 4, 5, 10} consecutive blocks. Coarser units retain fewer checkpoint boundaries across the rollout but produce larger peak memory during recomputation, whereas finer units have the opposite tradeoff. Figs. 7 (a) and 7(b) illustrate this tradeoff for D0 /U0 with tseg = 1 and tseg = 3, respectively. In particular, the coarse D0 /U0 units can substantially reduce boundary memory because their internal boundaries correspond to high-resolution feature maps. For each checkpoint unit i, MO profiles its boundary memory Ei and backward peak Pi . A checkpoint-only policy π selects the sampling granularity and Transformer segment length at each rollout lead. MO predicts the resulting peak memory as M (π) = Mbase + max {Ai (π) + Pi (π)} , i

(5)

$OORFDWHG*38PHPRU\*%

9

*% 

*%



*% *%

  

):'

%:'

(a) D0 , U0 , tseg = 1.

):'

%:'

):'

(b) D0 , U0 , tseg = 3.

%:'

(c) Checkpoint planning.

):'

%:'

(d) Budget-constrained offloading.

Fig. 7. Allocated-memory timelines for ten-lead finetuning under different policies.

where Ai (π) denotes the cumulative memory of checkpoint boundaries that remain resident immediately before the backward execution of unit i. The model accumulates boundaries that coexist across rollout leads and takes the maximum over the transient peaks created by recomputation. Mbase denotes the profiled fixed memory footprint, including input and label tensors and other persistent runtime allocations. MO minimizes Eq. (5) using a dynamic-programming method inspired by [20]. As illustrated in Fig. 7(c), the dynamicprogramming method selects coarser checkpoint units for early leads and finer-grained units for later leads in this configuration, reducing retained boundary memory and recomputation peaks, respectively. C. Budget-Constrained Offloading When the checkpoint-only policy π cannot satisfy a target memory budget Mb , MO constructs an offloading-based policy πo using the profiled boundary memories Ei and backward peaks Pi . As shown in Algorithm 2, MO employs a lightweight profile-guided greedy algorithm to satisfy the memory constraint without solving a globally optimal tensorlevel offloading schedule. MO first considers D0 /U0 and selects the largest tseg whose profile-predicted peak memory does not exceed Mb . The selected sampling granularity and Transformer segment template are then applied to all rollout leads. MO falls back to D1 /U1 only when no D0 /U0 configuration satisfies the memory constraint. MO then constructs πo over the K rollout leads. For each lead k, it appends the selected configuration to πo and inserts its down-sampling, Transformer, and up-sampling boundaries into a queue Q in forward creation order. Mk (πo ) denotes the profile-predicted peak obtained by replaying the backward events of lead k under the current offloading policy. It includes resident boundary memory, backward peak increments, and the boundary memory restored for offloaded checkpoint units. If Mk (πo ) exceeds Mb , MO removes the earliest boundary from the queue and marks it for offloading. Earlier boundaries normally have longer forward-to-backward lifetimes, making them suitable candidates for asynchronous offloading. The procedure continues until the memory constraint is satisfied. Fig. 7(d) illustrates the resulting offloading-based memory timeline under a 40 GB budget. Compared with the checkpoint-only policy in Fig. 7(c), boundary offloading reduces the measured peak from 60.3 GB to 39.1 GB, which satisfies the memory constraint.

Algorithm 2 Profile-guided boundary activation offloading 1: Input: profiled Ei , Pi , and Mbase , rollout length K, and memory budget Mb 2: Output: offloading-based policy πo 3: A configuration (D, tseg , U ) is feasible if Mbase + Ei + Pi ≤ Mb for every checkpoint unit i 4: Select the coarsest sampling granularity and the largest feasible tseg for (D, tseg , U ) 5: πo ← ∅ and Q ← ∅ 6: for k = 1 to K do 7: Append (D, tseg , U ) for lead k to πo 8: Append its checkpoint boundaries to Q in forward creation order 9: Compute the predicted peak Mk (πo ) 10: while Mk (πo ) > Mb do 11: Remove the earliest boundary i from Q 12: Mark boundary i for offloading in πo 13: Recompute the predicted peak Mk (πo ) 14: end while 15: end for 16: return πo

In the implementation, TERRA applies a boundary-offload wrapper to each checkpoint unit selected by MO. The wrapper extends the torch.utils.checkpoint interface with torch.autograd.graph.saved_tensors_hooks, which capture the input tensors saved at each checkpoint boundary. The pack hook offloads checkpoint boundaries to pinned host buffers through asynchronous device-to-host (D2H) transfers, while the unpack hook restores them through host-to-device (H2D) transfers before backward recomputation. This design requires minimal changes to the training code and is compatible with both FSDP and SAWSTP. D. Input Buffering Even with activation checkpointing and offloading, highresolution input and label data consume substantial GPU memory and constitute a major component of Mbase . This cost is particularly large at small DMP degrees, where each rank owns a larger spatial shard. For Wenhai, one FP16 field of shape 93×2041×4320 occupies approximately 1.53 GB. A ten-lead rollout requires one input field and ten label fields for loss computation. Therefore, simultaneously transferring all

10

TABLE IV E XPERIMENTAL PLATFORM .

Serial TERRA WP TERRA WP+SP TERRA WP+TP TERRA WP+SP+TP

Configuration Up to 12 nodes and 96 GPUs 8× NVIDIA H200 GPUs per node 2× Intel Xeon Platinum 8558, 96 cores per node 1.8 TB per node NVLink Mellanox RDMA network

0.15 0.10 0.05

M ODEL CONFIGURATIONS . Model Small Medium Large

L 10 10 10

d Nh 1024 8 4096 32 8192 64

p 4 4 4

α 2 2 2

w Number of Parameters 8 0.18B 8 2.85B 32 11.4B

these fields via H2D incurs both substantial transfer overhead and a large GPU-memory footprint. To address this issue, TERRA employs a lead-wise input and label prefetching pipeline with pinned host memory and non-blocking H2D transfers. Rather than transferring the entire rollout sequence to the GPU at once, the pipeline transfers tensors on demand and prefetches data for upcoming leads during rollout execution. Its input-prefetch stream is separate from the activation-offload stream, allowing H2D prefetching and checkpoint-boundary D2H offloading to progress independently when resources permit. Together, the input buffering method bounds the GPU-side lifetime of input and label tensors and hides part of the H2D latency behind rollout computation, thereby reducing GPU-memory consumption and end-to-end step time. VII. E VALUATION A. Experiment Setup 1) Hardware and Software: We conduct all experiments on a homogeneous cluster comprising up to 12 nodes and 96 NVIDIA H200 GPUs. Each GPU provides 141 GB of HBM3e memory and a peak dense FP16 throughput of 989 TFLOPS. Table IV summarizes the hardware environment. The software stack comprises PyTorch 2.4.0, CUDA 12.4, NCCL 2.20.5, FlashAttention 2.7.3 [41], and NVIDIA driver 570.133.20. 2) Workload and Models: We evaluate TERRA on the GLORYS-based Wenhai workload, in which each training sample contains 93 physical variables on a 2041 × 4320 global ocean grid. Table V summarizes the small, medium, and large hierarchical Wenhai model configurations used in our experiments. Models are trained using FP16 mixed precision and mean absolute error (MAE) loss. Unless otherwise specified, performance measurements use a micro-batch size of B = 1, while end-to-end training and forecasting-accuracy experiments use a global batch size of 16. Transformer checkpoint units use tseg = 1, and the down- and up-sampling modules use D1 /U1 checkpointing. For the small model, we apply a single FSDP wrapper to the entire model. For the medium and large models, we wrap each Transformer block and sampling module independently with FSDP. FlashAttention is enabled for configurations with window size w ≥ 32.

0.010

0.005

0

20

40

60

80

Training step

(a) Steps 0–100. TABLE V

0.015

Loss

Component Cluster scale GPU CPU Host memory Intra-node interconnect Inter-node interconnect

Loss

0.20

100

100 200

400

600

800

1000

Training step

(b) Steps 100–1000.

Fig. 8. Training-loss consistency between serial execution and TERRA parallel configurations.

3) Baselines and Metrics: We evaluate the scalability of TERRA through strong- and weak-scaling experiments on hierarchical Wenhai models. Unless otherwise specified, the reported step time includes the forward, backward, and optimizer-update phases but excludes the loading and H2D transfer of input and label tensors. For parameter-capacity scaling, we compare an AERIS-style WP+FSDP baseline (TP = 1) with TERRA configurations that combine FSDP and TP. For window assignment, we compare TERRA’s ragged Topo(mn, 1) and density-aware Topo(mn/s, 1, s) assignments with fixed AERIS-style WP and WP+SP topologies. For MO, we evaluate its effectiveness in reducing GPU memory consumption during rollout finetuning under different SAWSTP configurations. For forecasting accuracy, we compare models with different patch sizes and sampling configurations, as well as models finetuned with different rollout lengths. These comparisons quantify the forecasting benefits enabled by TERRA through finer patch granularity, hierarchical sampling, and longer rollout finetuning. We report the mean root mean square error (RMSE) over all predicted variables, where lower values indicate more accurate forecasts. B. Convergence Validation We first validate the correctness of TERRA’s parallel training framework. Fig. 8 compares the training-loss curve obtained from serial execution of the small model in Table V with the corresponding curves obtained from TERRA’s parallel configurations. The serial execution uses FSDP alone and does not employ SAWSTP. For a fair comparison, all runs use the same padded input resolution of 2304 × 4352 and DP = 2. The TERRA configurations use eight ranks per SAWSTP group and vary the Transformer topology across WP, SP, and TP. The WP, WP+SP, WP+TP, and WP+SP+TP configurations use ragged Topo(8, 1, 1, 1), Topo(4, 1, 2, 1), Topo(4, 1, 1, 2), and Topo(2, 1, 2, 2), respectively. As shown in Fig. 8, all TERRA configurations exhibit convergence behavior comparable to that of serial execution over the first 1000 training steps. This agreement validates gradient propagation through parallel sampling, token routing, and the evaluated combinations of WP, SP, and TP. C. Scalability We evaluate the strong- and weak-scaling performance of TERRA using the Wenhai models in Table V. To avoid OOM

11



6WHSWLPHV

VPDOO% PHGLXP% ODUJH%

 



 

 















1XPEHURI*38V

 











3HDNPHPRU\*%



3DUDOOHOHIILFLHQF\





 

 









(a) Strong scaling step time and efficiency.











VPDOO% PHGLXP% ODUJH%















1XPEHURI*38V







(b) Peak memory under strong scaling.

Fig. 9. Strong scaling performance.



VPDOO% PHGLXP% ODUJH%

6XVWDLQHGSHUIRUPDQFH3)/236

7KURXJKSXWVDPSOHVV













1XPEHURI*38V





(a) Weak scaling throughput.



VPDOO% PHGLXP% ODUJH%

 







 





 







1XPEHURI*38V





(b) Weak scaling sustained performance.

Fig. 10. Weak scaling performance.

errors, measurements for the large model start at 16 GPUs. For strong scaling, we fix DP = 4 and increase the DMP degree with the number of GPUs. For weak scaling, we fix DMP = 8 for the small and medium models and DMP = 16 for the large model, while increasing the DP degree with the number of GPUs. For all configurations, we set s = 1 and use the ragged Topo(mn, 1) assignment. Fig. 9 shows the strong-scaling results. Increasing the DMP degree reduces peak allocated GPU memory for all three models and consistently decreases the step time of the medium and large models. These results show that SAWSTP distributes the computation and activations of both sampling modules and Transformer blocks across multiple GPUs. The large model achieves a 3.9× speedup from 16 to 96 GPUs, corresponding to 65.0% strong-scaling efficiency. In contrast, the strong scalability of the small model is limited because the DMP stripe and the number of attention windows assigned to each rank become too small at high GPU counts, while communication and synchronization overheads become increasingly significant. Increasing the micro-batch size could provide more per-rank computation to amortize these overheads and thereby improve throughput. Fig. 9(b) further shows that activation and parameterrelated memory scale differently. As the SAWSTP group

size increases, DMP distributes high-resolution sampling activations, while WP distributes Transformer activations. The larger FSDP sharding group also reduces per-rank parameterrelated memory. Consequently, the peak memory of the small and medium models continues to decrease as more GPUs are used. However, peak memory remains around 40 GB at 64 and 96 GPUs for the large model. In this case, fullparameter materialization during FSDP all-gather accounts for an increasing fraction of the per-rank peak. This observation motivates combining TP with FSDP to scale larger hierarchical Swin Transformer models. Fig. 10 shows the weak-scaling results. The medium model increases aggregate throughput from 0.76 to 8.5 samples/s between 8 and 96 GPUs and achieves 39.76 PFLOPS, corresponding to 93.0% weak-scaling efficiency. The large model scales from 0.34 samples/s on 16 GPUs to 1.93 samples/s on 96 GPUs, achieving 94.1% weak-scaling efficiency. The small model increases throughput from 6.4 to 56.1 samples/s but reaches a lower efficiency of 73.0% at 96 GPUs because its smaller per-rank workload makes FSDP and SAWSTP communication more prominent. Fig. 11 evaluates how TP complements WP and FSDP for parameter scaling. To focus the evaluation on parameter scaling in the Transformer blocks, we use a Swin Transformer

12

TERRA (TP=2)

TERRA (TP=4)

OOM

400

Step time (ms)

Maximum Trainable Parameters (B)

AERIS (TP=1) 60 50 40 30 20

300

s=4

TERRA Ragged Topo(mn, 1) TERRA Density-aware Topo(mn/s, 1, s)

200 100 0

10 0

AERIS WP Topo(2, 4) AERIS WP+SP Topo(2, 2, 2)

8

32

64

128

Window size 16

32

Number of GPUs

64

(a) Training-step time on 16 GPUs.

Fig. 11. Largest trainable model size for the AERIS-style baseline and TERRA configurations with TP.

D. Window Assignment Performance Fig. 12 compares the window-assignment policies for the small model on 16 and 32 GPUs. Both configurations use DP = 2, with DMP = 8 on 16 GPUs and DMP = 16 on 32 GPUs. The AERIS-style baselines use fixed WP and WP+SP topologies. For TERRA, we report the ragged Topo(mn, 1) topology without SP and the fastest measured policy among density-aware Topo(mn/s, 1, s) candidates for each window size. Overall, TERRA achieves a 1.03×–1.18× speedup over the faster AERIS-style baseline. For dense window grids (w ≤ 64), TERRA selects the ragged Topo(8, 1) topology on 16 GPUs and Topo(16, 1) on 32 GPUs. In this regime, completewindow WP provides sufficient parallelism, while introducing SP incurs unnecessary intra-window all-to-all communication. The ragged assignments also avoid padding introduced solely to maintain uniform window ownership when the window grid is incompatible with a fixed topology. Consequently, TERRA

Step time (ms)

without sampling modules and fix L = 20, p = 4, and w = 8. We set DP = 2 and decompose the remaining ranks into WP and TP dimensions. The AERIS-style baseline retains AERIS’s window-parallel execution and uses FSDP in the same runtime with TP = 1. The TERRA configurations use the same setup while additionally enabling TP = 2 or 4. For each configuration, we increase the hidden dimension d until the next candidate model runs out of memory and report the largest tested trainable model. At 16 GPUs, the AERIS-style TP = 1 baseline supports 25.2B parameters, whereas the TERRA configurations with TP = 2 and TP = 4 support 16.1B parameters. At this scale, increasing the TP degree reduces the WP degree and causes each WP rank to process more windows. The activation allgather operations required by TP further limit the memory savings from parameter partitioning at this scale. At 32 GPUs, the TERRA configurations with TP = 2 and TP = 4 increase the largest tested model size from 25.2B to 39.3B and 49.3B parameters, respectively. At 64 GPUs, both TERRA configurations support 49.3B parameters, while the AERISstyle baseline remains limited to 25.2B. By combining TP with FSDP, SAWSTP supports up to 2.0× as many parameters as the TP = 1 baseline.

AERIS WP Topo(4, 4)

400 300

AERIS WP+SP Topo(2, 2, 4)

TERRA Ragged Topo(mn, 1) TERRA Density-aware Topo(mn/s, 1, s)

s=8 200 100 0

8

32

64

128

Window size

(b) Training-step time on 32 GPUs. Fig. 12. Training-step time under fixed AERIS-style topologies and TERRA’s ragged window assignments.

achieves 1.03×–1.14× speedups over the faster fixed baseline on 16 GPUs and 1.04×–1.10× speedups on 32 GPUs. For the sparse-window case at w = 128, TERRA selects the SP-enabled ragged topologies Topo(2, 1, 4) and Topo(2, 1, 8) on 16 and 32 GPUs, respectively. These policies reduce the step time to 306.4 ms and 200.0 ms, corresponding to 1.18× and 1.14× speedups over the faster fixed baseline. The worst case occurs for fixed WP on 32 GPUs. To maintain uniform window ownership, the fixed Topo(4, 4) baseline pads the minimally padded 2 × 5 window grid to a 4 × 8 grid containing 32 windows. This increases the padding overhead to 280.6%, compared with 18.9% under the selected ragged Topo(2, 1, 8) policy. Therefore, the selected policy achieves 2.36× and 1.14× speedups over the fixed Topo(4, 4) WPonly and Topo(2, 2, 4) WP+SP baselines, respectively. E. Memory Orchestration Performance Fig. 13 evaluates the checkpoint-planning stage of MO for ten-lead finetuning on a single H200 node, with activation offloading and input buffering disabled. We use the small model and evaluate two SAWSTP configurations for each patch size. The p = 2 experiments use DMP = 4 with DP = 2 and DMP = 8 with DP = 1, whereas the p = 4 experiments use DMP = 1 with DP = 8 and DMP = 2 with DP = 4. We compare the resulting lead-dependent plan with uniform policies that use D0 /U0 and a fixed tseg ∈ {1, 3, 5} for all leads. Across the four SAWSTP configurations, checkpoint planning reduces peak allocated GPU memory by 7.6%– 13.0% relative to the best feasible uniform policy. For p = 4, DMP = 1, and DP = 8, the uniform tseg = 1 and tseg = 3 policies run out of memory, while tseg = 5 reaches 124.0 GB.

13

Peak allocated GPU memory (GB)

tseg = 1

tseg = 5

tseg = 3

Checkpoint Planning

TABLE VI B REAKDOWN OF PER -GPU ACTIVATION OFFLOAD VOLUME PER TRAINING STEP.

OOM

140 124.0

120

114.0 105.2

100

Patch/DMP/DP

Mb (GB)

2/4/2

40 30 20 16 80 65 40 30

79.9

80

69.3 69.9 62.0 62.6

60.3

60

52.9

40

35.0 35.5

57.3

30.5

2/8/1

20 0 Patch2 DMP4 / DP2

Patch2 DMP8 / DP1

Patch4 DMP1 / DP8

4/1/8

Patch4 DMP2 / DP4

4/2/4 Fig. 13. Peak allocated GPU memory of checkpoint-only policies for ten-lead finetuning.

Checkpoint planning Input buffer

Budget offload 1 Budget offload 2

Patch2 / DMP4 / DP2 13.28s

11.26s

Peak allocated GPU memory (GB)

60

60.3

11.15s

30 39.1

40

Patch2 / DMP8 / DP1 6.80s 5.67s

11.17s

56.4

29.1

20

Step time OOM

30.5

5.61s

28.5 18.8

20

14.8

10

0

0 40 GB 30 GB

20 GB 16 GB

Patch4 / DMP1 / DP8

11.84s

11.35s

125 100 75

5.61s

11.37s

Patch4 / DMP2 / DP4 8.64s

11.80s

114.0

6.08s

60

5.84s

5.91s

57.3 49.5

98.4 74.9 61.1

50

40

38.9 30.5

20

25 0

0 80 GB 65 GB

40 GB 30 GB

Fig. 14. Peak allocated GPU memory and step time for input buffering and budget-constrained offloading.

Checkpoint planning remains feasible and reduces the peak to 114.0 GB. These results show that a uniform checkpoint granularity cannot consistently balance rollout-wide boundary accumulation against local recomputation peaks. Fig. 14 evaluates input buffering and budget-constrained boundary offloading using the same four ten-lead configurations as Fig. 13. The reported step time includes input and label H2D transfers as well as activation D2H and H2D transfers when offloading is enabled. During rollout execution, TERRA asynchronously prefetches the tensors required by the next lead and releases tensors that are no longer needed as execution advances. Across the four configurations, input buffering reduces peak allocated GPU memory by 6.4%– 13.7% and step time by 1.0%–4.1%. These results show that bounding the GPU-side lifetime of input and label tensors reduces memory consumption while overlapping part of their H2D transfers with computation. MO enables activation offloading when the checkpointonly policy exceeds the target budget Mb . Overall, relative to

Offload Volume (GB) Sampling Transformer 9.7 16.7 72.5 29.9 4.3 9.9 36.2 14.9 22.8 10.8 22.8 27.5 10.5 5.4 42.8 14.9

the corresponding checkpoint-only policies, the combination of input buffering and activation offloading reduces peak allocated GPU memory by 32.2%–51.8%, with at most 20.0% step-time overhead across the budget-feasible configurations. Under the primary budgets of 40, 20, 80, and 40 GB, all four configurations satisfy their targets, reducing peak memory by 32.2%–38.2%. When the budgets are tightened to 30, 16, 65, and 30 GB, the first three configurations remain feasible and reduce peak memory by 46.4%–51.8%, with at most 20.0% step-time overhead. The p = 4, DMP = 2, and DP = 4 configuration reaches a measured peak of 30.5 GB and therefore slightly exceeds its 30 GB budget. This difference results from a small prediction error in the profile-based memory model. In practice, OOM can be avoided by setting a slightly lower planning budget. Table VI provides a stage-wise breakdown of the per-GPU activation offload volume per training step. When a tighter budget cannot be satisfied using the coarse D0 /U0 units, MO switches to D1 /U1 , exposing additional sampling boundaries for offloading. For example, for p = 2, DMP = 4, and DP = 2, tightening Mb from 40 GB to 30 GB increases the total offload volume from 26.4 GB to 102.4 GB per GPU per step. This increase is dominated by sampling boundaries, whose offload volume grows from 9.7 GB to 72.5 GB, while the Transformer-boundary volume grows from 16.7 GB to 29.9 GB. The resulting increase in data transfer is consistent with the higher step times in Fig. 14. F. Forecasting Accuracy We finally evaluate the forecasting accuracy gains enabled by TERRA through hierarchical sampling, smaller patch sizes, and longer rollout finetuning. We train, finetune, and evaluate the small Wenhai model in Table V. We use GLORYS data from 1993 to 2018 for pretraining and finetuning, and data from 2020 for validation. We pretrain each model for 200 epochs and finetune each rollout variant for 5000 steps. We report the mean RMSE over all ocean variables at each forecast lead time, where lower values indicate more accurate forecasts. Fig. 15(a) compares pretrained models with different patch and sampling configurations. To match the number of tokens processed by the Transformer blocks, we use p = 8 for the model without hierarchical sampling and p = 4 with a downsampling stride of α = 2 for the hierarchical model. The p = 2 model with hierarchical sampling achieves the lowest

0.180

p = 4, with sampling p = 2, with sampling p = 8, no sampling

0.200 0.175

Mean RMSE across variables

Mean RMSE across variables

14

0.150 0.125 0.100 0.075 0.050 0.025 1

2

3

4

5

6

7

8

9

10

lead=1 lead=2 lead=4

0.160

lead=7 lead=10

0.140 0.120 0.100 0.080 0.060 0.040 1

Forecast lead time (day)

2

3

4

5

6

7

8

9

10

Forecast lead time (day)

(a) Pretraining with different patch and sampling configurations.

(b) Rollout finetuning with different rollout lengths.

Fig. 15. Mean RMSE of Wenhai forecasts at different lead times. Date: 2020-01-11 Ground Truth

p=2 With Sampling

p=4 With Sampling

p=8 No Sampling

20

SST/℃

25

15

10

Fig. 16. Day-10 SST forecasts over the western North Pacific.

RMSE at every evaluated forecast lead. With the sampling configuration fixed, reducing the patch size from p = 4 to p = 2 decreases the day-10 mean RMSE by 2.1%. Together, these results demonstrate the forecasting benefits of finer patch granularity and hierarchical sampling, highlighting the importance of TERRA’s support for the parallel execution of sampling modules. Fig. 15(b) fixes the p = 4 model with hierarchical sampling and compares the pretrained model with variants finetuned using different rollout lengths. Increasing the rollout length consistently improves long-lead accuracy. In particular, tenlead finetuning reduces the day-10 mean RMSE from 0.181 for the pretrained model to 0.156, corresponding to a 13.5% reduction. The improvement becomes smaller beyond seven leads. Overall, these results demonstrate that the longer rollout finetuning enabled by TERRA improves long-range forecasting accuracy. Fig. 16 presents ten-day sea surface temperature (SST) forecasts initialized from the GLORYS global ocean state on January 1, 2020. It compares the forecasts valid on January 11, 2020, from the pretrained configurations in Fig. 15(a) with the corresponding GLORYS ground truth. The model without hierarchical sampling produces the most spatially smoothed forecast, with a substantial loss of fine-scale SST structures. With hierarchical sampling, the p = 2 model resolves finer and clearer spatial details than the p = 4 model, which is consistent with the lower day-10 mean RMSE achieved by the p = 2 model.

VIII. C ONCLUSION We present TERRA in this paper, a hierarchical parallel training and memory orchestration framework for highresolution AI-based Earth forecasting models. TERRA’s SAWSTP parallelizes convolutional sampling modules and Swin Transformer blocks and connects them through differentiable token routes. Cost-aware window assignment reduces padding and shifted-window communication, while MO reduces the memory footprint of long-lead rollout finetuning. Experiments on Wenhai show that TERRA scales the 11.4B model to 96 H200 GPUs with 65.0% strong-scaling efficiency. Compared with checkpoint-only policies, MO reduces peak allocated GPU memory by 32.2%–51.8% with at most 20.0% steptime overhead. These capabilities make smaller patch sizes and longer rollout finetuning feasible, thereby improving long-lead forecasting accuracy. R EFERENCES [1] T. Kurth, S. Subramanian, P. Harrington, J. Pathak, M. Mardani, D. Hall, A. Miele, K. Kashinath, and A. Anandkumar, “Fourcastnet: Accelerating global high-resolution weather forecasting using adaptive fourier neural operators,” in Proceedings of the platform for advanced scientific computing conference, 2023, pp. 1–11. [2] K. Bi, L. Xie, H. Zhang, X. Chen, X. Gu, and Q. Tian, “Accurate medium-range global weather forecasting with 3d neural networks,” Nature, vol. 619, no. 7970, pp. 533–538, 2023. [3] L. Chen, X. Zhong, F. Zhang, Y. Cheng, Y. Xu, Y. Qi, and H. Li, “Fuxi: A cascade machine learning forecasting system for 15-day global weather forecast,” npj climate and atmospheric science, vol. 6, no. 1, p. 190, 2023.

15

[4] R. Lam, A. Sanchez-Gonzalez, M. Willson, P. Wirnsberger, M. Fortunato, F. Alet, S. Ravuri, T. Ewalds, Z. Eaton-Rosen, W. Hu et al., “Learning skillful medium-range global weather forecasting,” Science, vol. 382, no. 6677, pp. 1416–1421, 2023. [5] I. Price, A. Sanchez-Gonzalez, F. Alet, T. R. Andersson, A. El-Kadi, D. Masters, T. Ewalds, J. Stott, S. Mohamed, P. Battaglia et al., “Probabilistic weather forecasting with machine learning,” Nature, vol. 637, no. 8044, pp. 84–90, 2025. [6] Y. Cui, R. Wu, X. Zhang, Z. Zhu, B. Liu, J. Shi, J. Chen, H. Liu, S. Zhou, L. Su et al., “Forecasting the eddying ocean with a deep neural network,” Nature Communications, vol. 16, no. 1, p. 2268, 2025. [7] C. Bodnar, W. P. Bruinsma, A. Lucic, M. Stanley, A. Allen, J. Brandstetter, P. Garvan, M. Riechert, J. A. Weyn, H. Dong et al., “A foundation model for the earth system,” Nature, vol. 641, no. 8065, pp. 1180–1187, 2025. [8] X. Wang, S. Liu, A. Tsaris, J.-Y. Choi, A. M. Aji, M. Fan, W. Zhang, J. Yin, M. Ashfaq, D. Lu et al., “Orbit: Oak ridge base foundation model for earth system predictability,” in SC24: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2024, pp. 1–11. [9] Y. Zhao, A. Gu, R. Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer et al., “Pytorch fsdp: experiences on scaling fully sharded data parallel,” arXiv preprint arXiv:2304.11277, 2023. [10] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-lm: Training multi-billion parameter language models using model parallelism,” arXiv preprint arXiv:1909.08053, 2019. [11] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly et al., “An image is worth 16x16 words: Transformers for image recognition at scale,” arXiv preprint arXiv:2010.11929, 2020. [12] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo, “Swin transformer: Hierarchical vision transformer using shifted windows,” in Proceedings of the IEEE/CVF international conference on computer vision, 2021, pp. 10 012–10 022. [13] V. Hatanpää, E. Ku, J. Stock, M. Emani, S. Foreman, C. Jung, S. Madireddy, T. Nguyen, V. Sastry, R. A. Sinurat et al., “Aeris: Argonne earth systems model for reliable and skillful predictions,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2025, pp. 72–85. [14] S. A. Jacobs, M. Tanaka, C. Zhang, M. Zhang, S. L. Song, S. Rajbhandari, and Y. He, “Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models,” arXiv preprint arXiv:2309.14509, 2023. [15] Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu et al., “Gpipe: Efficient training of giant neural networks using pipeline parallelism,” Advances in neural information processing systems, vol. 32, 2019. [16] T. Nguyen, R. Shah, H. Bansal, T. Arcomano, R. Maulik, V. Kotamarthi, I. Foster, S. Madireddy, and A. Grover, “Scaling transformer neural networks for skillful and reliable medium-range weather forecasting,” Advances in Neural Information Processing Systems, vol. 37, pp. 68 740–68 771, 2024. [17] L. Jean-Michel, G. Eric, B.-B. Romain, G. Gilles, M. Angélique, D. Marie, B. Clément, H. Mathieu, L. G. Olivier, R. Charly et al., “The copernicus global 1/12 oceanic and sea ice glorys12 reanalysis,” Frontiers in Earth Science, vol. 9, p. 698876, 2021. [18] L. Hou, Y. Cheng, N. Shazeer, N. Parmar, Y. Li, P. Korfiatis, T. M. Drucker, D. J. Blezek, and X. Song, “High resolution medical image analysis with spatial partitioning,” arXiv preprint arXiv:1909.03108, 2019. [19] T. Chen, B. Xu, C. Zhang, and C. Guestrin, “Training deep nets with sublinear memory cost,” arXiv preprint arXiv:1604.06174, 2016. [20] A. Gruslys, R. Munos, I. Danihelka, M. Lanctot, and A. Graves, “Memory-efficient backpropagation through time,” Advances in neural information processing systems, vol. 29, 2016. [21] M. Kirisame, S. Lyubomirsky, A. Haan, J. Brennan, M. He, J. Roesch, T. Chen, and Z. Tatlock, “Dynamic tensor rematerialization,” arXiv preprint arXiv:2006.09616, 2020. [22] X. Zhao, T. Le Hellard, L. Eyraud-Dubois, J. Gusak, and O. Beaumont, “Rockmate: an efficient, fast, automatic and generic tool for rematerialization in pytorch,” in International Conference on Machine Learning. PMLR, 2023, pp. 42 018–42 045. [23] M. Rhu, N. Gimelshein, J. Clemons, A. Zulfiqar, and S. W. Keckler, “vdnn: Virtualized deep neural networks for scalable, memory-efficient neural network design,” in 2016 49th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO). IEEE, 2016, pp. 1–13.

[24] S. Shriram, A. Garg, and P. Kulkarni, “Dynamic memory management for gpu-based training of deep neural networks,” in 2019 IEEE International Parallel and Distributed Processing Symposium (IPDPS). IEEE, 2019, pp. 200–209. [25] O. Beaumont, L. Eyraud-Dubois, and A. Shilova, “Efficient combination of rematerialization and offloading for training dnns,” Advances in Neural Information Processing Systems, vol. 34, pp. 23 844–23 857, 2021. [26] Y. Tang, Q. Li, L. Yin, D. Li, Y. Zhang, C. Wang, X. Zhang, L. Qiao, Z. Zhang, and K. Lu, “Delta: Memory-efficient training via dynamic fine-grained recomputation and swapping,” ACM Transactions on Architecture and Code Optimization, vol. 21, no. 4, pp. 1–25, 2024. [27] H. Hersbach, B. Bell, P. Berrisford, S. Hirahara, A. Horányi, J. MuñozSabater, J. Nicolas, C. Peubey, R. Radu, D. Schepers et al., “The era5 global reanalysis,” Quarterly journal of the royal meteorological society, vol. 146, no. 730, pp. 1999–2049, 2020. [28] D. Kieckhefen, M. Götz, L. H. Heyen, A. Streit, and C. Debus, “Jigsaw: Training multi-billion-parameter ai weather models with optimized model parallelism,” arXiv preprint arXiv:2507.05753, 2025. [29] X. Wang, J.-Y. Choi, T. Kurihaya, I. Lyngaas, H.-J. Yoon, X. Xiao, D. Pugmire, M. Fan, N. M. Nafi, A. Tsaris et al., “Orbit-2: Scaling exascale vision foundation models for weather and climate downscaling,” in proceedings of the International Conference for high performance computing, networking, storage and analysis, 2025, pp. 86–98. [30] C. Adams, P. Harrington, A. Subramaniam, M. S. Abbas, J. Pathak, M. Pritchard, and S. Choudhry, “Shardtensor: Domain parallelism for scientific machine learning,” arXiv preprint arXiv:2605.11111, 2026. [31] J. Fang, Z. Zhu, S. Li, H. Su, Y. Yu, J. Zhou, and Y. You, “Parallel training of pre-trained models via chunk-based dynamic memory management,” IEEE Transactions on Parallel and Distributed Systems, vol. 34, no. 1, pp. 304–315, 2022. [32] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He, “Zero: Memory optimizations toward training trillion parameter models,” in SC20: international conference for high performance computing, networking, storage and analysis. IEEE, 2020, pp. 1–16. [33] J. Ren, S. Rajbhandari, R. Y. Aminabadi, O. Ruwase, S. Yang, M. Zhang, D. Li, and Y. He, “{Zero-offload}: Democratizing {billion-scale} model training,” in 2021 USENIX Annual Technical Conference (USENIX ATC 21), 2021, pp. 551–564. [34] S. Rajbhandari, O. Ruwase, J. Rasley, S. Smith, and Y. He, “Zeroinfinity: Breaking the gpu memory wall for extreme scale deep learning,” in Proceedings of the international conference for high performance computing, networking, storage and analysis, 2021, pp. 1–14. [35] X. Lian, M. Tanaka, O. Ruwase, and M. Zhang, “Superoffload: Unleashing the power of large-scale llm training on superchips,” in Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, 2026, pp. 249–264. [36] V. A. Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro, “Reducing activation recomputation in large transformer models,” Proceedings of Machine Learning and Systems, vol. 5, pp. 341–353, 2023. [37] N. Dryden, N. Maruyama, T. Moon, T. Benson, M. Snir, and B. Van Essen, “Channel and filter parallelism for large-scale cnn training,” in Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2019, pp. 1–20. [38] Y. Oyama, N. Maruyama, N. Dryden, E. McCarthy, P. Harrington, J. Balewski, S. Matsuoka, P. Nugent, and B. Van Essen, “The case for strong scaling in deep learning: Training large 3d cnns with hybrid parallelism,” IEEE Transactions on Parallel and Distributed Systems, vol. 32, no. 7, pp. 1641–1652, 2020. [39] NCAR, “MILES-CREDIT,” https://github.com/NCAR/miles-credit, accessed: 2026-06-22. [40] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga et al., “Pytorch: An imperative style, high-performance deep learning library,” Advances in neural information processing systems, vol. 32, 2019. [41] T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré, “Flashattention: Fast and memory-efficient exact attention with io-awareness,” Advances in neural information processing systems, vol. 35, pp. 16 344–16 359, 2022.

Related documents

• TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling#612633
• TERRA: A Hierarchical Parallel Training and Memory Orchestration Framework for High-Resolution AI-based Earth Modeling#611156
• A Checklist to assess the energy and carbon impacts of ML/AI applications in Earth System Modeling#820851
• MarsCast: Transfer Learning of AI Weather Foundation Models to Planetary Atmospheres#429360
• Optimal scenario design for climate emulation#1141574