Adapting a Trillion-Parameter Sparse MoE Model for Inference on Domestic AI Accelerators

1. Background: Inference Challenges of Ultra-Large Sparse Models

As large model parameter counts evolve toward the trillion scale, their capabilities continue to improve across scenarios such as code generation, long-horizon agent tasks, knowledge-based Q&A, financial research, and 3D content creation. At the same time, the challenges on the inference side have grown substantially. An increase in parameter count means a sharp expansion in the size of weight files—single-GPU memory is no longer sufficient to hold all parameters, making multi-node, multi-GPU coordination the norm. Sparse activation structures reduce the amount of computation in a single forward pass, but they place higher demands on weight distribution, expert routing, and communication efficiency.

In this context, whether a hardware platform can complete basic inference adaptation as soon as model weights are released has become a key indicator of the maturity of its software ecosystem. The industry typically refers to this “ready at release” capability as Day-0 support. Day-0 support does not merely mean running a simple example—it requires stable completion of model loading, service startup, and the full basic inference path without modifying the core model architecture. This provides a reliable foundation for subsequent performance tuning and production deployment.

2. Adaptation Objectives and Overall Approach

The target of this adaptation effort is a sparse Mixture-of-Experts (MoE) model with three trillion parameters. Compared with previous dense or smaller sparse models, this model differs significantly in the following aspects:

1. Massive parameter scale: The full weights require a large amount of storage space. The loading process must account for distributed file reading and multi-level caching mechanisms.
2. MoE architecture: During the forward pass, only a subset of experts is dynamically selected and activated based on input tokens. The weights of unselected experts are not involved in the current computation, requiring an inference engine with efficient expert routing and weight swapping capabilities.
3. Dual-modality input: The model accepts both text and image inputs. The inference pipeline must handle text encoding and image encoding separately, then perform cross-modal fusion at specific layers.
4. Long-context capability: The model supports a context window of up to one million tokens, imposing additional requirements on positional encoding, cache management in the attention mechanism, and memory allocation strategies.

Given these characteristics, the core strategy for this adaptation work can be summarized as follows: use a general-purpose software stack as the foundation to migrate and adapt the critical inference paths; use the AI training-inference integrated accelerator as the hardware base to validate the complete flow from model loading to service response; and target correct and reproducible basic inference for both text and image modalities.

3. The Role of the Software Stack in Inference Adaptation

In large model inference adaptation, the software stack plays a bridging role. Upward, it must remain compatible with model representations and operator sets from mainstream deep learning frameworks. Downward, it must efficiently map computational tasks onto the accelerator’s compute units and memory hierarchy.

A complete inference adaptation process typically includes the following layers:

3.1 Model Representation and Graph Optimization

The first step is converting pretrained weights from a general-purpose format into an internal representation recognizable by the inference engine. This stage involves tensor parallelism partitioning, expert weight grouping, computational graph rewriting, and redundant operator elimination. For sparse MoE models, the organization of expert weights significantly impacts subsequent loading efficiency. A reasonable grouping strategy can reduce the number of weight exchanges during inference and lower end-to-end latency.

3.2 Operator Adaptation and Precision Alignment

Accelerators differ from general-purpose GPUs in instruction sets, thread scheduling, and numerical precision behavior. During adaptation, core operators—including attention computation, feed-forward networks, layer normalization, and activation functions—must be validated layer by layer to ensure that the output error relative to the reference implementation remains within an acceptable range. Precision alignment is a prerequisite for correctness validation and the basis for subsequent performance tuning.

3.3 Memory Management and Communication Optimization

Inference with trillion-parameter models demands far more memory than a single card can provide. In practice, a combination of multi-card tensor parallelism and pipeline parallelism is typically used. The adaptation phase must focus on verifying the correctness and stability of communication primitives across different parallelism modes, while also designing a sound memory pooling mechanism to reduce fragmentation and improve memory reuse. In long-context scenarios, the KV cache management strategy directly affects the maximum sequence length that can be served and the level of concurrent throughput.

4. Hardware Platform and Dual-Modality Inference Pipeline

The hardware platform used in this adaptation is an AI training-inference integrated accelerator. This type of hardware supports both training and inference workloads on a single card, offering high compute density and large memory capacity, making it suitable as a validation and deployment platform for large model inference.

4.1 Model Loading Pipeline

