Day-0 Adaptation: Running a 3-Trillion Parameter MoE Model on Domestic GPUs

    This article explores the engineering challenges and solutions in adapting a 3-trillion parameter sparse Mixture-of-Experts (MoE) model on a domestic AI accelerator. It covers the technical hurdles of massive parameter storage and communication bottlenecks, detailing the hardware architecture, software stack (MUSA), and the end-to-end Day-0 adaptation process from weight loading to inference. Key optimization strategies like communication-computation overlap and PagedAttention are also discussed, offering actionable insights for AI infrastructure engineers.

    Background: Engineering Challenges of Ultra-Large-Scale Models

    As large language model parameter scales surpass the trillion-level mark, AI infrastructure is facing unprecedented engineering challenges. Taking the recently open-sourced 3-trillion parameter sparse Mixture-of-Experts (MoE) model as an example, its actual deployment involves memory footprints, complex routing computations, and cross-device communication overheads far exceeding those of traditional dense models. For any AI compute provider, the ability to complete adaptation on Day-0—the exact day model weights are officially released—has become a core metric for measuring the maturity of its software stack and the foresight of its hardware design.

    Day-0 adaptation refers to the complete process where, at the exact moment model weights are publicly released, the hardware platform can immediately load the model, start inference services, and verify basic capabilities. This not only tests the memory capacity and compute precision of a single GPU but also challenges the scheduling efficiency and fault tolerance of the software stack for large-scale distributed inference. This article breaks down the key engineering phases of deploying ultra-large-scale models based on a recent adaptation case completed by a domestic GPU vendor, exploring its implications for autonomous AI infrastructure development.

    Sparse Expert Models: A Dual Test of Memory and Compute

    Current mainstream trillion-level parameter models predominantly adopt the sparse Mixture-of-Experts (MoE) architecture. Unlike dense models that activate all parameters during every inference, MoE models selectively activate a portion of expert networks via a gating routing mechanism. This design allows the total parameter count to far exceed the memory limit of a single card, but it also introduces two core technical challenges:

    1. Parameter Storage and Dynamic Loading: Complete model weights can reach tens of TB (calculated in FP8 precision), far exceeding single-node memory limits. Distributed tensor parallelism and expert parallelism strategies must be employed to slice and assign different expert layers across compute nodes. This requires the software stack to feature automated model sharding tools capable of generating optimal partitioning schemes based on actual cluster topologies.
    2. Routing Computation and Communication Bottlenecks: Every token needs to be routed to its corresponding expert’s device via all-to-all communication at each layer. When the number of experts reaches hundreds, communication volume grows exponentially. Traditional PCIe buses or Ethernet interconnects quickly become bottlenecks, necessitating high-speed NVLink-like buses or dedicated RDMA networks. Concurrently, the software stack must achieve overlap between communication and computation.

    Hardware Foundation: Architecture of the AI Compute Card

    In this adaptation case, the hardware platform utilized an AI compute card designed for both training and inference scenarios, equipped with an independently architected GPU core. Based on public information, the card possesses several architectural features closely related to ultra-large model adaptation:

    • High-Capacity, High-Bandwidth Memory: The single-card memory capacity reaches tens of GBs and is equipped with HBM2e or HBM3-class high-bandwidth memory. This ensures that a large proportion of expert parameters can reside on a single card, reducing cross-node access frequency. Furthermore, high bandwidth (typically exceeding 1TB/s) is a prerequisite for maintaining high-throughput inference for MoE models.
    • Flexible Tensor Cores: Supports multiple precision modes including FP16, BF16, and FP8. In inference scenarios, FP8 precision can double the effective memory capacity of a single card while increasing compute throughput several times. During adaptation, precision sensitivity analysis must be conducted on critical operators (such as matrix multiplication and normalization) to prevent output quality degradation caused by quantization.
    • MUSA Software Stack: Unified Programming and Scheduling Framework: This is the software foundation of the entire adaptation effort, consisting of three core layers:
      • MUSA Toolkit: Provides programming interfaces similar to the industry-standard CUDA, enabling developers to write high-performance kernels in C/C++ and achieve source-code-level compatible migration.
      • MUSA Libs: Contains high-performance libraries for General Matrix Multiply (GEMM), convolution, and normalization. Autonomous driving-level tuning ensures these operators perform no worse than mature industry solutions.
      • Distributed Framework Plugins: Provides adaptation layers for mainstream open-source AI frameworks (such as PyTorch and TensorFlow), achieving efficient collective communication required for multi-card parallelism via a unified communication library.

    The Day-0 Adaptation Process: From Weight Loading to Inference Service

    Completing the Day-0 adaptation for a 3-trillion parameter model is far from a simple environment configuration; it is a comprehensive systems engineering effort involving model parsing, topology-aware sharding, precision validation, and performance tuning. The entire process can be broken down into the following four core phases.

    Phase 1: Model Parsing and Format Normalization

    Open-source models are typically released in snapshot formats (such as Safetensors). The primary task is to develop a dedicated parser to losslessly convert model weights into the native storage format of the MUSA software stack. This phase requires careful handling of the following details:

    • Sparse Weight Mapping Table: The weight storage in MoE models is sparse; it is necessary to identify the index relationships of expert lists and shared experts in storage.
    • Metadata Consistency Check: Ensure that the expert indices, layer numbers, and computation graph definitions corresponding to each sharded weight block are completely consistent.

    Phase 2: Distributed Inference Scheme Design

    Based on the cluster topology (e.g., 8 cards per server), the software automatically generates distributed inference configuration files. A typical scheme employs a hybrid strategy of “Tensor Parallelism + Expert Parallelism”:

    • For shared components like attention mechanisms, tensor parallelism is used to slice matrices by row/column and distribute them across multiple cards for collaborative computation.
    • For MoE expert layers, an expert parallelism strategy is used, distributing different experts across different cards and utilizing all-to-all communication to complete token routing.

    Command-line launch schematic (presented as pseudocode without specific brands):

    1. # Model loading and distributed inference script example
    2. musa_launch --nnodes 4 --nproc_per_node 8 \
    3. --model_path /data/kimi-k3-pro \
    4. --tensor_parallel_size 8 \
    5. --expert_parallel_size 4 \
    6. --dtype float8 \
    7. --use_flash_attention

    Phase 3: Memory Optimization and Loading Acceleration

    Facing 3 trillion parameters, the original loading strategy (direct full mapping) would result in minute-level startup latency. Optimization strategies include:

    • Asynchronous Pipeline Loading: Parameters are split into smaller batches by layer, read simultaneously from NVMe storage via multi-threading, and written directly to memory, hiding I/O latency.
    • Dynamic Memory Allocator: The default PyTorch caching allocator is replaced with a customized weight preloading strategy. This prioritizes retaining high-frequency expert parameters (similar to a cache replacement algorithm) and evicting low-frequency parameters to host memory, addressing dynamic memory pressure during large context windows.

    Phase 4: Precision Calibration and Inference Verification

    After enabling basic text and image inference pipelines, strict numerical precision comparisons must be conducted. Specific methods include:

    • Using FP32 high-precision inference results as Ground Truth to compare the cosine similarity of FP8/FP16 inference results (threshold set above 0.99).
    • Verifying routing consistency specific to MoE models—ensuring that under FP8 precision, the expert selection provided by the gating network remains at least 95% consistent with the selection under FP32 mode, otherwise inference result drift will occur.

    Key Bottlenecks and Tuning Strategies

    During this adaptation process, the engineering team resolved several critical runtime performance bottlenecks. These experiences offer direct reference value for similar projects.

    Bottleneck 1: All-to-All Communication Blocking

    When the inference batch size is large, tokens are evenly distributed across hundreds of expert cards, requiring each GPU to communicate with all other GPUs in the cluster. This instantly saturates interconnect bandwidth, leaving GPU compute units idle for extended periods.

    Tuning Solution: Employing communication-computation overlap technology—subdividing MoE layer computations into “local” and “remote” components. While local experts process a portion of tokens, the background DMA engine asynchronously prepares the transmission and reception of remote tokens. By utilizing the event stream mechanism provided by the MUSA software stack, communication operations and kernel execution occur concurrently, which empirically improved GPU utilization by approximately 42%.

    Bottleneck 2: KV Cache Explosion Caused by Sequence Length

    Even with high memory capacity, the KV Cache can occupy tens of GBs of memory when inferring long text sequences (e.g., exceeding 128K tokens). This requires the inference service to possess memory management capabilities from a PagedAttention perspective.

    Implementation Mechanism: The KV Cache is divided into fixed-size blocks connected via index linked lists, enabling dynamic utilization of non-contiguous memory spaces. Implementing this mechanism requires rewriting the attention kernel, making it one of the most complex parts of the inference engine.

    Implications and Outlook for AI Infrastructure

    The successful adaptation of the 3-trillion parameter model demonstrates that domestic AI hardware and software are now capable of competing with mainstream international solutions on key technical metrics. The implications are multidimensional:

    1. Strategic Importance of the Software Stack: The upper limit of hardware performance is dictated by software stack maturity. The three-layer architecture of the MUSA software stack—“unified programming interface + high-performance operator library + distributed plugins”—effectively lowers the migration barrier for developers, making Day-0 adaptation possible.
    2. Ecosystem Compatibility is Essential: The current large model ecosystem is concentrated on a few open-source frameworks. Evaluating a hardware platform’s value relies heavily on its transparent replacement support for these frameworks. Through interface-level compatibility, developers can switch hardware without modifying underlying training scripts, significantly reducing migration risk.
    3. The Leap from “Usable” to “Optimized”: The completeness of Day-0 adaptation is not just about running the model, but achieving high performance and stability. The precision calibration strategies and communication overlap optimizations in this practice demonstrate a pragmatic attitude toward pursuing卓越 performance.

    Moving forward, as model parameters progress toward the 10-trillion mark, compute cluster requirements will extend beyond single GPU specifications to focus more on large-scale cluster linear scalability, automatic fault recovery, and green energy consumption control. If domestic compute solutions continue to invest heavily in interconnect technology (e.g., optical interconnects) and adaptive inference engines, they are poised to gain the upper hand in the next phase of competition.

    For AI developers and enterprise architects, this adaptation case provides a reusable methodology: when introducing a new hardware stack, prioritize verifying whether it possesses end-to-end observability tools (such as operator profilers and memory profilers) and comprehensive debugging diagnostic tools, which are often the core elements for accelerating problem identification and reducing tuning costs.