Method for evaluating computing resources of a meteorological large model

Abstract

This invention discloses a method for evaluating the computational resources of a large-scale meteorological model, comprising: step (1) evaluation of the parameters of the large-scale meteorological model; step (2) evaluation of the FLOPs of the large-scale meteorological model; step (3) evaluation of the memory usage of the large-scale meteorological model; and step (4) evaluation of the distributed communication of the large-scale meteorological model. This invention proposes for the first time a multi-granularity joint evaluation framework for computational resources. This framework establishes a parameter calculation model, a spatiotemporal awareness FLOPs evaluation model, a memory usage model, and a distributed communication analysis model by dividing the data into modules, and combines these with the spatiotemporal heterogeneity characteristics of meteorological data to quantify the hardware resources required for the large-scale meteorological model.

Description

Technical Field

[0001] This invention relates to the field of meteorological large-scale model application technology. Specifically, it relates to a method for evaluating the computational resources of a meteorological large-scale model. Background Technology

[0002] In recent years, breakthroughs in deep learning technology have driven the rapid penetration of large-scale meteorological models. Large-scale meteorological models, such as FourCastNet, GraphCast, Pangu-Weather, Fengwu, and Fuxi, have demonstrated their potential to surpass traditional numerical methods in tasks like high-resolution weather forecasting and extreme event simulation through data-driven approaches. However, their large-scale training and deployment face severe computational resource challenges, especially in high-resolution weather forecasting scenarios. A single inference involves the spatiotemporal multidimensional characteristics of the atmospheric field (such as longitude, latitude, altitude, time step, and coupling of multiple physical variables), resulting in input data volumes reaching gigabytes, exacerbating the computational burden. Existing resource assessment methods are mainly designed for large-scale models in the field of Natural Language Processing (NLP) and are ill-suited to the dynamic computational needs of meteorological tasks (such as spatiotemporal multidimensional characteristics) and the unique architecture of meteorological models, leading to low resource utilization and high computational costs.

[0003] Currently, research on computational resource evaluation for large-scale deep learning models mainly focuses on the field of natural language processing (NLP). Language models represented by GPT and BERT are usually based on the Transformer architecture. Their computational resource analysis methods mainly revolve around indicators such as the number of model parameters, GPU memory usage, and floating-point operations (FLOPs), and improve computational efficiency by optimizing GPU memory allocation and batch processing strategies.

[0004] While large language models based on the Transformer architecture (such as GPT-3 and LLaMa) have made groundbreaking progress in the field of natural language processing, their computational resource evaluation methods have formed a structured theoretical system. Researchers have quantified the hardware resource requirements of large models, such as memory, computing power, and bandwidth, through parameter modeling, computational load decomposition, and distributed communication analysis. For example, Brown et al. proposed a formula for calculating the number of language model parameters, covering the embedding layer, Transformer layer, and output layer, providing a foundation for evaluating memory usage and training costs. In addition, research shows that the computational cost per iteration (FLOPs) can be correlated with the number of model parameters and the length of the input sequence, and the prediction error using this method is relatively small. However, existing NLP computational resource evaluation methods have certain limitations. These methods assume a fixed input sequence length, making it difficult to capture resource fluctuations caused by dynamic batch processing or sparse attention mechanisms. Furthermore, communication overhead is significantly affected by hardware topology. Gradient synchronization in distributed training depends on the hardware topology (such as NVLink and InfiniBand), but existing distributed training tools lack cross-platform generalization capabilities. For example, DeepSpeed's communication optimization strategy in heterogeneous hardware environments lacks universality. These limitations stem from the fact that NLP tasks typically process one-dimensional text data, and the computational complexity is mainly determined by the sequence length and the number of layers. However, the computational load of large meteorological models involves spatiotemporal coupling, and the computational complexity varies with data resolution and time step, making it difficult to directly apply existing evaluation tools.

[0005] The computational resource assessment of large-scale language models mainly revolves around parameter modeling, computational load analysis, memory usage prediction, and distributed training overhead. The core issue lies in quantifying the dynamic balance between model size and hardware resources. Taking GPT-3 (175B parameters) as an example, its parameter count can be decomposed into a weighted sum of the embedding layer, multi-layer Transformer structure, and output layer, as calculated below:

[0006] (1) Embedding layer parameters: vocabulary size With hidden dimensions The product of, i.e. ;

[0007] (2) Transformer parameters: single-layer self-attention With feedforward network ,total If it contains A multi-layer transformer contains parameters. After 96 layers are stacked, it reaches parameter;

[0008] (3) Output layer parameters: linear mapping from hidden layer to vocabulary The final total number of parameters is: Nparams = V∙d +12 ld 2 + d∙V ≈175 billion.

[0009] The size of this parameter directly determines the computational load and memory requirements of the model.

[0010] The computational load during the training phase is primarily measured in floating-point operations (FLOPs), encompassing forward propagation, back propagation, and gradient updates. The computational cost per iteration (FLOPs) can be approximated as FLOPs. iter ≈6× N total × S × B ( S For sequence length, B (for batch size), including forward (2) N total S ) and backpropagation (4) N total S The accumulation of () is performed when the length of the input sequence is... S =2048, B When the training efficiency is 1, a single iteration of GPT-3 training requires 2.15 × 10⁻⁶ units. 15 FLOPs, if calculated continuously using 1000 A100 GPUs (312 TFLOPS), theoretically take about 3.68 hours (in practice, the computing power utilization rate of the large model after optimization reaches 52%), highlighting the extremely high computing power requirements and severe challenges of the training cycle during the training phase.

[0011] Memory requirements map the number of parameters to hardware constraints. At FP32 precision, 175 billion parameters occupy approximately 700GB of memory, the equivalent gradients generated by backpropagation add up to 700GB, and the Adam optimizer requires additional storage for momentum and variance (1400GB), totaling 2.8TB. Furthermore, during forward propagation, the memory usage of intermediate activation values ​​increases linearly with sequence length (S) and the number of layers (L), where… For the number of bytes (e.g., FP32, b=4): ;when and At that time, the activated video memory exceeded 7.5TB, far exceeding the capacity of a single GPU, forcing the adoption of a distributed training strategy to distribute the storage pressure.

[0012] The communication overhead of distributed training is determined by both data parallelism and model parallelism strategies. Data parallelism relies on the RingAll-Reduce algorithm to synchronize gradients, and its communication overhead can be modeled as follows:

[0013] ;

[0014] Taking GPT-3 as an example, when 1024 GPUs synchronize 700GB gradients using a 200Gbps InfiniBand network, this may account for 15% to 20% of the total training time, with the specific time varying depending on the parallelization strategy and optimization. In model parallelism, the bubble time proportion of pipeline parallelism increases with the number of stages. With micro-batch quantity Increase and rise:

[0015] ;

[0016] With a 32-stage pipeline and a 64-micro-batch configuration, device idle time reaches as high as 32.63%, becoming a significant bottleneck for scalability. If tensor parallelism (matrix multiplication split by attention head) is added, each layer needs to exchange intermediate results of dimension d across devices, resulting in a single-layer communication volume of [amount missing]. The cumulative communication latency of 96 layers can offset the computational gain, and a parallel strategy that combines NVLink and InfiniBand hybrid topology optimization is required.

[0017] Resource assessment during the inference phase focuses on the trade-off between computational density and memory efficiency. The FLOPs of a single forward propagation are approximately one-third of those during the training phase. During inference, GPT-3 needs to generate 2048 tokens. The A100 GPU theoretically takes 2.4 seconds. Memory usage is significantly reduced to 1.5TB due to precision compression (FP 16) and activation value pruning, but throughput optimization is still constrained by I / O bandwidth.

[0018] When the batch size At that time, the measured throughput reached 12,000 tokens / second, but the memory requirement doubled accordingly, reflecting the Pareto frontier of computing and storage in inference scenarios.

[0019] Resource evaluation for large language models is a multi-dimensional optimization problem, involving a comprehensive consideration of parameter count, computational load, memory usage, and communication overhead. (Parameter count) It is a core metric for model size, directly driving the FLOPs requirements for training and inference, and affecting efficiency through memory usage and communication overhead.

