Engineering Practices and Optimizations for Unified Training and Inference Adaptation of Ultra-Large Models

1. Background: The Continuing Growth of Large Model Parameter Scales

In recent years, the parameter scales of large language models and multimodal models have continued to rise, evolving from tens of billions to hundreds of billions and gradually toward the trillion-level. This growth not only increases model capacity but also imposes higher requirements on training and inference infrastructure. For engineering teams, a common challenge is: when the model weights themselves already exceed the memory capacity of a single card, how do you accomplish loading, sharding, inference, and service deployment?

In this context, adaptation work should not stop at the level of “can it run” but should address the complete engineering pipeline, including:

• Model weight format parsing and conversion;
• Multi-card memory planning and distributed loading;
• Selection of inference parallelism strategies;
• Service startup and request scheduling;
• Connecting inference pipelines for different inputs such as text and images;
• Optimization of latency, throughput, and stability.

This article uses trillion-level sparse or mixture-of-experts models as reference objects and discusses general adaptation methods on unified training-and-inference compute cards and unified software stacks. The hardware and software involved are described in a neutral manner and do not point to specific vendors or models.

2. Core Difficulties in Ultra-Large Model Inference Adaptation

1. Weight Scale Far Exceeds Single-Card Memory

Taking a three-trillion-parameter model as an example, even with lower-precision storage—for instance, 4 bytes or fewer bits of quantization per parameter—the raw weight volume may reach several TB to over ten TB. The memory of a single compute card typically ranges from tens of GB to around a hundred GB, so the complete weights cannot be placed into a single card at once.

Practical deployment needs to address:

• How weights are distributed across multiple cards;
• How to avoid memory overflow during the loading phase;
• How to load weight shards on demand during inference;
• How to reduce communication overhead caused by parameter transfer.

2. Complex Model Structures

Ultra-large models are usually not simply stacked Transformer layers; they may adopt mixture-of-experts structures, multimodal encoders, sparse attention, or grouped-query attention designs. Adaptation cannot consider only the generic language model path but must also handle:

• Preprocessing of inputs from different modalities;
• Activation and computation of different expert branches;
• Output alignment between image encoders and text encoders;
• Switching between different task heads.

3. High Requirements for Inference Services

Production environments impose far higher requirements on inference services than offline scripts. A service needs:

• Fast startup capability;
• Stable concurrent processing capability;
• Controllable memory usage;
• A clear error-return mechanism;
• Observable logs and metrics.

Therefore, adaptation work must extend from “can infer” to “can serve.”

3. General Adaptation Workflow Design

On unified training-and-inference compute devices and unified software stacks, the adaptation workflow can be divided into the following stages.

1. Model Format Identification and Conversion

First, the model weight format needs to be identified. Common formats include:

• Sharded weight files;
• Weights after tensor-parallel splitting;
• A directory structure containing configuration, tokenizer, and weight index.

During adaptation, weights usually need to be converted into a format recognizable by the current software stack. During conversion, tensor names, dimensions, and precision information should be kept consistent to avoid loading failures caused by naming inconsistencies.

A general approach is: first parse the model configuration to confirm parameters such as the number of layers, hidden dimension, number of attention heads, number of experts, and vocabulary size; then build the model structure according to the configuration; finally map the weights layer by layer to the corresponding modules.

2. Multi-Card Memory Planning

Before loading an ultra-large model, it is necessary to plan in advance which weight shards are stored on each card. Common strategies include:

• Layer-wise splitting: different cards store different Transformer layers;
• Tensor-wise splitting: weight matrices within the same layer are split by rows or columns;
• Hybrid strategy: splitting both across layers and within layers.

For example, suppose a model has 128 layers and is configured with 16 cards. A simple scheme is for each card to manage 8 layers; another scheme is for each card to store part of the tensors of all layers. The former is simple to implement but may cause load imbalance; the latter has more complex communication but helps improve parallelism.

Memory planning also needs to reserve space for activations, input data, output buffers, and temporary tensors. It is recommended to use small inputs for memory probing in the early stages of adaptation and gradually approach a stable configuration.

3. Sharded Weight Loading

To avoid memory overflow during loading, streaming loading or memory mapping is recommended. Specific steps include:

1. Read the weight index to determine the parameter name corresponding to each shard;
2. Map parameters to target devices according to the parallelism strategy;
3. Load shard by shard and immediately transfer to the corresponding device;
4. Release host-side temporary caches after loading is complete.