Model loading is the process of reading weight files from persistent storage into accelerator memory. For a three-trillion-parameter model, the full weight set typically reaches hundreds of gigabytes or more. The loading pipeline must support the following capabilities:

• Distributed parallel reading: Fetch weight shards in parallel from multi-node storage to shorten load time.
• Weight conversion and validation: Perform data format conversion and integrity checks after reading.
• On-demand loading strategy: For sparse expert weights, lazy loading can be used—experts are loaded into compute units only when actually triggered during inference.

In practice, the adaptation first validated model loading on a single-node, multi-card environment, then extended to multi-node scenarios, confirming that the loading path remains stable under distributed conditions.

4.2 Service Startup Process

After model loading, the model is wrapped into a runtime instance that can serve inference requests externally. The service startup process typically includes the following steps:

Initialize the inference engine
Load the model configuration file
Allocate memory and compute resources
Load weights and complete parallel partitioning
Build the inference computation graph
Start the HTTP/RPC service port
Run health checks and warm-up inference

The adaptation phase focused on the stability and repeatability of service startup. Repeated startup tests showed that, with reasonable configuration, the model initializes within the expected time and enters a ready-to-serve state. Health check endpoints returned normal results, and warm-up inference requests produced correct outputs.

4.3 Text Inference Pipeline

Text inference is the most fundamental use case. Input is tokenized into a sequence of tokens, then passed through the embedding layer, multiple Transformer encoder layers, and the output layer to generate response tokens. For MoE models, the feed-forward network in each layer activates a subset of experts based on token routing results.

During text inference adaptation, the following aspects were verified:

• Consistency between short and long text inputs;
• Reasonableness of output under different decoding parameters;
• Correctness of both streaming and non-streaming output modes;
• Concurrent stability under batched requests.

Test results confirmed that the model performs text inference correctly on the accelerator, with outputs semantically consistent with the reference implementation.

4.4 Image Inference Pipeline

The image inference pipeline is comparatively more complex. The input image first passes through a vision encoder to extract features, which are then fused with text features at specific layers before entering the unified multimodal decoding process. Image inference places higher demands on memory bandwidth and computational precision, especially in high-resolution input scenarios where the computational load of the vision encoding section increases significantly.

During image inference adaptation, the correctness of image preprocessing, visual feature extraction, cross-modal projection, and fusion decoding was validated. Since the vision encoder and text encoder differ in architecture, the adaptation work required handling the computation mapping of the two subnetworks separately while ensuring that tensor dimensions match at the cross-modal fusion layers.

5. Adaptation Validation and Basic Evaluation

After completing the pipeline, a set of standard evaluations was used to verify correctness and basic usability. The evaluation is divided into functional validation and stability validation.

Functional validation focuses on whether model output quality meets expectations. For text inference, a standard set of Q&A pairs was used to check the relevance and logical consistency of responses. For image inference, mixed inputs containing both text and images were used to verify that the model correctly understands and responds. Evaluation results confirmed that the model’s inference outputs on the accelerator meet expectations without systematic deviations.

Stability validation focuses on behavior under sustained operation and concurrent load. Continuous running tests covered input sequences of varying lengths, different batch sizes, and alternating text and image requests. Throughout the test period, the service remained stable, memory usage fluctuated within a reasonable range, and no memory leaks or computational anomalies were observed.

6. The Significance of Day-0 Support and Future Work

The core value of Day-0 support lies in reducing the time between model release and production usability. For enterprise users and developers, whether a model can run on their own hardware platform immediately after weight release directly affects technology selection and business planning. Especially during the critical stage of building the domestic accelerator ecosystem, rapidly completing adaptation validation for mainstream large models helps build a more complete software-hardware co-optimization capability.

The current adaptation achievements focus on opening up the basic inference path and validating stability. Future work should continue in the following directions:

1. Deep performance optimization: Optimize operator implementations based on the computational characteristics of MoE models, improve compute utilization, and reduce end-to-end latency.
2. Long-context validation: Conduct stress tests at the million-token context window and optimize KV cache management and memory allocation strategies.
3. Multi-node scaling validation: Further validate inference scalability and communication efficiency on larger clusters.
4. Business scenario adaptation: Develop deployment solutions and engineering practices tailored to practical scenarios such as knowledge Q&A, code generation, and financial research.

Through continuous adaptation and iterative optimization, the inference capability of ultra-large sparse models on domestic accelerators will gradually progress from “usable” to “production-ready,” providing reliable technical support for a wider range of real-world business scenarios.