[0020] The rapid development of large-scale meteorological models poses multidimensional challenges to traditional resource assessment methods. Compared to the text sequence processing of language models, meteorological models take on high-dimensional, multivariate spatiotemporal grid data (such as temperature, pressure, and humidity fields) as input, with data volumes reaching terabytes (TB) and dynamic spatiotemporal resolution (e.g., typhoon forecasting requires simultaneous processing of a 1km local fine-grained grid and a 10km global background field). Furthermore, the computational complexity of meteorological models increases non-linearly, whereas traditional methods based on fixed inputs (such as language data) have become increasingly complex. Figure 1 The estimation methods for FLOPs at different sizes can lead to significant differences, primarily due to the increased computational load caused by spatiotemporal convolution in meteorological models. The real-time requirements of extreme weather forecasting (such as minute-level typhoon track updates) further necessitate a multi-objective optimization framework to balance computational latency, prediction accuracy, and energy consumption. While tools like Megatron-LM perform well in language modeling scenarios, their static load assumptions and hardware adaptation strategies are insufficient for operational meteorological needs. For instance, the high-resolution computational demands of a region can trigger sudden load spikes, and existing tools lack adaptive resource scheduling mechanisms based on physical field characteristics.

[0021] Therefore, directly applying the computational resource assessment method for large-scale language models to the meteorological field has significant limitations:

[0022] (1) Meteorological data has strong temporal correlation. The development and evolution of different weather systems will cause the computational load to fluctuate significantly during the reasoning process. NLP tasks mainly deal with text sequences of fixed length, and the computational complexity is usually relatively stable. Existing evaluation methods are difficult to characterize the dynamic computational needs in meteorological tasks.

[0023] (2) Meteorological large models widely adopt architectures such as UNet, Swin Transformer, and GNN. The computational load of these architectures is affected by factors such as data resolution, spatial dimension changes, and inter-layer communication overhead. However, the existing evaluation framework is mainly geared towards language large models and is difficult to adapt to the computational characteristics of meteorological models.

[0024] (3) Weather forecasting tasks have extremely high requirements for computational timeliness. For example, typhoon path prediction needs to be updated every minute. The allocation of computing resources must be balanced between computational accuracy and inference latency. However, the optimization goal of NLP usually only focuses on minimizing latency. This strategy is difficult to meet the needs of actual meteorological operations.

[0025] Due to the aforementioned challenges, the current allocation of computing resources for large-scale meteorological models still relies heavily on human experience, such as fixed GPU quotas, resulting in consistently low resource utilization and increased costs associated with redundant computing power. Furthermore, as the resolution of meteorological models continues to increase (e.g., from 10km to 1km), traditional experience-based resource assessment methods are insufficient to predict future computing power demands, severely hindering the technological iteration and large-scale application of large-scale meteorological models. Therefore, developing a computing resource assessment method suitable for large-scale meteorological models is not only crucial for improving computing power utilization and reducing computing costs, but also an important support for ensuring the efficient and stable operation of large-scale meteorological models and promoting the development of intelligent meteorological services. Summary of the Invention

[0026] Therefore, the technical problem to be solved by this invention is to provide a multi-granularity computing resource joint evaluation method for large meteorological models that can adapt to the multi-scale characteristics of meteorological data and significantly improve the utilization rate of computing resources, and to construct a multi-granularity computing resource joint evaluation framework (MGCRJEF). Based on the computing characteristics of meteorological models, this framework establishes a four-dimensional computing resource evaluation system including the number of model parameters, computational complexity, memory usage, and distributed communication overhead.

[0027] To solve the above-mentioned technical problems, the present invention provides the following technical solution:

[0028] A method for evaluating the computational resources of a large meteorological model includes the following steps: Step (1) Evaluation of the number of parameters of the large meteorological model; Step (2) Evaluation of the number of FLOPs of the large meteorological model; Step (3) Evaluation of the memory usage of the large meteorological model; Step (4) Evaluation of the distributed communication of the large meteorological model.

[0029] In the above-mentioned computational resource assessment method for the large meteorological model, step (1) involves the total number of parameters in the Unet within the large meteorological model. Calculation:

[0030] ;

[0031] In the formula, L enc The number of encoder layers. L dec The number of decoder layers. L skip The number of skip connection layers. The number of parameters for each encoder layer's convolutional layer. The number of parameters for the convolutional layer of each decoder layer. For the parameters of the batch normalization layer, The parameters of the skip connection layer;

[0032] Calculation of parameters for the convolutional layers in the encoder and decoder layers:

[0033] The number of parameters in a 2D convolutional layer ;

[0034] The number of parameters in a 3D convolutional layer ;

[0035] Parameters of batch normalization layer Calculation: ;

[0036] Parameters of the skip connection layer Calculation: ;

[0037] In the formula, Input the number of channels. c is the number of output channels. The height of the convolution kernel. The width of the convolution kernel. The depth of the 3D convolution kernel.

[0038] In the above-mentioned computational resource assessment method for large meteorological models, step (1) involves the total number of parameters of the Swin-transformer in the large meteorological model. Calculation:

[0039] ;

[0040] In the formula, The number of parameters for the Patch Embedding layer. For the number of stages, The number of parameters for the Patch Merging layer. The number of blocks in each stage, The parameter count for the Swing TransformerBlock layer. For the parameters of the output head;

[0041] Number of parameters in the Patch Embedding layer: ;

[0042] Number of parameters in the Patch Merging layer: ;

[0043] Number of parameters in the Swin Transformer Block layer:

[0044] ;

[0045] Output header parameters: ;

[0046] In the formula, C For the number of meteorological variables, p For Patch Size, choose 4×4 or 8×8. D For the hidden layer dimension, The scaling factor of the multilayer perceptron. This represents the number of output channels.

[0047] In the above-mentioned computational resource assessment method for large meteorological models, step (1) involves the total number of parameters in the graph neural network of the large meteorological model. Calculation:

[0048] ;

[0049] In the formula, This refers to the number of GCN or GAT layers in the graph neural network. For parameters of the GCN layer or GAT layer, The parameters for edge operations;

[0050] Parameter calculation of GCN layer: ;

[0051] Parameter calculation of GAT layer: ;

[0052] Calculation of parameters for edge operations: ;

[0053] In the formula, The input node feature dimension, To output the feature dimensions of the nodes, The number of heads for multi-head attention. The input edge feature dimension, This is the dimension of the output edge features.

[0054] In the above-mentioned computational resource assessment method for large meteorological models, step (1) involves the total FLOPs of Unet in the large meteorological model. Calculation:

[0055] ;

[0056] In the formula, The number of encoder layers. The number of decoder layers. FLOPs of the convolutional layers for each encoder layer FLOPs of the convolutional layers for each decoder layer For batch normalized layers of FLOPs;

[0057] FLOPs calculation for convolutional layers in encoder and decoder layers:

[0058] ; ;

[0059] FLOPs calculation for batch normalized layers:

[0060] ;

[0061] In the formula, To output the high-resolution feature map, To determine the width of the output feature map, Input the number of channels. Number of output channels The height of the convolution kernel, The width of the convolution kernel, The length of the time dimension of the output feature map;

[0062] The calculation of FLOPs for transposed convolutional layers is the same as that for ordinary convolutional layers, but the output size... , These are determined by the step size and the fill size, respectively.

[0063] In the above-mentioned computational resource assessment method for large meteorological models, step (2) involves the total FLOPs of the Swin-transformer in the large meteorological model. Calculation:

[0064] ;

[0065] In the formula, For the FLOPs of the Patch Embedding layer, For the number of stages, The number of blocks in each stage, For FLOPs of the multi-head self-attention mechanism layer of windows, For FLOPs of a multilayer perceptron, For the FLOPs of the Patch Merging layer, To output the FLOPs of the head layer;

[0066] Calculation of FLOPs in the Patch Embedding layer:

[0067] ;