A pseudocode illustration is as follows:

for shard in weight_shards:
    tensors = load_shard(shard)
    for name, tensor in tensors.items():
        target_device = map_to_device(name)
        tensor = tensor.to(target_device)
        model.load_parameter(name, tensor)
    release_host_cache(tensors)

This approach can effectively control peak host memory and device memory usage.

4. Parallel Inference Strategy Configuration

During inference, parallel strategies need to be selected based on model scale and hardware topology. Common strategies include tensor parallelism, pipeline parallelism, and expert parallelism.

• Tensor parallelism is suitable for splitting single-layer weights across multiple cards to reduce per-card memory usage;
• Pipeline parallelism is suitable for distributing different layers across different devices;
• Expert parallelism is suitable for mixture-of-experts structures, routing different experts to different devices.

In practical adaptation, the three strategies are often combined. For example: tensor parallelism for attention layers, expert parallelism for feed-forward expert layers, and pipeline parallelism between different stages.

When configuring parallel strategies, communication cost must be considered. In a single inference pass, tensor parallelism may introduce multiple collective communications; pipeline parallelism may introduce batch padding overhead; expert parallelism depends on routing results and may cause load imbalance. Therefore, splitting granularity needs to be adjusted based on measured results.

5. Service Startup and Request Handling

After model loading is complete, inference capability needs to be encapsulated as a service. The service startup process usually includes:

• Initializing the runtime environment;
• Reading the model configuration;
• Loading weights;
• Warming up inference;
• Binding ports and receiving requests.

The request handling pipeline should preferably adopt a queue and scheduling mechanism: after receiving a request, perform validity checks first, then enter the inference queue; return the result after inference completes. For multimodal models, both text and image inputs need to be supported. Image inputs generally pass through an image encoder to produce features, which then enter the subsequent network together with text inputs.

A simplified request handling flow is as follows:

receive request -> parameter validation -> text tokenization -> image preprocessing -> encoding -> joint inference -> result decoding -> return response

6. Inference Pipeline Validation

Day-0 adaptation usually aims to get the basic pipeline working. Validation content may include:

• Whether text generation works normally;
• Whether image inputs can be correctly encoded;
• Whether multimodal inputs can produce reasonable outputs;
• Whether long text inputs trigger memory anomalies;
• Whether concurrent requests cause service crashes.

Early validation does not require fully achieving production performance, but it needs to confirm that the critical path is available and provide a baseline for subsequent optimization.

4. Recommendations for Inference Performance Optimization

1. Quantization and Low-Precision Inference

For ultra-large models, quantization is an important means to reduce memory usage and improve inference speed. Common approaches include:

• Static weight quantization;
• Dynamic activation quantization;
• Mixed-precision inference.

In the early stages of adaptation, single-precision loading can be used first to ensure correctness; low-precision schemes can then be introduced and output errors compared.

2. Communication and Computation Overlap

In multi-card inference, communication is often the performance bottleneck. By overlapping communication with computation, overall latency can be reduced. For example, in tensor parallelism, collective communication of previous-round results can be performed while computing partial results.

3. Caching and Reuse

For repeated inference scenarios, prefix caching can be used to store key-value tensors of common prefixes, reducing redundant computation. For multimodal models, image feature caching can also significantly reduce latency.

4. Dynamic Batching

In service scenarios, merging multiple requests into a batch can improve throughput. However, ultra-large models are sensitive to memory, so batch size needs to be dynamically adjusted based on input length to avoid memory overflow.

5. Stability and Observability

Ultra-large model inference services need good observability. It is recommended to add the following capabilities during adaptation:

• Memory usage monitoring;
• Inference latency statistics;
• Request success rate recording;
• Model loading time recording;
• Exception log retention.

When service startup fails or inference is abnormal, logs can help quickly determine whether the issue is a model format problem, insufficient memory, communication timeout, or invalid input.

6. Summary

Unified training-and-inference adaptation for trillion-parameter models is a systems engineering effort involving multiple stages such as model format, memory planning, parallel strategies, service encapsulation, multimodal pipelines, and performance optimization. Completing Day-0 adaptation is not the endpoint but the starting point for subsequent stability optimization and performance tuning.

For developers, establishing a reusable loading and inference validation workflow is far more valuable than getting a single model to run. By reasonably designing weight loading, parallel splitting, and request handling pipelines, it is possible to provide a feasible path for ultra-large model inference on limited hardware resources and lay the foundation for subsequent large-scale deployment.