MaxText-Slurm: Production-Grade LLM Training with Built-In Observability

2026-03-02 · Source: AMD ROCm Blogs · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Cloud Computing & IT Infrastructure, Data Science & Analytics · Depth: Intermediate, long

Summary

MaxText-Slurm is an open-source launch system and observability stack designed to streamline large language model (LLM) training on AMD Instinct GPU clusters managed by Slurm. It addresses operational challenges like multi-node coordination, environment containerization, runtime tuning, and failure diagnosis, which are not typically handled by training frameworks like MaxText. The system features a single-command workflow for launching MaxText training and an observability stack, powered by Ray and Prometheus, that provides real-time and post-mortem visibility into system layers. This stack collects GPU, host, network, and training metrics into a unified Prometheus time-series database, ensuring data persistence even upon unexpected job termination. MaxText-Slurm also supports AI-assisted diagnosis through agentic workflows that interpret metrics to identify root causes of issues like RCCL hangs.

Key takeaway

For NLP Engineers deploying MaxText on AMD Instinct GPU clusters, MaxText-Slurm simplifies distributed training and enhances operational reliability. You should integrate this system to gain real-time visibility into GPU, host, network, and training metrics, enabling faster diagnosis of complex issues like NCCL hangs or thermal throttling. This will reduce operational overhead and improve the stability of your LLM training workflows, especially for long-running production jobs.

Key insights

MaxText-Slurm provides a unified, observable, and extensible system for production-grade LLM training on AMD GPUs.

Principles

• Isolate concerns in swappable tiers.
• Unify metrics for comprehensive diagnosis.
• Ensure zero steady-state overhead for observability.

Method

MaxText-Slurm uses a layered architecture for orchestration, containerization, and training. Its observability stack, powered by Ray and Prometheus, collects metrics via auto-discovered plugins, running training in a subprocess to avoid overhead.

In practice

• Use `RAY=1` to activate full observability.
• Browse post-mortem TSDB with `utils/prometheus.sh view`.
• Add custom metrics via `*_metrics_plugin.sh` scripts.

Topics

• LLM Training
• Distributed Training
• ML Observability
• AMD Instinct GPUs
• AI-Assisted Diagnostics

Code references

• AMD-AGI/maxtext-slurm
• AI-Hypercomputer/maxtext
• kubernetes-sigs/jobset

Best for: NLP Engineer, MLOps Engineer, AI Engineer, Machine Learning Engineer

Related on AIssential

• Streamlining Recommendation Model Training on AMD Instinct™ GPUs — ROCm training docker simplifies recommendation model training on AMD Instinct GPUs by optimizing sparse embedding handling.
• Train Models Faster with JAX and MaxText Using NVFP4 on NVIDIA Blackwell — NVFP4 enables significantly faster, accurate 4-bit LLM pre-training on NVIDIA Blackwell GPUs.
• Researchers say they trained a foundation model from scratch for about $1,500 — HRM-Text offers a cost-effective, sample-efficient architecture for training specialized 1B-parameter reasoning models from scratch.
• Open model, open metrics: How Lambda and the Olmo team trained Olmo Hybrid — Hybrid RNN-transformer architectures can enhance model expressivity and performance, particularly in structured reasoning tasks.
• Boost Training Goodput: How Continuous Checkpointing Optimizes Reliability in Orbax and MaxText — Continuous checkpointing in Orbax and MaxText dynamically optimizes training reliability and performance by adapting save frequency.
• Agentic Diagnosis for LLM Training at Scale — AI agents with tool access can autonomously diagnose complex LLM training issues using structured diagnostic skills and unified telemetry.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by AMD ROCm Blogs.