[0068] Calculation of FLOPs in the multi-head self-attention mechanism layer of the window:

[0069] ;

[0070] Calculation of FLOPs for a multilayer perceptron layer: ;

[0071] FLOPs calculation for the Patch Merging layer: ;

[0072] Calculation of FLOPs in the output head layer: ;

[0073] In the formula, C For the number of meteorological variables, To output the high-resolution feature map, To determine the width of the output feature map, For the hidden layer dimension, Number of output channels For the number of windows, This refers to the window size.

[0074] In the above-mentioned computational resource assessment method for large meteorological models, step (2) involves the total FLOPs of the graph neural network in the large meteorological model. Calculation:

[0075] ;

[0076] In the formula, This refers to the number of GCN or GAT layers in the graph neural network. For FLOPs of GCN or GAT layers, FLOPs for edge operations;

[0077] Calculation of FLOPs in the GCN layer:

[0078] ;

[0079] Calculation of FLOPs in the GAT layer:

[0080] ;

[0081] Calculation of FLOPs for edge operations: ;

[0082] In the formula, For the number of nodes, Let the number of edges be denoted as '('). The input node feature dimension, To output the feature dimensions of the nodes, The number of heads receiving multi-head attention.

[0083] In the above-mentioned method for evaluating the computational resources of the large meteorological model, step (3) includes the memory usage of the large meteorological model, which includes the memory usage of training and the memory usage of inference.

[0084] Calculation of training memory usage:

[0085] ;

[0086] Calculation of inference memory usage: ;

[0087] In the formula, For parameters, Storage precision, For sequence length, For the hidden layer dimension, This represents the number of model layers.

[0088] The above-mentioned method for evaluating the computational resources of the meteorological big model includes the following steps in step (4): Step (4-1) Analysis of data parallel communication overhead; Step (4-2) Evaluation of model parallel communication overhead; Step (4-3) Analysis of the impact of network topology on communication overhead; Step (4-4) Evaluation of total training communication overhead; Step (4-5) Evaluation of latency and throughput.

[0089] In the computational resource assessment method for the above-mentioned large meteorological model, step (4-1) includes: total communication time c Calculation: ;

[0090] In the formula, For model runtime computation time, The overlap ratio, Total communication time;

[0091] ;

[0092] ;

[0093] In the formula, P For the number of model parameters, b For parameter accuracy, B The total gradient is divided into buckets; the effective bandwidth is the actual available network bandwidth; the latency includes network transmission latency, protocol overhead, and synchronization wait time; and the hop count is the number of intermediate nodes in the data transmission path.

[0094] In step (4-2): assuming the embedding dimension is... The sequence length is The number of GPUs is The communication volume at each layer is ;

[0095] Bubble Time This refers to the idle time the equipment spends while waiting for other stages to complete. ,in Q This represents the number of stages in the model. For single-stage calculation time; ; The average stage time, For the first The computation time for each stage;

[0096] Total communication time Calculation: ,in For cross-stage communication time, , For the first The first stage to the second The amount of data to be transmitted in each stage; This refers to the quantity in a micro-batch. ;

[0097] In step (4-4): the total communication overhead of the meteorological large model Calculation:

[0098] ;

[0099] In steps (4-5), the single inference delay Including calculation time and data transmission time : , For parameters, The sequence length;

[0100] Throughput tokens / s calculation: ; Maximum token length This represents the number of parallel devices.

[0101] The technical solution of the present invention achieves the following beneficial technical effects:

[0102] 1. The computational resource assessment method for large-scale meteorological models in this invention includes: constructing a parameter calculation framework for meteorological models, calculating model parameters for UNet, Swin Transformer, and GNN architectures respectively, establishing a general parameter calculation formula to provide basic data support for resource assessment; proposing a spatiotemporal awareness-based FLOPs calculation method, analyzing the computational complexity of different model architectures, and improving the accuracy of computational resource assessment; establishing a memory usage prediction model, calculating memory usage during the training and inference stages respectively, and performing accurate modeling in conjunction with resolution, batch size, etc., to provide data support for computational resource allocation; and proposing a distributed communication assessment method for large-scale meteorological models, analyzing the impact of data parallelism, model parallelism, and communication overhead on computational resource overhead, and performing quantitative analysis of communication overhead, latency, and throughput to provide decision-making basis for distributed deployment optimization.

[0103] 2. The computational resource assessment method for large-scale meteorological models in this invention not only provides new methodological support for the optimization of computational resources in large-scale meteorological models, but also lays the foundation for efficient training and inference of next-generation intelligent meteorological systems. Through precise computational resource assessment, MGCRJEF provides new methodological support for the optimization of computational resources in large-scale meteorological models, further promoting the application of meteorological artificial intelligence technology.

[0104] 3. Meteorological models are inputted with high-dimensional, multivariate spatiotemporal grid data (such as temperature, pressure, and humidity fields), with data volumes reaching terabytes and dynamic spatiotemporal resolution. These characteristics lead to prediction errors exceeding 40% for traditional FLOPs formulas (based on fixed sequence lengths in language). This invention addresses the unique needs of large-scale meteorological models by, for the first time, incorporating the spatiotemporal dimension of meteorological data into the evaluation index design. By constructing a resolution-sensitive memory-consumption model-dependent FLOPs prediction method, it achieves fine-grained resource quantification, analyzes the resource consumption patterns of the model at different spatiotemporal granularities, and supports meteorological agencies in independently selecting the optimal configuration scheme in operational scenarios. This framework provides a scalable methodological foundation for the sustainable deployment of large-scale meteorological models.

[0105] 4. The computational resource assessment method for large-scale meteorological models in this invention achieves fine-grained resource quantification from model structure to hardware usage by dynamically mapping parameter quantities to computational resources. This allows for the analysis of resource consumption of large-scale meteorological models at different spatiotemporal resolutions, helping meteorological agencies balance model performance and deployment costs based on operational needs. Furthermore, based on parameter-driven joint modeling technology, this method can adapt to the multi-scale characteristics of meteorological data, significantly improving computational resource utilization. It also provides quantitative basis for model compression and distributed strategy optimization, promoting the efficient iteration and sustainable application of large-scale meteorological models. Attached Figure Description

[0106] Figure 1A schematic diagram of language data (sequence position) in the background art of this invention;

[0107] Figure 2 This invention provides a framework for assessing computational resources for large-scale meteorological models. Detailed Implementation

[0108] Evaluation of computational resources for large-scale meteorological models presents unique challenges, primarily due to their hybrid architecture (e.g., local feature extraction using UNet, global attention modeling using Swin-Transformer, and graph-structured reasoning using GNN). To address this, this embodiment presents a computational resource evaluation method for large-scale meteorological models (Multi-Granularity Computing Resource Joint Evaluation Framework, MGCRJEF). This method considers four core dimensions: model parameter count, computational complexity, memory usage, and communication overhead. It also incorporates the spatiotemporal characteristics of meteorology to construct a dynamic resource prediction model (see...). Figure 2 ).

[0109] I. Calculation of Parameters for Large Meteorological Model

[0110] The number of model parameters is fundamental to computational resource assessment, determining memory usage, computational load, and communication costs. Different architectures of large-scale meteorological models exhibit significant differences in parameter distribution and computational patterns. This embodiment performs parameter calculation modeling for UNet (U-net, a deep learning architecture), Swin-Transformer structure, and GNN (Graph Neural Network) to accurately assess their parameter counts. This helps optimize computational resource allocation and improve the inference and training efficiency of large-scale meteorological models.

[0111] Calculation of Unet parameters

[0112] Encoder-decoder structures, represented by UNet, are widely used in weather forecasting, such as high-resolution precipitation forecasting. Their local feature extraction capabilities enable them to effectively capture spatial details in meteorological fields. Their parameters consist of convolutional layers, batch normalization layers (BN layers), and skip connection layers.

[0113] (1) Formula for calculating the number of parameters in a convolutional layer

[0114] Number of parameters per 2D convolutional layer : .

[0115] Number of parameters per 3D convolutional layer : .

