Summary (Overview)

• This paper presents SLAI T-Rex, a full-stack optimization framework for full-parameter post-training of trillion-parameter Mixture-of-Experts (MoE) models on Ascend NPU SuperPOD infrastructure, achieving 34.22% Model FLOPs Utilization (MFU), a 2.93× improvement over the open-source baseline (from 11.67%).
• The optimization spans three levels: model-level parallelism, computation-communication orchestration, and low-level kernel execution, including a novel AuraKernel performance optimization agent for complex Ascend kernels.
• Building on the optimized infrastructure, the paper establishes a Continued Pre-Training (CPT) and Supervised Fine-Tuning (SFT) workflow specialized for Operations Research (OR) tasks, using DeepSeek-V4-Flash as the base model.
• The specialized model achieves 71.81% zero-shot Pass@1 on OR benchmarks, outperforming GPT-5.4-Mini by 3.98 percentage points and the base DeepSeek-V4-Flash by 11.27 percentage points.
• The paper demonstrates that CPT establishes transferable OR-domain modeling priors, improving both solver-facing feasibility and structural equivalence after SFT.

Introduction and Theoretical Foundation

The paper addresses the substantial system-level challenges of full-parameter post-training for trillion-parameter-scale MoE models on non-GPU infrastructure. Key challenges include:

• Severe memory pressure from large model parameters and optimizer states.
• Non-overlapped communication overhead due to the all-to-all communication patterns in MoE expert parallelism.
• Inefficient kernel execution on Ascend NPUs, which lack the mature, optimized libraries (e.g., cuBLAS, flash-attention) available on GPU clusters.

The theoretical foundation lies in developing a hierarchical optimization framework that jointly optimizes:

1. Model-level parallelism (e.g., expert parallelism, sequence parallelism).
2. Computation–communication orchestration (e.g., overlapping communication with computation).
3. Low-level kernel execution (e.g., kernel fusion, rewrite, and auto-tuning).

The target domain is Operations Research (OR) modeling, which requires joint capability in mathematical formulation, structured reasoning, and code generation. This makes it a challenging and representative domain for evaluating the model's post-training capabilities.

The paper defines Model FLOPs Utilization (MFU) as:

$$\text{MFU} = \frac{\text{FLOPs achieved}}{\text{Peak FLOPs}} \times 100\%$$

The optimization rate is defined as:

$$\text{Optimization Rate} = \frac{\text{MFU}_{optimized}}{\text{MFU}_{baseline}}$$

Methodology

The methodology is divided into two main parts: infrastructure optimization and OR-oriented post-training.

1. Infrastructure Optimization on Ascend SuperPOD

• Parallelism Strategy: Optimized parallelism configuration balancing data, tensor, sequence, and expert parallelism to minimize communication overhead and memory fragmentation.
• Computation-Communication Overlap: Advanced scheduling techniques to overlap all-to-all communication in MoE layers with computation, reducing idle time.
• NPU-CPU Co-optimization: Offloading certain operations (e.g., data preprocessing, logging) to CPU to reduce NPU idle time.
• Kernel Optimization:
  • AuraKernel: A performance optimization agent that automatically analyzes bottleneck kernels and generates optimized code. It performs:
    1. Kernel profiling and bottleneck analysis.
    2. Automatic kernel rewriting using domain-specific knowledge of Ascend NPU architecture.
    3. Kernel fusion to combine multiple small kernels into a single, more efficient kernel.
  • Ascend NPU-aware Kernel Fusion and Rewrite: Manual and automatic fusion of operators (e.g., attention, layer norm, activation) to reduce memory access and launch overhead.

2. OR-Oriented Post-Training Workflow

The workflow is based on a two-stage training adaptation: CPT then SFT.

• Capability Boundary Analysis: A detailed analysis of the base model's weaknesses in OR tasks, identifying errors in mathematical formulation, reasoning, and code generation.
• OR-Oriented Continued Pre-Training (CPT):
  • Data Construction: A domain-dominant mixture of collected OR-related resources (e.g., textbooks, papers, code repositories) and solver-verified synthetic optimization documents.
  • Training Execution: CPT is performed on the Ascend SuperPOD using the optimized infrastructure.
  • Checkpoint Selection: Checkpoints are selected based on their transferability to the downstream SFT task, validated by a small-scale SFT proxy.
