Adapting a Trillion-Parameter Open-Source Model: From Inference Migration to Stable Deployment

1. Background and Challenges: Crossing the Trillion-Parameter Threshold

A large language model with 3 trillion parameters was recently open-sourced. As the first model in its series, it demonstrates capabilities approaching or matching top-tier closed-source models across multiple dimensions, including code generation, 3D content creation, financial research, knowledge-based question answering, and long-horizon agent tasks. For instance, it scores 65.7% on a specific terminal command parsing benchmark and 61.2% on a long-context code understanding benchmark, showing significant improvements over its predecessors.

However, such a leap in model capability comes with a steep engineering curve. What does 3 trillion parameters actually mean? Stored in FP16 precision, the model weights exceed 6 TB. Even with current mainstream 8-card high-speed interconnect servers, loading the full weights takes minutes. The more critical challenges are:

• Memory Wall: The limited memory capacity of a single accelerator necessitates tensor parallelism, pipeline parallelism, or multi-node clusters to hold the model.
• Communication Wall: During inference with a trillion-parameter model, every generated token requires synchronizing all tensor-parallel shards, making communication overhead a frequent throughput bottleneck.
• Long-Context Pressure: Supporting a 1-million-token context window means the computational complexity of attention grows quadratically, while the Key-Value (KV) cache consumes enormous memory space.
• Software Stack Adaptation: New hardware architectures require mature compiler and runtime support to efficiently map the model’s mathematical operators onto the underlying instruction set.

This article uses the open-source adaptation of this model as a case study to explore the technical path and engineering considerations for achieving end-to-end deployment on domestic AI acceleration hardware.

2. Model Architecture and Computational Characteristics

Before starting the adaptation, a deep understanding of the model’s structural features dictates the direction of subsequent optimizations. The 3-trillion-parameter model employs the industry-standard Mixture-of-Experts (MoE) architecture, using a sparse activation mechanism to control the computational cost per inference while maintaining the model’s overall capacity. Its core characteristics can be summarized as follows:

• Massive-Scale MoE Layers: Most of the feed-forward networks in the model are replaced with multiple parallel expert modules, with only a few experts activated per inference. While this significantly reduces computation, it requires the inference engine to have efficient expert routing and load-balancing capabilities.
• Long-Sequence Attention Optimization: To support a 1-million-token context, the model incorporates sparse attention patterns (e.g., combining sliding windows with global anchor tokens) in its attention layers. This requires the underlying libraries to provide efficient operator support for these non-standard attention patterns.
• Post-Training with Reinforcement Learning: The model underwent extensive RLHF (Reinforcement Learning from Human Feedback) and online policy optimization, resulting in output distribution characteristics that differ from purely pre-trained models. This requires recalibrating sampling parameters (such as temperature and top-p) for online inference.

3. Core Adaptation: Migration and Restructuring of the Inference Path

Based on the model characteristics above, the first step in adaptation is to complete the migration of critical inference paths. This process is not a simple model file format conversion but a systematic engineering effort across three layers: computational graphs, operator mapping, and memory management.

3.1 Computational Graph Parsing and Partitioning

The original model is released in a general-purpose framework format, containing hundreds of Transformer layers and thousands of expert modules. We first use the deep learning compiler’s frontend parser to load the model as an intermediate representation (IR), then perform the following key operations:

1. Operator Fusion: Fuse adjacent linear layers and normalization layers to reduce kernel launch overhead. For example, fuse “column-parallel matrix multiplication + bias + LayerNorm” into a single kernel.
2. Expert Parallel Partitioning: Distribute the experts in the MoE layers evenly across compute chips according to a preset parallelism degree. A strategy called “grouped shuffling” is used to ensure that routed experts are likely located on local or near-neighbor nodes, reducing cross-chip communication.
3. Pipeline Stage Partitioning: Divide the 200+ Transformer blocks into multiple pipeline stages and insert point-to-point communication primitives (e.g., send/recv) between stages. This allows different chips to asynchronously and parallelly execute forward computations for different layers.

3.2 Operator Mapping Strategy on the MUSA Software Stack

The domestic accelerator used for adaptation follows the MUSA (a general-purpose GPU-like computing architecture) software standard. This stack abstracts away the underlying chip microarchitecture differences, providing CUDA-like programming interfaces and high-performance operator libraries. The adaptation relied on three levels of this stack:

• Standard Operator Library Coverage: Basic operators involved in attention computation, such as matrix multiplication, Softmax, and scaled dot-product, already have highly optimized implementations in the library. The key is to automatically select the optimal tiling strategy (e.g., blocking a 128×128 matrix multiply into 64×64 tiles) based on the actual tensor shapes.
• Custom Kernel Complementation: Operations like “Top-k selection + data movement” in MoE routing and “sliding window mask” generation for sparse attention lack ready-made implementations in the standard library. We wrote several compute kernels using MUSA’s low-level inline assembly interface. For example, the sparse attention mask generation kernel only generates mask bits for the window positions that require computation, avoiding the memory waste of a full mask tensor.
• Just-In-Time Compilation and Caching: For model-specific fused operators (e.g., “expert gating + weighted sum”), we use runtime compilation techniques to generate code dynamically and cache the compiled artifacts on disk to expedite subsequent loading.