[0116] (2) Formula for calculating batch normalized layer parameters

[0117] Each channel in a batch normalization layer has two parameters: scale and shift. Therefore, .

[0118] (3) Formula for calculating parameters of skip connection layer

[0119] When UNet's skip connection layer needs to align the number of channels using a 1×1 convolution, the additional parameters are the number of parameters in the skip connection layer. The same applies to the number of parameters in the deconvolution layer of the decoder.

[0120] In this embodiment, Input the number of channels (e.g., the number of channels for different meteorological elements). This represents the number of output channels. The height of the convolution kernel is typically the size of the window along the latitudinal direction. The width of the convolution kernel is typically the window size along the longitude direction. The depth of the 3D convolution kernel is typically the size along the time dimension.

[0121] Therefore, the total number of parameters in UNet is the sum of the parameters of the encoder, decoder, and all layers in the skip connections. Assume the UNet model contains... L enc Each encoder layer L dec Each decoder layer and L skip A skip connection layer, the total number of parameters of Unet Represented as:

[0122] ;

[0123] In the formula, The number of parameters in the convolutional layer representing each encoder layer. This represents the number of parameters in the convolutional layer of each decoder layer.

[0124] (4) Example of UNet parameter calculation

[0125] This embodiment uses the classic 4-layer UNet as an example to calculate the number of UNet parameters.

[0126] A. Encoder section

[0127] Each layer contains two convolutional layers (typically with 3×3 convolutional kernels) and a batch normalization (BN) layer, with the number of channels doubling with each subsequent layer.

[0128] Level 1 (Channels 1 → 64):

[0129] .

[0130] Level 2 (Channels 64 → 128):

[0131] .

[0132] Level 3 (Channels 128 → 256):

[0133] .

[0134] Level 4 (Channels 256 → 512):

[0135] .

[0136] Therefore, the total number of encoder parameters = 37696 + 221696 + 885760 + 3540992 = 4686144

[0137] B. Bottleneck layer

[0138] The intermediate layer connecting the encoder and decoder is the Bottleneck layer; Bottleneck layer (512 → 1024 channels):

[0139] .

[0140] C. Decoder section

[0141] Each layer contains a transposed convolution (2×2 kernels) and two convolutional layers (3×3 kernels), with the number of channels halved with each subsequent layer.

[0142] Level 1 (Channels 1024 → 512):

[0143] .

[0144] Level 2 (Channels 512 → 256):

[0145] .

[0146] Level 3 (Channels 256 → 128):

[0147] .

[0148] Level 4 (128 → 64 channels):

[0149] .

[0150] Therefore, the total number of decoder parameters = 917088 + 2294784 + 573952 + 143616 = 12189440.

[0151] D. Output layer

[0152] Assume the output has 2 categories:

[0153] .

[0154] Therefore, the total number of parameters in UNet = encoder parameters + BN layer parameters + decoder parameters + output layer parameters .

[0155] Calculation of Swin-transformer structural parameters

[0156] The Swin-Transformer captures multi-scale meteorological features through a hierarchical window attention mechanism, making it particularly suitable for efficiently modeling meteorological data, such as large models like the Pangu-weather and wind-crow models. Many large meteorological models currently employ the Swin-Transformer structure. Unlike the pure Transformer structure used in large language models, the Swin-Transformer typically includes a Patch Embedding layer, a Swin Transformer Block layer, a Patch Merging layer, and an output head.

[0157] (1) Formula for calculating the number of parameters in the Patch Embedding layer

[0158] The key component, Patch embedding, can process input meteorological data (such as grid fields, dimensions [...]). B, C, H, W ]) divided into p × p The patches are linearly mapped to the hidden dimensions. D Therefore, the number of parameters in Patch Embedding P embedding The calculation formula is:

[0159] ;

[0160] Number of input channels C For meteorological variables (such as temperature, pressure, etc.); p This is the Patch Size, which is usually 4×4 or 8×8.

[0161] (2) Formula for calculating the parameters of the Swing Transformer Block layer

[0162] Each Swing Transformer Block layer contains a normalization layer (LayerNorm), a window multi-head self-attention mechanism layer (W-MSA), and a multilayer perceptron layer (MLP).

[0163] a. Normalization Layer: Each Swing Transformer Block layer contains two LayerNorm layers, each with 2 parameters. D (Scaling factor γ and bias β). Therefore, the number of parameters in the LayerNorm layer. .

[0164] b. Multi-head Self-Attention Mechanism Layer: The parameter calculation of the multi-head window attention mechanism layer is similar to that of the self-attention layer in the large language model, but the input is the feature dimension. The windowed multi-head self-attention QKV projection layer contains three linear transformation matrices with 3D parameters. 2 It also includes an output projection layer with D parameters. 2 ,therefore, .

[0165] c. Multilayer Perceptron Layer: In a Multilayer Perceptron (MLP), the input dimension D of the extended layer is converted to the output dimension... Its parameter count is The parameters of the contraction layer are the same as those of the expansion layer. If the partially linear layer contains a bias term, the D parameter is increased for each linear layer (e.g., the MLP expansion layer and contraction layer are increased respectively). And D); therefore, adding the parameter quantities of the above structures together, the parameter quantity of a single-layer Swing Transformer Block The calculation formula is as follows:

[0166] ;

[0167] in For MLP expansion ratio, typically ;therefore It can be simplified to .

[0168] (3) Formula for calculating the number of parameters in the Patch Merging layer

[0169] Patch merging layers can merge adjacent 2×2 patches, halving the spatial size and doubling the number of channels. The parameter calculation of patch merging layers mainly includes two parts: linear transformation and normalization layer.

[0170] a. Linear Transformation Layer: The Patch Merging layer merges adjacent 2×2 patches together, resulting in a patch with the following dimension: (in (This refers to the number of input channels); then, through a linear transformation, the dimensionality is reduced to... Therefore, the number of parameters in the linear transformation layer is .

[0171] b. Normalization Layer: Patch merging layers typically use a normalization layer to normalize the merged features. The number of parameters in the normalization layer is usually twice the number of the input dimensions; therefore, the number of parameters in the normalization layer is... .

[0172] Therefore, the total number of parameters in the Patch Merging layer It is to add the parameters of the two parts mentioned above, that is:

[0173] .

[0174] (4) Calculation formula for output head parameters

[0175] The output header is used to map the final features to the predicted state. Output header parameter count. , This represents the number of output channels.

[0176] In summary, the total number of parameters in the Swin-transformer structure is... The calculation formula is as follows: .

[0177] In the formula, For the number of stages, The number of blocks in each stage.

[0178] Calculation of parameters in Graph Neural Networks (GNNs)

[0179] The core of GNN is to update node representations by aggregating information from nodes and edges. It is suitable for graph structure modeling in meteorological data. Common variants include GCN (Graph Convolutional Network) and GAT (Graph Attention Network). The following is a method for calculating its parameters:

[0180] (1) Formula for calculating the number of parameters in GCN layer

[0181] ;

[0182] In the formula, It is the feature dimension of the input node. It is the output node feature dimension. This is the weight matrix. This refers to bias.

[0183] (2) Formula for calculating the parameters of the GAT layer

[0184] ;

[0185] In the formula, This is the weight matrix. For attention parameters, multi-head attention: each head is calculated independently and needs to be multiplied by the number of heads. Attention parameters: Each head needs to learn two weight vectors to calculate the attention score.

[0186] (3) Formula for calculating side operation parameters

[0187] Edge operation parameters, such as edge weights. If the model contains learnable edge weights (such as the correlation strength between stations in a weather map), the linear transformation parameters of the edge features need to be calculated additionally, i.e., the number of edge operation parameters. .

[0188] .in, The input edge feature dimension, This is the dimension of the output edge features.

[0189] In summary, if the GNN model contains If the number of layers is GCN or GAT, then the total number of parameters in GNN is... .

[0190] II. FLOPs Evaluation of Spatiotemporal Perception in Meteorological Large Models

[0191] 1. Calculation of FLOPs (floating-point operations per second) in UNet

[0192] UNet's FLOPs mainly consist of convolutional layers, batch normalization (BN) layers, and transposed convolutional layers.

[0193] (1) Formula for calculating FLOPs of convolutional layers (2D / 3D)

[0194] ;

[0195] In the above formula, To output the height and width of the feature map, Number of input and output channels, The height and width of the convolution kernel are given. Since the convolutional layer needs to perform multiplication and addition operations, it needs to be multiplied twice during the calculation.

[0196] If a 3D convolutional layer is used, then

[0197] ;

[0198] The length of the time dimension of the output feature map. To output the height and width of the feature map, Number of input and output channels, These represent the height and width of the convolution kernel. Since the convolutional layer requires multiplication and addition operations, the kernel needs to be multiplied twice during computation.

[0199] (2) Formula for calculating FLOPs of BN layer

[0200] Each channel in the Batch Normalization (BN) layer requires one multiply-accumulate operation for scaling and shifting. Therefore, .

[0201] (3) Formula for calculating FLOPs of transposed convolution (deconvolution)

[0202] The FLOPs of transposed convolution are calculated in the same way as those of ordinary convolution, but the output size is different. It is determined by the step size and the fill factor. That is: .

[0203] Therefore, assuming the UNet model contains L enc Each encoder layer L dec There are one decoder layer; the total FLOPs of UNet can be calculated as follows:

[0204] ; FLOPs of the convolutional layers for each encoder layer FLOPs for each convolutional layer of the decoder layer.

[0205] (4) FLOPs calculation of UNet (example)

[0206] This embodiment uses the classic 4-layer UNet as an example to calculate the FLOPs of UNet.

[0207] A. Encoder section

[0208] Each layer contains two 3×3 convolutional layers and a BN layer.

[0209] Level 1 (Channels 1 → 64):

[0210] .

[0211] Layer 2 (64→128 channels, input size changes) ):

