RATrain: A Resource-Aware Training Runtime for Large Language Models on Bandwidth-Constrained Heterogeneous Supercomputing Platforms

2026-06-09

AI Summary

This work addresses the challenges of training large language models on bandwidth-constrained heterogeneous supercomputing platforms—such as the MT-3000—characterized by limited GPU memory, low communication bandwidth, and a lack of tailored runtime systems. The authors propose the first resource-aware LLM training runtime specifically designed for such environments, formulating non-interleaved 1F1B training as a training state lifecycle scheduling problem. Their approach co-schedules gradient synchronization, parameter updates, prefetching, and activation recomputation across both layer and pipeline stage granularities, integrating a platform-customized execution backend with a resource-constrained planner. Leveraging hierarchical state scheduling, FP16 GEMM and attention backward optimizations, and explicit data movement, the system enables end-to-end training of LLaMA-2 models ranging from 7B to 70B under a 20GB DDR memory constraint, achieving up to 1.35× speedup. Furthermore, LLaMA-2-7B scales to 1,024 nodes at 112,790.55 tokens/s with 97.0% scaling efficiency and loss deviation below 0.081%.

Abstract

Production heterogeneous supercomputing platforms are increasingly used to host large language model (LLM) training workloads. However, existing GPU-oriented training runtimes typically rely on high-bandwidth device memory, fast interconnects, and mature collective communication libraries, making them difficult to directly adapt to MT-3000, a platform with an explicit memory hierarchy, limited usable DDR capacity, and constrained inter-cluster communication. This paper presents RATrain, a resource-aware training runtime for dense LLMs on bandwidth-constrained heterogeneous supercomputing platforms. RATrain formulates standard non-interleaved 1F1B training as a training-state lifecycle scheduling problem, and schedules gradient synchronization, parameter update, parameter-view prefetching, and activation recovery at layer-level and stage-local granularity. RATrain further combines an MT-3000-aware execution backend for efficient and predictable FP16 GEMM, Attention Backward, and explicit data movement with a resource-aware planner that selects feasible training configurations under the 20GB usable-DDR constraint per compute cluster. We implement RATrain on a real MT-3000 platform and evaluate it using LLaMA-2-7B, Baichuan2-13B, Qwen2.5-32B, and LLaMA-2-70B configurations. Results show that RATrain achieves up to 1.35$\times$ end-to-end speedup over MT-3000-adapted GPU-style training strategies. For LLaMA-2-7B, RATrain scales to 1024 compute clusters, reaches 112,790.55 tokens/s, and achieves 97.0\% scaling efficiency. A further 1.028B-token correctness run shows that RATrain preserves the loss trajectory of a semantically equivalent Baseline-1F1B run, with a maximum relative loss deviation of 0.081\%.

Problem

• Large Language Models
• Heterogeneous Supercomputing
• Bandwidth-Constrained
• Resource-Aware Training
• Memory Hierarchy

Innovation

• resource-aware training
• heterogeneous supercomputing
• bandwidth-constrained
• LLM training runtime
• memory hierarchy optimization