4. Parallel Strategies and Memory Optimization: Fitting and Accelerating 3 Trillion Parameters

The physical foundation for model adaptation is memory. With 80 GB of memory per card, purely data-parallel approaches with no memory optimization would require at least 80 cards—infeasible for production environments. Therefore, a combined, multi-pronged memory management strategy is essential.

4.1 Hybrid Parallelism: Synergy of Tensor and Expert Parallelism

Given the MoE structure of the model, a hybrid parallel approach is commonly used in the industry:

• Attention Layers: Employ Tensor Parallelism (TP) by splitting the attention weight matrices along row/column dimensions across cards within a TP group. All-to-All communication happens over high-speed interconnects within the group.
• MoE Layers: Switch to Expert Parallelism (EP). Each card holds only a replica of a subset of experts. The routing mechanism distributes all tokens to the card hosting their corresponding expert for computation. This mode reduces communication from a full tensor broadcast per layer (as in TP mode) to a sparse broadcast of only the routed tokens (as in EP mode), significantly improving communication efficiency.

4.2 Trading Compute for Memory: Mixed Precision and Quantization

To accommodate weights exceeding 6 TB in limited memory, a mixed-precision strategy is necessary:

• Weight Storage: Core weights are stored in BF16 (Brain Floating Point 16) format, halving the memory footprint compared to FP32.
• KV Cache Quantization: During long-context inference, the KV cache expands rapidly. We quantize it from FP16 to INT8 and temporarily restore precision via dequantization during attention computation. Our tests show this reduces KV cache memory usage by 40% with negligible impact on final generation quality.
• Activation Recomputation: In certain deep stages (during forward inference or if end-side fine-tuning is involved), intermediate activations are not kept in memory but recomputed when needed. This adds roughly 20% redundant computation but significantly lowers activation memory peaks.

5. Stability Engineering for Long Contexts

Supporting a 1-million-token context is a core selling point of the model and the most error-prone part of the adaptation process. Stability in long contexts covers three aspects: memory fluctuation, numerical precision drift, and inference latency jitter.

• Memory Circuit Breaker Protection: For extremely long inputs, we dynamically monitor the memory pool watermark for each layer at runtime. If it approaches the physical memory limit, the layer’s computation is automatically switched to a “blocked streaming” mode—processing attention not for the full sequence at once, but by loading context in blocks and updating the KV cache in a rolling manner.
• Numerical Stability Handling: When sequence length exceeds a certain threshold (e.g., 200K tokens), the variance of attention scores tends to increase, leading to Softmax saturation. We introduce an adaptive scaling factor adjustment mechanism and enforce FP32 accumulators for summation in critical attention operators to avoid precision divergence from low-precision arithmetic.
• Long-Horizon Agent Task Persistence: To address the frequent interruption of long-horizon agent tasks in practical testing, we added an intelligent checkpoint-resume mechanism at the outer layer of the inference service. This mechanism saves the current KV cache snapshot to disk after generating a complete response (not per token). If an upper-layer application request is interrupted or an abnormal signal like high chip temperature is detected, it can automatically load the most recent snapshot and continue generating. This ensures continuity for time-consuming tasks like financial research and multi-turn code debugging.

6. Performance Tuning and Recommendations

After completing the foundational adaptation, we measured single-card throughput and multi-card scaling efficiency. The following tuning guidelines are summarized based on our practice:

Configuration Item	Recommended Initial Value	Tuning Strategy
Tensor Parallel Size	4	Can be increased to 8 if interconnect bandwidth is above 400 GB/s; keep at 2 if below 200 GB/s.
Expert Parallel Size	Equal to the number of physical cards	Since the number of experts far exceeds card count, no further splitting is needed.
Batch Size	1 (for long sequences)	Can be increased to 8 for short sequences (<4K) to boost throughput.
KV Cache Quantization	INT8	For precision-sensitive tasks, downgrade to FP16 and add a sliding window mean compensation layer.
Communication Backend	Polling mode	Switch to interrupt mode in high-concurrency scenarios to reduce CPU usage.

In practice, we observed a single-card generation throughput of 12 tokens/s at a sequence length of 32K. When scaling to a 64-card cluster handling a 128K context, the overall throughput degradation was kept under 20%, validating the effectiveness of the parallel strategy.

7. Conclusion and Outlook

Hardware adaptation for ultra-large-scale models is no longer a simple “install and run” exercise. It is a complex systems engineering effort involving co-design across compilers, parallel computing, system architecture, and algorithms. The Day-0 adaptation of this 3-trillion-parameter model on AI training/inference cards and the MUSA software stack demonstrates that the domestic software and hardware ecosystem’s capability to support cutting-edge models is rapidly advancing.

Looking ahead, as model scales move towards ten trillion parameters and beyond, the marginal benefits of purely manual tuning will diminish. The industry will need more intelligent automated compilation pipelines—ones that can automatically derive optimal parallel strategies from model structure, automatically generate fused operators suited to hardware characteristics, and even dynamically sense network congestion at runtime and adjust communication routes. This will be the core battleground for AI infrastructure competition in the next phase.