[0212] .

[0213] The other layers are similar to layers 1 and 2, but the number of channels is doubled. Therefore, the total FLOPs of the encoder are...

[0214] .

[0215] B. Bottleneck layer

[0216] .

[0217] C. Decoder section

[0218] Each layer contains a transposed convolution (2×2) and two 3×3 convolutions.

[0219] .

[0220] The other layers are similar to layers 1 and 2, but the number of channels is halved. Therefore, the total FLOPs of the decoder are...

[0221] .

[0222] D. Output Layer

[0223] 1×1 convolution (64→1 channel):

[0224] .

[0225] Therefore, UNet's total FLOPs = encoder's total FLOPs + Bottleneck layer FLOPs + decoder's total FLOPs + output layer FLOPs .

[0226] FLOPs calculation of the Swing Transformer structure

[0227] The FLOPs of the Swin-Transformer dynamically adjust as the local complexity of the weather field changes. Mutual information between variables affects the choice of the effective number of attention heads, thus significantly impacting the computational load.

[0228] (1) Formula for calculating FLOPs of the Patch Embedding layer

[0229] Swin-Transformer's FLOPs calculation begins with Patch Embedding, which divides the input meteorological data into patches and performs a linear mapping. The formula for calculating FLOPs in the Patch Embedding layer is as follows:

[0230] ;

[0231] in, The size of the output feature. C Where C is the number of input channels (number of meteorological variables), and D is the dimension of the hidden layer; Patch Embedding is equivalent to a linear layer (input dimension C, output dimension D).

[0232] (2) Formula for calculating FLOPs of the Swing Transformer Block layer

[0233] a. Window Multi-Head Self-Attention Mechanism Layer (W-MSA):

[0234] ;

[0235] Where N is the number of windows, W 2 For window size (e.g., 7×7); 4×N×D 2 The computational cost of FLOPs for QKV projection and output projection is 2 × N × D × W. 2 Calculate the attention score.

[0236] b. Multilayer Perceptron (MLP): Equivalent to two linear layers, with a typical scaling ratio of 4; therefore, .

[0237] (3) Formula for calculating FLOPs of Patch Merging layer

[0238] Patch merging combines adjacent 2×2 patches and performs a linear transformation to reduce dimensionality. The FLOPs of the Patch merging layer are: Where 4C represents the merged input dimension and 2C represents the output dimension.

[0239] (4) Formula for calculating FLOPs in the head layer

[0240] The output head can map the final features to the predicted state, and its FLOPs are calculated using the formula... for: ; This represents the final number of output channels.

[0241] In summary, the formula for calculating the total FLOPs of the Swin Transformer structure is:

[0242] ;

[0243] In the formula, For the number of stages, The number of blocks in each stage. This hierarchical calculation method fully demonstrates the complexity and flexibility of Swin-Transformer in processing meteorological data.

[0244] FLOPs calculation for Graph Neural Networks (GNNs):

[0245] The FLOPs of graph neural networks (GNNs) in meteorological tasks mainly come from graph convolutional network (GCN) layers, graph attention network (GAT) layers, and edge weight operations.

[0246] (1) Formula for calculating FLOPs of GCN layer

[0247] ;

[0248] in, For the number of nodes, Let the number of edges be denoted as '('). For characteristic transformation matrix multiplication, For the corresponding adjacency matrix aggregation operation, where For input / output feature dimensions.

[0249] (2) Formula for calculating FLOPs of GAT layer

[0250] ;

[0251] Where H represents the number of attention heads, the first term is the computational cost of multi-head projection, and the second term is the overhead of attention score calculation and aggregation. For input / output feature dimensions.

[0252] (3) Formula for calculating the FLOPs of edge weights

[0253] If the model involves edge weight learning, the learnable weights for each edge require a multiplication-addition operation. Therefore, the FLOPs calculation for edge weights needs to be multiplied by 2. That is, the formula for calculating edge weight FLOPs is:

[0254] .

[0255] In summary, assuming the GNN model contains L The formula for calculating the total FLOPs of a graph neural network (GNN) with GCN or GAT layers is:

[0256] .

[0257] This calculation method reflects the iterative overhead of GNNs when processing irregular grid data in meteorological tasks, which is closely related to the size and complexity of the graph.

[0258] III. Prediction of Memory Usage in Large-Scale Meteorological Models

[0259] GPU memory usage is a core resource metric in the training and inference processes of deep learning models, and its dynamic changes directly affect the model's performance and efficiency, especially when processing large-scale meteorological data. This embodiment explores methods for predicting GPU memory usage during the training and inference phases, providing a systematic framework for resource management and optimization by decomposing GPU memory components and providing calculation formulas.

[0260] Training memory usage

[0261] The memory consumption during training mainly consists of four parts: model parameters, backpropagation gradients, optimizer states, and intermediate activation values.

[0262] (1) Model parameter video memory: The video memory usage of model parameters is determined by the number of parameters P and the storage precision b (in bytes), and the calculation formula is as follows:

[0263] ;