• Data-Efficient Supervised Fine-Tuning (SFT):
  • Solver-Oriented SFT Objective: The model is fine-tuned to generate code that is directly executable by a solver (e.g., Gurobi, CPLEX), ensuring feasibility and correctness.
  • Self-Distilled OR Modeling Data Flywheel: A pipeline that uses the model itself to generate new SFT samples, which are then verified by a solver and used for further training, creating a data flywheel.
  • Contract-Aware Cleaning and CoT Enhancement: Data cleaning based on "contracts" (e.g., constraint satisfaction, objective function correctness) and enhancement with Chain-of-Thought (CoT) reasoning.
  • Progressive SFT with General Capability Anchoring: A training strategy that gradually introduces OR-specific data while maintaining general language capabilities.
• AI-Assisted Data Construction and Refinement: Using the model itself and other LLMs to generate and refine OR problem statements, solutions, and code.

Empirical Validation / Results

The paper presents extensive experimental results demonstrating the effectiveness of the proposed methods.

1. Efficient Kernels via AuraKernel

Table 1: AuraKernel Optimization Results

AuraKernel achieves significant speedups across various bottleneck kernels, with the most dramatic improvement in the MoE Gate operation (2.44×).

2. Performance Improvements from Kernel Fusion

Kernel fusion reduces kernel launch overhead and memory traffic. The fused attention kernel (combining QKV projection, attention computation, and output projection) shows a 1.93× speedup. The overall system benefits from a cascading effect, where kernel optimization contributes to the total MFU improvement.

3. Long-term Training Stability with Ascend SuperPOD

The optimized system maintains stable training for an extended period with no loss spikes or divergence. The MFU remains consistently high (approximately 34.22%) throughout the training, as shown in Figure 1(a).

4. SFT Scaling and Cleaning Analysis

Table 2: SFT Data Scaling and Cleaning Effects

Contract-aware cleaning provides a consistent improvement across all data scales. The 10K contract-aware SFT dataset yields the best performance.

5. Effect of CPT Initialization under Matched SFT Conditions

Table 3: CPT Initialization Effect

CPT provides a strong initialization for SFT, improving the final performance. The optimal CPT data size is 5K tokens, with diminishing returns and potential overfitting beyond that.

6. Performance of the Proposed Post-Training Pipeline

Table 4: Final Model Performance on OR Benchmarks (Zero-shot Pass@1)

SLAI T-Rex achieves the highest average zero-shot Pass@1 score across all four benchmarks, demonstrating its superior capability in OR modeling.

Theoretical and Practical Implications

• Theoretical Implications: The paper validates a hierarchical optimization framework for training large models on non-GPU infrastructure. The success of AuraKernel suggests that automated, architecture-aware kernel optimization is a promising direction. The CPT-to-SFT transfer learning paradigm is validated for the complex domain of OR, confirming that domain-specific pre-training can establish strong, transferable priors.
• Practical Implications:
  • For AI Infrastructure: The paper provides a practical recipe for achieving high MFU (34.22%) on Ascend NPU clusters, demonstrating that these systems can be viable alternatives to GPU clusters for large-scale LLM training.
  • For Domain Specialization: The data-efficient CPT-SFT pipeline (10K SFT samples) provides a blueprint for creating specialized models for complex reasoning domains like OR, without requiring massive amounts of domain-specific data.
  • For OR Communities: The specialized model, SLAI T-Rex, offers a significant improvement over existing general-purpose models for solver-grounded mathematical modeling, potentially accelerating research and application in OR.

Conclusion

This paper introduces SLAI T-Rex, a full-stack system for training and specializing trillion-parameter MoE models on Ascend SuperPOD infrastructure. The key contributions are:

• A hierarchical optimization framework achieving 34.22% MFU, a 2.93× improvement over the baseline.
• The AuraKernel performance optimization agent for automated kernel optimization on Ascend NPUs.
• A data-efficient CPT-SFT workflow for specializing in OR tasks, achieving 71.81% zero-shot Pass@1 on OR benchmarks, outperforming GPT-5.4-Mini.

Limitations and Future Directions:

• The OR specialization is limited to a single domain; future work should explore multi-domain specialization.
• The data flywheel is human-in-the-loop; future work could explore fully automated closed-loop data generation.
• The system validation is on a single Ascend SuperPOD; generalization to other hardware and cluster topologies needs further study.
• The paper focuses on post-training; future work can extend the optimization to pre-training from scratch.