Self-Attention Mechanisms in HPC Job Scheduling: A Novel Framework Combining Gated Transformers and Enhanced PPO
Source abstract
Abstract
In HPC systems, job scheduling plays a critical role in determining resource allocation and task execution order. With the continuous expansion of computing scale and increasing system complexity, modern HPC scheduling faces two major challenges: a massive decision space consisting of tens of thousands of computing nodes and a huge job queue, as well as complex temporal dependencies between jobs and dynamically changing resource states.Traditional heuristic algorithms and basic reinforcement learning methods often struggle to effectively address these challenges in dynamic HPC environments. This study proposes a novel scheduling framework that combines GTrXL with PPO, achieving significant performance improvements through multiple technical innovations. The framework leverages the sequence modeling capabilities of the Transformer architecture and selectively filters relevant historical scheduling information through a dual-gate mechanism, improving long sequence modeling efficiency compared to standard Transformers. The proposed SECT module further enhances resource awareness through dynamic feature recalibration, achieving improved system utilization compared to similar attention mechanisms. Experimental results on multiple datasets (ANL-Intrepid, Alibaba, SDSC-SP2) demonstrate that the proposed components achieve significant performance improvements over baseline PPO implementations. Comprehensive evaluations on synthetic workloads and real HPC trace data show improvements in resource utilization and waiting time, particularly under high-load conditions, while maintaining good robustness across various cluster configurations.
Source: crossref · Quality: baseline
Research context
OpenAlex topics: Cloud Computing and Resource Management · Distributed and Parallel Computing Systems · Parallel Computing and Optimization Techniques
SciVX discovery tags are terminology-based. OpenAlex topics are source-provided classifications.
Affiliations
School of Computer and Artificial Intelligence, Zhengzhou University, Zhengzhou 450001, China
National Supercomputing Center in Zhengzhou, Zhengzhou 450001, China
Funding
National Key Research and Development Program of China
References (64)
Show source-provided references
1. Demystifying Parallel and Distributed Deep Learning: An In-depth Concurrency Analysis
2. Deep Learning in Remote Sensing: A Comprehensive Review and List of Resources10.1109/MGRS.2017.2762307 ↗
3. DRAS: Deep Reinforcement Learning for Cluster Scheduling in High Performance Computing10.1109/TPDS.2022.3205325 ↗
4. Multi-objective scheduling for scientific workflow in multicloud environment10.1016/j.jnca.2018.03.028 ↗
5. Neo: A learned query optimizer10.14778/3342263.3342644 ↗
6. Narayanan, D., Santhanam, K., Kazhamiaka, F., Phanishayee, A., and Zaharia, M. (2020, January 4–6). Heterogeneity-Aware Cluster Scheduling Policies for Deep Learning Workloads. Proceedings of the 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20), New York, NY, USA.
7. Cheng, M., Li, J., and Nazarian, S. (2018, January 22–25). DRL-cloud: Deep reinforcement learning-based resource provisioning and task scheduling for cloud service providers. Proceedings of the 2018 23rd Asia and South Pacific Design Automation Conference (ASP-DAC), Jeju, Republic of Korea.10.1109/ASPDAC.2018.8297294 ↗
8. Barve, M., Sinha, S., Hardikar, R.P., Gunturu, A., and Mallik, W. (2022, January 24–26). Workload Characterization in HPC Environment for Auto-scaling of Resources—Preliminary Study. Proceedings of the 2022 IEEE 19th India Council International Conference (INDICON), Kochi, India.10.1109/INDICON56171.2022.10040124 ↗
9. HPC AI500 V3.0: A scalable HPC AI benchmarking framework10.1016/j.tbench.2022.100083 ↗
10. Hardware resources contention-aware scheduling of hard real-time multiprocessor systems10.1016/j.sysarc.2021.102223 ↗
11. Amaral, M., Polo, J., Carrera, D., Seelam, S., and Steinder, M. (2017, January 12–17). Topology-aware GPU scheduling for learning workloads in cloud environments. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, New York, NY, USA. SC ’17.10.1145/3126908.3126933 ↗
12. Huang, J., Xiao, C., and Wu, W. (2020, January 21–24). RLSK: A Job Scheduler for Federated Kubernetes Clusters based on Reinforcement Learning. Proceedings of the 2020 IEEE International Conference on Cloud Engineering (IC2E), Sydney, NSW, Australia.10.1109/IC2E48712.2020.00019 ↗
13. Gómez, C., Martınez, F., Armejach, A., Moretó, M., Mantovani, F., and Casas, M. (2019, January 20–24). Design Space Exploration of Next-Generation HPC Machines. Proceedings of the 2019 IEEE International Parallel and Distributed Processing Symposium (IPDPS), Rio de Janeiro, Brazil.10.1109/IPDPS.2019.00017 ↗
14. Towards Metaheuristic Scheduling Techniques in Cloud and Fog: An Extensive Taxonomic Review10.1145/3494520 ↗
15. Achieving Fairness-Aware Two-Level Scheduling for Heterogeneous Distributed Systems10.1109/TSC.2018.2836444 ↗
16. Tirmazi, M., Barker, A., Deng, N., Haque, M.E., and Wilkes, J. (2020, January 27–30). Borg: The next generation. Proceedings of the EuroSys ’20: Fifteenth EuroSys Conference 2020, Heraklion, Greece.
17. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł., and Polosukhin, I. (2017, January 4–9). Attention is all you need. Proceedings of the 31st International Conference on Neural Information Processing Systems, Long Beach, CA, USA. NIPS’17.
18. Ermis, B., Zappella, G., Wistuba, M., Rawal, A., and Archambeau, C. (2022). Continual Learning with Transformers for Image Classification. arXiv.10.1109/CVPRW56347.2022.00422 ↗
19. Parisotto, E., Song, H.F., Rae, J.W., Pascanu, R., Gulcehre, C., Jayakumar, S.M., Jaderberg, M., Kaufman, R.L., Clark, A., and Noury, S. (2019). Stabilizing Transformers for Reinforcement Learning. arXiv.
20. Hu, J., Shen, L., and Sun, G. (2018, January 18–23). Squeeze-and-Excitation Networks. Proceedings of the 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, Salt Lake City, UT, USA.10.1109/CVPR.2018.00745 ↗
21. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017). Proximal Policy Optimization Algorithms. arXiv.
22. Wang, Y., He, H., Wen, C., and Tan, X. (2020). Truly Proximal Policy Optimization. arXiv.
23. Fan, Y., Lan, Z., Rich, P., Allcock, W., and Papka, M.E. (June, January 30). Hybrid Workload Scheduling on HPC Systems. Proceedings of the 2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS), Lyon, France.10.1109/IPDPS53621.2022.00052 ↗
24. STORM: Scalable Resource Management for Large-Scale Parallel Computers10.1109/TC.2006.206 ↗
25. A scheduling algorithm based on reinforcement learning for heterogeneous environments10.1016/j.asoc.2022.109707 ↗
26. Lifka, D. (1995). The ANL/IBM SP Scheduling System.10.1007/3-540-60153-8_35 ↗
27. Jette, M., Yoo, A., and Grondona, M. (2003). SLURM: Simple Linux Utility for Resource Management.
28. Fairness in parallel job scheduling10.1002/1099-1425(200009/10)3:5<297::AID-JOS50>3.0.CO;2-D ↗
29. Tang, W., Lan, Z., Desai, N., and Buettner, D. (September, January 31). Fault-aware, utility-based job scheduling on Blue, Gene/P systems. Proceedings of the 2009 IEEE International Conference on Cluster Computing and Workshops, New Orleans, LA, USA.10.1109/CLUSTR.2009.5289206 ↗
30. Feitelson, D.G., and Rudolph, L. (2001). Core Algorithms of the Maui Scheduler. Job Scheduling Strategies for Parallel Processing, Springer.10.1007/3-540-45540-X ↗
31. Utilization, predictability, workloads, and user runtime estimates in scheduling the IBM SP2 with backfilling10.1109/71.932708 ↗
32. Mao, H., Alizadeh, M., Menache, I., and Kandula, S. (2016, January 9–10). Resource Management with Deep Reinforcement Learning. Proceedings of the 15th ACM Workshop on Hot Topics in Networks, Atlanta, GA, USA.10.1145/3005745.3005750 ↗
33. Mao, H., Schwarzkopf, M., Venkatakrishnan, S.B., Meng, Z., and Alizadeh, M. (2019, January 19–23). Learning scheduling algorithms for data processing clusters. Proceedings of the ACM Special Interest Group on Data Communication, Beijing, China.10.1145/3341302.3342080 ↗
34. Wong, A.K., and Goscinski, A.M. (2007, January 17–20). Evaluating the EASY-backfill job scheduling of static workloads on clusters. Proceedings of the 2007 IEEE International Conference on Cluster Computing, Austin, TX, USA.10.1109/CLUSTR.2007.4629218 ↗
35. Chadha, M., John, J., and Gerndt, M. (2021). Extending SLURM for Dynamic Resource-Aware Adaptive Batch Scheduling. arXiv.10.1109/HiPC50609.2020.00036 ↗
36. DRAS-CQSim: A Reinforcement Learning based Framework for HPC Cluster Scheduling10.1016/j.simpa.2021.100077 ↗
37. Human-level control through deep reinforcement learning10.1038/nature14236 ↗
38. Lillicrap, T.P., Hunt, J.J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., Silver, D., and Wierstra, D. (2016, January 2–4). Continuous control with deep reinforcement learning. Proceedings of the 4th International Conference on Learning Representations, San Juan, Puerto Rico.
39. Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. (2018). Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor. arXiv.
40. Espeholt, L., Marinier, R., Stanczyk, P., Wang, K., and Michalski, M. (2020). SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference. arXiv.
41. Deep reinforcement learning for dynamic distributed job shop scheduling problem with transfers10.1016/j.eswa.2024.123970 ↗
42. Task Allocation on Layered Multiagent Systems: When Evolutionary Many-Objective Optimization Meets Deep Q-Learning10.1109/TEVC.2021.3049131 ↗
43. Distributed Subgradient Methods for Multi-Agent Optimization10.1109/TAC.2008.2009515 ↗
44. Gu, Y., Liu, Z., Dai, S., Liu, C., Wang, Y., and Wang, S. (2025). Deep Reinforcement Learning for Job Scheduling and Resource Management in Cloud Computing: An Algorithm-Level Review. arXiv.
45. An integrated optimization method to task scheduling and VM placement for green datacenters10.1016/j.simpat.2024.102962 ↗
46. Job Scheduling in Hybrid Clouds with Privacy Constraints: A Deep Reinforcement Learning Approach10.1002/cpe.8307 ↗
47. Transformer in Transformer
48. Apache Spark: A unified engine for big data processing10.1145/2934664 ↗
49. Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q., and Salakhutdinov, R. (August, January 28). Transformer-XL: Attentive Language Models beyond a Fixed-Length Context. Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, Florence, Italy.
50. Deng, K., He, Z., Zhang, H., Lin, H., and Wang, D. (2023). TSNet-SAC: Leveraging Transformers for Efficient Task Scheduling. arXiv.
51. Lee, J., Kee, S., Janakiram, M., and Runger, G. (2024). Attention-based Reinforcement Learning for Combinatorial Optimization: Application to Job Shop Scheduling Problem. arXiv.10.2139/ssrn.4830118 ↗
52. Hastie, T., Tibshirani, R., and Friedman, J. (2009). The Elements of Statistical Learning: Data Mining, Inference, and Prediction, Springer. [2nd ed.].10.1007/978-0-387-84858-7 ↗
53. El-Mahdy, M., Sakr, N., and Carrasco, R. (2025). Improving Mixed-Criticality Scheduling with Reinforcement Learning. arXiv.
54. PyTorch distributed: Experiences on accelerating data parallel training10.14778/3415478.3415530 ↗
55. Tang, H., Gan, S., Awan, A.A., Rajbhandari, S., Li, C., Lian, X., Liu, J., Zhang, C., and He, Y. (2021, January 18–24). 1-bit Adam: Communication Efficient Large-Scale Training with Adam’s Convergence Speed. Proceedings of the International Conference on Machine Learning, Virtual.
56. Peng, Y., Bao, Y., Chen, Y., Wu, C., and Guo, C. (2017, January 28). A Generic Communication Scheduler for Distributed DNN Training Acceleration. Proceedings of the 26th Symposium on Operating Systems Principles, Shanghai, China.
57. Argonne National Laboratory (2024, August 01). ANL-Intrepid Workload Archive, Available online: https://www.cs.huji.ac.il/labs/parallel/workload/l_anl_int/.
58. Experience with using the Parallel Workloads Archive10.1016/j.jpdc.2014.06.013 ↗
59. Optimizing spatial distribution of watershed-scale hydrologic models using Gaussian Mixture Models10.1016/j.envsoft.2021.105076 ↗
60. Li, B., Fan, Y., Dearing, M., Lan, Z., Rich, P., and Allcock, W. (2022, January 5–8). MRSch: Multi-Resource Scheduling for HPC. Proceedings of the 2022 IEEE International Conference on Cluster Computing (CLUSTER), Heidelberg, Germany.10.1109/CLUSTER51413.2022.00020 ↗
61. Li, J., Liu, J., Fan, Y., and Lan, Z. (2021). CQGym: A Reinforcement Learning Environment for Cluster Scheduling. arXiv.
62. Woo, S., Park, J., Lee, J.Y., and Kweon, I.S. (2018, January 8–14). CBAM: Convolutional Block Attention Module. Proceedings of the European Conference on Computer Vision (ECCV), Munich, Germany.10.1007/978-3-030-01234-2_1 ↗
63. Wang, Q., Wu, B., Zhu, P., Li, P., Zuo, W., and Hu, Q. (2020, January 13–19). ECA-Net: Efficient Channel Attention for Deep Convolutional Neural Networks. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), Seattle, WA, USA.10.1109/CVPR42600.2020.01155 ↗
64. Online over time processing of combinatorial problems10.1007/s10601-018-9287-4 ↗