[0264] For example, in FP32 precision, each parameter occupies 4 bytes ( In FP16 precision, it is 2 bytes ( In mixed-precision training, if copies of the FP32 parameters need to be retained (e.g., for gradient updates), the actual GPU memory requirement will double, i.e. .

[0265] (2) Backpropagation gradient memory: The gradient memory occupied by backpropagation is the same as that of the parameters, which is equal to the memory of the model parameters. The gradient precision is usually matched with the parameter precision (FP32 or FP16). If mixed precision training is used, the gradient is usually stored in FP32 to avoid numerical stability problems.

[0266] (3) Optimizer memory: Taking the Adam optimizer as an example, it is necessary to maintain the first-order momentum and second-order momentum states for each parameter. At FP32 precision, each parameter occupies 8 bytes (4 bytes each for momentum and variance), and the memory is:

[0267] ;

[0268] (4) Intermediate activation value memory: The intermediate activation values ​​generated during forward propagation need to be reused during backward propagation. Their memory usage is strongly related to the model structure and input size. The memory usage can be approximated as:

[0269] ;

[0270] Where S is the sequence length, D is the hidden layer dimension, L is the number of model layers, and the coefficient 2 indicates that intermediate results of forward and backward propagation are stored simultaneously.

[0271] The total memory usage for training is the sum of the four parts mentioned above:

[0272] ;

[0273] In large models (such as weather models), optimizer states and intermediate activation values ​​are often the main bottlenecks in GPU memory usage. For example, the number of parameters... The optimizer state alone requires 8GB of video memory; if the sequence length Hidden layer dimensions number of floors The activation value of the video memory at FP32 precision will be as high as Therefore, in actual training, it is necessary to combine methods such as mixed precision, gradient accumulation (reducing batch size), activation recomputation (recompiling intermediate values ​​on demand), and model parallelism (distributed storage of parameters) to compress the memory requirements to a range acceptable for a single card.

[0274] 2. Memory usage during the inference stage

[0275] The memory requirements during the inference phase are relatively simple, mainly consisting of model parameters and intermediate activation values, without involving gradients and optimizer states.

[0276] The memory usage of the model parameters is the same as that of the training phase. The memory usage of intermediate activation values ​​in inference is as follows: Intermediate activation values ​​are generated only by forward propagation. ;

[0277] Therefore, the total memory usage for inference is: .

[0278] IV. Distributed Communication Evaluation of Meteorological Large Model

[0279] In the distributed training of large-scale meteorological models, communication overhead is one of the key factors affecting training efficiency and inference performance. Evaluating communication overhead requires comprehensive consideration of communication volume, communication mode, network bandwidth, latency, and the specific implementation of parallel strategies (such as data parallelism and model parallelism). This section analyzes the communication overhead of large-scale meteorological models in data parallelism and model parallelism, explores the impact of network topology, and provides a method for evaluating total communication overhead. Furthermore, the analysis of latency and throughput during the inference phase provides theoretical support for resource allocation and communication optimization. Future research will explore further improvements in training efficiency through dynamic scheduling and adaptive parallelism strategies.

[0280] Data parallel communication overhead analysis

[0281] Data parallelism is a common strategy in distributed training, especially suitable for scenarios where large meteorological models are processing large datasets. In data parallelism, each GPU (Graphics Processing Unit) independently computes its local gradient, and then synchronizes the global gradient through communication operations. Gradient synchronization in data parallelism typically employs an All-Reduce operation, with Ring All-Reduce being a common algorithm. Assuming there are N GPUs, P model parameters, and b parameter precision, the amount of data sent and received by each GPU in Ring All-Reduce is:

[0282] ;

[0283] Communication time is affected by network bandwidth and latency; the actual communication time can be expressed as:

[0284] ;

[0285] The effective bandwidth refers to the actual available network bandwidth (such as 600GB / s for NVLink or 100Gbps for InfiniBand). The hop count is the number of intermediate nodes in the data transmission path. For example, communication between GPUs within the same server via PCIe / NVLink connection results in a hop count of 0 (direct communication). If cross-node communication requires a switch, the hop count is typically greater than 1, depending on the network topology.

[0286] For large-scale meteorological models, the number of parameters P can reach billions or even tens of billions. The communication volume increases significantly with the number of GPUs N, making the optimization of communication algorithms and network configuration crucial. To reduce communication frequency, gradient synchronization often employs a bucketing strategy, dividing the total gradient into B buckets. The communication volume per bucket is... Total communication time The sum of communication times for all buckets:

[0287] ;

[0288] Latency includes network transmission latency, protocol overhead, and synchronization wait time. This bucketing strategy reduces the number of communications but increases memory usage, requiring a trade-off between communication efficiency and GPU memory utilization. Furthermore, overlapping communication and computation can further optimize efficiency. Let the model's runtime be... The overlap ratio is η (0 ≤ η ≤ 1), then the total time is: In large-scale meteorological models, the data volume is large and the calculations are complex, with a high degree of overlap. η Improving efficiency (e.g., through asynchronous communication) is crucial for training efficiency.

[0289] Model parallel communication overhead evaluation

[0290] When large-scale meteorological models become too large for a single GPU's memory to accommodate, model parallelism becomes essential. Model parallelism includes tensor parallelism and pipelined parallelism. Tensor parallelism distributes each layer's operation (such as matrix multiplication) across multiple GPUs, requiring the exchange of some results across devices. Let the embedding dimension be D, the sequence length be S, and the number of GPUs be... The communication volume at each layer is Communication time also depends on effective bandwidth and network latency. Pipeline parallelism divides the model into multiple stages, with different GPUs processing different stages, requiring the transmission of intermediate activation values ​​and gradients. The communication volume is... Where M is the number of micro-batches. The relatively long length S of meteorological data sequences (such as long-term series forecasting) significantly increases the communication overhead of pipeline parallelism. In pipeline parallelism, bubble time is the idle time of equipment while waiting for other stages to complete, which can be expressed as… ,in Q For the number of stages, This refers to the computation time per stage. If the model is not divided evenly, the computation time per step will vary in different stages.

[0291] Let the first The calculation time for each stage is The slowest phase lasts for 10 days. Average stage time The computation time of the slowest stage The total time is determined as follows: ,in This refers to the inter-stage communication time. The formula for inter-stage communication time is: , For the first The first stage to the second The amount of data that needs to be transmitted at each stage. The complex hierarchical structure of the meteorological large model may lead to... Fluctuations need to be optimized through load balancing.

[0292] The impact of network topology on communication overhead

[0293] Network topology significantly impacts communication efficiency. Intra-node communication is achieved via NVLink, boasting bandwidth up to 600GB / s, zero hops, and extremely low latency. Cross-node communication, however, utilizes the InfiniBand network, offering approximately 100Gbps bandwidth with latency ranging from microseconds to milliseconds; the hop count depends on the distance between nodes and the number of switches. Large-scale meteorological models often require cross-node training, and communication overhead increases significantly with the number of nodes. Optimizing network configuration (e.g., reducing hop count) is crucial.

[0294] Total training communication overhead assessment

[0295] Taking into account the gradient synchronization of data parallelism, inter-layer communication of tensor parallelism, inter-stage communication of pipeline parallelism, and the impact of bubble time, it is suitable for distributed training scenarios of large meteorological models.

[0296] The total communication overhead of the meteorological big data model can be expressed as:

[0297] =Data parallelism overhead + Tensor parallelism overhead (intra-node) + Pipeline parallelism overhead (cross-node) + Other overhead .

[0298] Latency and throughput assessment

[0299] In the inference phase of large-scale meteorological models, latency and throughput are key performance indicators.

[0300] (1) The single inference latency includes computation time and data transmission time:

[0301] ,

[0302] in Related to I / O bandwidth, the high-dimensionality of meteorological data may increase transmission overhead. The estimation of FLOPs in a single inference iteration is related to the model structure. ,in For parameters, The sequence length is given.

[0303] (2) The formula for calculating throughput (tokens / s) is:

[0304] ;

[0305] in, Maximum token length. Increase the number of parallel devices. Increasing batch size can significantly improve throughput, but a balance between latency and throughput needs to be struck in real-time weather forecasting.

[0306] V. Case Study on the Computational Resource Assessment Framework for Large-Scale Meteorological Models

[0307] We employ the Multi-Granularity Computational Resource Joint Evaluation Framework (MGCRJEF) to quantify the core hardware resource requirements of the model from four dimensions: parameter quantity, computational complexity, memory usage, and communication overhead, combined with the spatiotemporal heterogeneity of meteorological data. This section uses the Pangu-Weather meteorological model as a typical example for analysis for two reasons: the Pangu-Weather model is based on the Swing Transformer architecture, while current mainstream meteorological models (such as Fengwu and Fuxi) all adopt a Swing-like hierarchical attention mechanism, making their computational patterns and resource consumption characteristics universally applicable; the publicly disclosed architectural details and hyperparameter configurations of the Pangu-Weather model provide a reproducible benchmark for resource evaluation, avoiding the incomparability caused by structural differences in other models (such as UNet). This section will clarify the differences in resource requirements of this model in practical applications through specific calculations, providing theoretical basis and practical guidance for hardware configuration and resource optimization.

[0308] The Pangu-Weather model is a high-precision weather forecasting model based on the Transformer architecture. It employs a 3D Earth-Specific Transformer (3DEST) and utilizes a Swin transformer encoder-decoder architecture, which excels at capturing the spatiotemporal characteristics of meteorological data. The resource requirements are assessed from four dimensions below.

[0309] (1) Parameter calculation

[0310] The parameters of Pangu-Weather mainly come from the Patch Embedding layer, the Swin blockTransformer layer, the Patch Merging layer, and the output layer. According to the article (BI, K., XIE, L., ZHANG, H., et al. Accurate medium-range global weather forecasting with 3D neural networks[J]. Nature, 2023, 619(7970): 533-538), the model's hyperparameters are configured as shown in Table 1. The number of upper-air meteorological variables input to the Pangu model is... The number of surface meteorological variables is (All based on the ERA5 dataset), the patch size for the high-altitude variable is [size missing]. The patch size for ground variables is The article does not provide the hidden layer dimension, but it can be viewed by examining the structure of the open-source ONNX inference model on GitHub. .

[0311] Table 1 Pangu-Weather Hyperparameter Configuration

[0312]

[0313] Therefore, the number of parameters in the Patch Embedding layer is:

[0314] .

[0315] The parameters of a Swin Transformer Block consist of LayerNorm layers, Window Attention (W-MSA), and a Multilayer Perceptron (MLP) matrix. According to section 3.1.1, the parameters of a single-layer Swin Transformer Block are: According to the article analysis, generally Therefore, the number of parameters in a single layer According to the article, the 3DEST encoder and decoder each have 8 layers, totaling 16 layers. Therefore, the total number of parameters in the Swin Transformer Block is approximately 255 million. The Patch Merging layer will merge adjacent... The patches are merged together, and the resulting patch has the following dimensions: Then, the dimension is reduced to [value] through a linear transformation. Therefore, the number of parameters in Patch Merging is: The output head layer is used to map the final features to the predicted state. Its parameter count is Summing the above four parts, we can deduce that the number of parameters in the Pangu-weather model is approximately 266 million, which differs from the 256 million stated in the paper by 3.90%. This demonstrates the rationality of our algorithm and provides a reliable theoretical basis for the deployment of large-scale meteorological models.

[0316] (2) Floating-point operation quantity (FLOPs) calculation

[0317] The calculation of FLOPs for the Pangu-Weather training model needs to take into account Patch Embedding, Windowed Multi-Head Self-Attention (W-MSA), MLP, Patch Merging, and the computational load of the output head.

[0318] In Pangu-Weather, the resolution of the upper-air variable in the input meteorological data is [resolution value missing]. After the Patch Embedding layer is embedded, the dimension is The input ground variable is... The patch size is 4×4 after embedding. Therefore, the FLOPs of the Patch Embedding layer are: TFLOPs, among which The FLOPs calculation of the Swin Transformer Block considers LayerNorm layers, window attention (W-MSA), and multilayer perceptron (MLP). According to section 3.2.2, the FLOPs of a single-layer Swin Transformer Block are calculated using the following parameters: According to the article analysis, window size for , The size after is The number of windows can be estimated. It is 3720. Typically, r=4, therefore the number of FLOPs per layer is approximately... TFLOPs. Based on 8 layers each for the encoder and decoder, totaling 16 layers, the total floating-point operations of the Swin Transformer Block are approximately 328 TFLOPs.

[0319] The FLOPs of the Patch Merging layer are a parameter number. The total FLOPs for the Patch Merging layer are 47.4 TFLOPs, as this layer requires 3 computations.

[0320] The output layer is used to map the final features to the predicted state, and its FLOPs are: GFLOPs. In summary, the total floating-point operations for a single-sample inference iteration of the model are approximately 418 TFLOPs. Training FLOPs calculations include forward propagation, backpropagation, and parameter updates, typically three times that of inference. Therefore, the total floating-point operations for a single-sample iteration of training of the Pangu-weather model can be estimated to be approximately 1254 TFLOPs.

[0321] (3) Memory usage analysis

[0322] Pangu-weather model (number of parameters) The memory usage (266 million) is broken down into the following values ​​during the training and inference phases: During the training phase, the memory usage for model parameters is... The gradient memory usage is the same as the model parameter memory. ， The optimizer's state memory usage (taking the commonly used Adam optimizer as an example) is as follows: Intermediate activation value memory usage and input sequence length Hidden layer dimension and number of layers Regarding the memory requirements for training at FP32 precision: In summary, the total training memory usage is 59.19 GB. During the inference phase, only the memory used for model parameters and intermediate activation values ​​is calculated, therefore the total inference memory usage is 28.53 GB.

[0323] (4) Assessment of communication overhead and hardware requirements

[0324] In the distributed training of the Pangu-weather model, the paper used 192 V100 GPUs, therefore the communication volume per iteration is approximately Intra-node communication latency based on NVLink (300GB / s) is negligible, while cross-node InfiniBand (100 Gbps) communication time is... The bandwidth utilization rate is assumed to be 80%.

[0325] (5) Analysis of training resource allocation based on V100 FP32 performance

[0326] The article shows that the Pangu-weather training used 192 NVIDIA V100 chips for 15 days, based on an effective computing power utilization rate of 35% (usually 30%-40%). Based on the official configuration (192 V100 GPUs, 100 epochs completed in 15 days, batch size=1), and considering the V100's theoretical FP32 computing power (14 TFLOPS) and actual training efficiency, the calculation logic is as follows: the number of floating-point operations per training iteration is 1254 TFLOPs, therefore the computation time for a single sample (batch size=1) iteration is approximately 256 seconds. According to the article, the training time is 100 epochs, and the dataset size is 56940 (39 years of ERA5 data at 6-hour intervals). Considering the overall training requirements, the 192 V100 GPUs complete 100 epochs of training in parallel, using NVLink intra-group interconnect + InfiniBand cross-node topology, and enabling gradient bucketing and asynchronous communication. Based on the above, the communication time is 2.5 seconds, so the total time for a single iteration is 258.5 seconds.

[0327] Considering the use of 192 V100 chips, the total training days are approximately The number of days is close to the 16 days mentioned in the article, reflecting the rationality of the algorithm. A summary of the computational resource requirements for each part is shown in Table 2.

[0328] Table 2 Pangu-Weather Computing Resource Requirements Overview

[0329]

[0330] 5. Conclusion

[0331] This embodiment systematically constructs a technical framework for evaluating the computational resources of large-scale meteorological models. It proposes an analytical method adapted to the spatiotemporal characteristics of meteorological data, considering four dimensions: parameter quantity, computational complexity, memory usage, and communication overhead. For different model architectures (such as U-Net, Swin-Transformer, and graph neural networks), parameter quantity calculation models are established, revealing the growth pattern of parameter quantity with the dimensionality of meteorological variables and network depth. In the computational complexity evaluation, a spatiotemporally aware FLOPs calculation framework is used to quantify the differences in resource requirements for convolution, attention mechanisms, and graph convolution operations. For example, the window attention mechanism of Swin-Transformer reduces the spatiotemporal correlation of FLOPs through local computation. The memory usage prediction model further distinguishes the memory bottlenecks in the training and inference stages, proving that intermediate activation value storage is the core source of memory pressure in large-scale meteorological models. Furthermore, by analyzing the coupling effect of data parallelism, model parallelism, and network topology in distributed training, the constraint mechanism of communication overhead on the efficiency of supercomputing clusters is clarified. A case study using the Pangu meteorological model verifies the universality of the framework, providing theoretical support for hardware adaptation and resource allocation of meteorological models.

[0332] In summary, this embodiment constructs a computational resource evaluation framework for large-scale meteorological models. By quantifying the model's parameter count, computational load, memory usage, and communication overhead, it provides a theoretical basis for hardware configuration and resource allocation, thereby reducing computational costs and ensuring the efficient and stable development and operation of large-scale meteorological models. Specifically, this embodiment proposes a Multi-Granularity Computing Resource Joint Evaluation Framework (MGCRJEF). This framework establishes a parameter count calculation model, a spatiotemporal awareness FLOPs evaluation model, a memory usage model, and a distributed communication analysis model in modules, and comprehensively evaluates the core hardware resource requirements of large-scale meteorological models by combining the spatiotemporal heterogeneity characteristics of meteorological data. This embodiment uses the Pangu-Weather meteorological model based on the Swin-Transformer architecture as a case study. The framework reveals the resource requirements of this model, such as increased memory usage with high-resolution input and communication overhead becoming a performance bottleneck in multi-node training. These findings provide practical guidance for resource optimization. Furthermore, the resource requirements calculated by the framework are basically consistent with the actual resource consumption, verifying the rationality and effectiveness of the framework. The MGCRJEF framework provides a standardized method for assessing resource requirements of large meteorological models, supports resource planning in intelligent computing hardware environments, and provides theoretical basis and practical reference for model deployment and hardware optimization in the meteorological field.

Claims

1. A method for assessing computational resources for a large-scale meteorological model, characterized in that, Includes the following steps: Step (1): Calculation and evaluation of meteorological large model parameters; In step (1), the total number of parameters in the Unet in the meteorological large model Calculation: ； In the formula, L enc The number of encoder layers. L dec The number of decoder layers. L skip The number of skip connection layers. The number of parameters for each encoder layer's convolutional layer. The number of parameters for the convolutional layer of each decoder layer. For the parameters of the batch normalization layer, The parameters of the skip connection layer; Calculation of parameters for the convolutional layers in the encoder and decoder layers: The number of parameters in a 2D convolutional layer : ; The number of parameters in a 3D convolutional layer : ； Parameters of batch normalization layer Calculation: ; Parameters of the skip connection layer Calculation: ; In the formula, Input the number of channels. Number of output channels The height of the convolution kernel. The width of the convolution kernel. The depth of the 3D convolution kernel; It also includes the total number of parameters of the Swin-transformer in the large meteorological model. The total number of parameters in the graph neural network in the calculation and meteorological big data model Calculation; Step (2), Calculation and evaluation of FLOPs in the large meteorological model; in step (2), the total FLOPs of Unet in the large meteorological model. Calculation: ； In the formula, The number of encoder layers. The number of decoder layers. FLOPs of the convolutional layers for each encoder layer FLOPs of the convolutional layers for each decoder layer For batch normalized layers of FLOPs; FLOPs calculation for convolutional layers in encoder and decoder layers: ； ； FLOPs calculation for batch normalized layers: ; In the formula, To output the high-resolution feature map, To determine the width of the output feature map, Input the number of channels. Number of output channels The height of the convolution kernel, The width of the convolution kernel, The length of the time dimension of the output feature map; The calculation of FLOPs for transposed convolutional layers is the same as that for ordinary convolutional layers, but the output size... , Determined by step size and fill size respectively; This also includes the total FLOPs of the Swin-transformer in the large meteorological model. The total FLOPs of graph neural networks in the computation and meteorological large model Calculation; Step (3): Calculation and evaluation of the memory usage of the large meteorological model; The memory usage of the large meteorological model includes the memory usage for training and the memory usage for inference; Step (4): Distributed communication evaluation of the meteorological big data model; In step (4), the distributed communication evaluation of the meteorological big model includes the following steps: Step (4-1), Analysis of data parallel communication overhead; Step (4-2), Evaluation of model parallel communication overhead; Step (4-3): Analysis of the impact of network topology on communication overhead; Step (4-4), total training communication overhead assessment; Steps (4-5): Latency and throughput assessment.

2. The computational resource assessment method for the large-scale meteorological model according to claim 1, characterized in that, In step (1), the total number of parameters of the Swin-transformer in the meteorological large model Calculation: = ； In the formula, The number of parameters for the Patch Embedding layer. For the number of stages, The number of parameters for the PatchMerging layer. The number of blocks in each stage, For the parameters of the Swing Transformer Block layer, The parameters of the output header; Number of parameters in the Patch Embedding layer: ; Number of parameters in the Patch Merging layer: ; Number of parameters in the Swin Transformer Block layer: ； Output header parameters: ; In the formula, C For the number of meteorological variables, p For Patch Size, choose 4×4 or 8×8. D Let r be the dimension of the hidden layers, and r be the scaling factor of the multilayer perceptron. This represents the number of output channels.

3. The computational resource assessment method for the large meteorological model according to claim 1, characterized in that, In step (1), the total number of parameters in the graph neural network of the meteorological large model is... Calculation: ； In the formula, This refers to the number of GCN or GAT layers in the graph neural network. For parameters of the GCN layer or GAT layer, The parameters for edge operations; Parameter calculation of GCN layer: ; Parameter calculation of GAT layer: ; Calculation of parameters for edge operations: ; In the formula, The input node feature dimension, To output the feature dimensions of the nodes, H The number of heads for multi-head attention. The input edge feature dimension, This is the dimension of the output edge features.

4. The computational resource assessment method for the large meteorological model according to claim 1, characterized in that, In step (2), the total FLOPs of the Swin-transformer in the large meteorological model are... Calculation: ； In the formula, For the FLOPs of the Patch Embedding layer, For the number of stages, The number of blocks in each stage, For FLOPs of the multi-head self-attention mechanism layer of windows, For FLOPs of a multilayer perceptron, For the FLOPs of the Patch Merging layer, To output the FLOPs of the head layer; Calculation of FLOPs in the Patch Embedding layer: ； Calculation of FLOPs in the multi-head self-attention mechanism layer of the window: ； Calculation of FLOPs for a multilayer perceptron layer: ; FLOPs calculation for the Patch Merging layer: ; Calculation of FLOPs in the output head layer: ; In the formula, C For the number of meteorological variables, To output the high-resolution feature map, To determine the width of the output feature map, D For the hidden layer dimension, Number of output channels N For the number of windows, This refers to the window size.

5. The computational resource assessment method for the large meteorological model according to claim 1, characterized in that, In step (2), the total FLOPs of the graph neural network in the meteorological large model are... Calculation: ； In the formula, This refers to the number of GCN or GAT layers in the graph neural network. For FLOPs of GCN or GAT layers, FLOPs for edge operations; Calculation of FLOPs in the GCN layer: ； Calculation of FLOPs in the GAT layer: ； Calculation of FLOPs for edge operations: ; In the formula, For the number of nodes, Let the number of edges be denoted as '('). The input node feature dimension, To output the feature dimensions of the nodes, H The number of heads receiving multi-head attention.

6. The computational resource assessment method for the large meteorological model according to claim 1, characterized in that, In step (3): Calculation of training memory usage: ； Calculation of inference memory usage: ; In the formula, P For parameters, b Storage precision, S For sequence length, D For the hidden layer dimension, L This represents the number of model layers.

7. The computational resource assessment method for the large meteorological model according to claim 1, characterized in that, In step (4-1): Total communication time ; In the formula, For model runtime computation time, The overlap ratio, Total communication time; ； ； In the formula, P For the number of model parameters, b For parameter accuracy, B The number of buckets into which the total gradient is divided; Effective bandwidth is the actual available network bandwidth; latency includes network transmission delay, protocol overhead, and synchronization wait time; hop count is the number of intermediate nodes in the data transmission path. In step (4-2): assuming the embedding dimension is... D The sequence length is S The number of GPUs is Then the communication volume at each layer is ; Bubble Time This refers to the idle time the equipment spends while waiting for other stages to complete. ,in This represents the number of stages in the model. For single-stage calculation time; ; The average stage time, For the first The computation time for each stage; Total communication time Calculation: ,in For cross-stage communication time, , For the first The first stage to the second The amount of data to be transmitted in each stage; This refers to the quantity in a micro-batch. ; In step (4-4): the total communication overhead of the meteorological large model Calculation: ； In steps (4-5), the single inference delay Including calculation time and data transmission time ； Throughput tokens / s calculation: ； The maximum token length. This represents the number of parallel devices.

Images