Building the Foundation: Setting Up a Production ML Environment on Ubuntu 24.04

Introduction

In the world of production machine learning serving, the quality of the infrastructure foundation directly determines the reliability and performance of the deployed models. Before a single inference request can be served, before any throughput benchmarking or latency tuning can begin, the underlying environment must be meticulously constructed — NVIDIA drivers installed, CUDA toolkits configured, Python virtual environments created, and complex CUDA extensions compiled. This article synthesizes a chunk of an opencode coding session that documents exactly such a foundation-building effort: the complete setup of an ML development and serving environment on Ubuntu 24.04, culminating in the deployment of the GLM-5-NVFP4 model using SGLang.

The work described in this chunk spans the full arc from bare-metal system configuration to application deployment. It begins with the installation of NVIDIA drivers and CUDA Toolkit 13.1 on a fresh Ubuntu 24.04 machine, proceeds through the creation of a Python virtual environment using uv, navigates the treacherous waters of compiling flash-attn from source (a process that tests the limits of system memory and CUDA compatibility), and ultimately pivots to deploying a production-grade language model across eight GPUs. Along the way, the assistant and user collaborate through multiple rounds of trial, error, and refinement — a testament to the iterative nature of ML infrastructure engineering.

Installing NVIDIA Drivers and CUDA Toolkit

The first task was to install the latest NVIDIA drivers and CUDA toolkit on a remote Ubuntu 24.04 machine. The assistant connected to the machine and began by checking the available GPU hardware. The system was equipped with two RTX PRO 6000 Blackwell GPUs — high-end professional GPUs designed for demanding compute workloads.

The assistant installed NVIDIA drivers version 590.48.01, the latest available at the time, along with CUDA Toolkit 13.1. This process involved downloading the CUDA runfile from NVIDIA's repository, executing the installer, and configuring environment paths. The assistant carefully set PATH, LD_LIBRARY_PATH, and CUDA_HOME to point to the new CUDA installation, ensuring that subsequent builds would find the correct toolkit.

Verification was a critical step. The assistant ran nvidia-smi to confirm both GPUs were recognized and operational, and nvcc --version to verify the CUDA compiler was correctly installed. The output showed both RTX PRO 6000 Blackwell GPUs online, each with substantial memory capacity, and CUDA 13.1 ready for use.

Creating the Python Environment with uv

With the GPU drivers and CUDA toolkit in place, the next step was to set up a Python virtual environment. The assistant chose uv — a modern, fast Python package manager — over the traditional pip/venv combination. This decision reflected a preference for speed and reliability in dependency management, especially important when dealing with the complex dependency graphs of ML frameworks.

The assistant created a virtual environment, activated it, and installed core packages. PyTorch was the primary dependency, and the assistant installed it with CUDA support, ensuring the version was compatible with the installed CUDA toolkit. The installation was verified by running a simple PyTorch script that checked CUDA availability and GPU count — confirming that PyTorch could see both GPUs and was using the correct CUDA version.

The flash-attn Build Odyssey

The installation of flash-attn — a critical CUDA extension that implements the Flash Attention algorithm for efficient transformer inference — became the central challenge of this phase. Flash attention is a memory-efficient attention mechanism that reduces the quadratic memory complexity of standard attention to linear, enabling larger batch sizes and longer sequences. However, building it from source on a new system is notoriously difficult.

The first obstacle was a CUDA version mismatch. The system had CUDA 13.1 installed, but PyTorch's prebuilt binaries were compiled against CUDA 12.8. When flash-attn attempted to build against CUDA 13.1, it produced binaries incompatible with PyTorch's expectations. The solution was to install a secondary CUDA 12.8 toolkit alongside the primary 13.1 installation, and configure the build to use the compatible version.

The second obstacle was memory exhaustion during compilation. Building flash-attn involves compiling hundreds of CUDA kernel variants, and the default parallel compilation setting (MAX_JOBS) was set to 128 — the number of CPU threads available on the system. Each compilation job requires significant memory, and 128 concurrent jobs quickly exhausted the available RAM, causing the build to fail with out-of-memory errors.

The assistant and user collaborated on a trial-and-error reduction of MAX_JOBS. Starting from 128, they tried progressively lower values — 64, 48, 32 — each time the build failed or took impractically long. The breakthrough came when the machine was rebooted with an expanded 432GB of RAM, allowing MAX_JOBS=20 to succeed. This experience highlighted a crucial lesson: building CUDA extensions at scale requires careful resource management, and the default parallelism settings are often inappropriate for memory-constrained builds.

Dependency Conflicts and the vLLM Interaction

The challenges did not end with a successful flash-attn build. The installation of vLLM — another critical component of the serving stack — introduced a new dependency conflict. vLLM's installation process downgraded PyTorch from version 2.9.1 to an older version, breaking the flash-attn binary that had been compiled against the original PyTorch. The compiled CUDA kernels were incompatible with the new PyTorch runtime, causing import errors and runtime crashes.

The solution was a targeted rebuild of flash-attn against the correct PyTorch version. The assistant first ensured that PyTorch 2.9.1 was reinstated (preventing vLLM from downgrading it), then rebuilt flash-attn from source against this version. This required careful ordering of the installation steps: PyTorch first, then flash-attn, then vLLM — with explicit version pinning to prevent unintended upgrades or downgrades.

The Final Environment Stack

After navigating these challenges, the environment was stabilized with a fully compatible software stack:

Pivot to Application Deployment

With the environment stabilized, the conversation pivoted sharply from infrastructure setup to application deployment. The machine was upgraded from 2 to 8 GPUs, and the user tasked the assistant with deploying the GLM-5-NVFP4 model using a nightly build of SGLang — a high-performance inference engine for large language models.

The GLM-5-NVFP4 model is a large language model that uses NVFP4 (NVIDIA FP4) quantization for efficient inference. Deploying it required configuring SGLang's model loading pipeline, setting up the serving endpoints, and preparing for performance tuning and load testing. The assistant began by examining the SGLang configuration options, understanding the model's architecture requirements, and planning the deployment strategy.

This pivot from environment setup to model deployment marked a significant transition in the session's work. The foundation had been laid — the drivers, CUDA toolkit, Python environment, and critical CUDA extensions were all in place. Now the focus shifted to the application layer: getting the model loaded, serving requests, and optimizing performance.

Key Insights and Lessons

Several important lessons emerge from this chunk of work:

1. CUDA version compatibility is a first-class concern. The mismatch between CUDA 13.1 (system) and CUDA 12.8 (PyTorch's base) required a secondary toolkit installation. This is a common pain point in ML infrastructure: the CUDA version used by the system and the one used by PyTorch's prebuilt binaries are often different, and CUDA extensions must be compiled against the version PyTorch expects.

2. Parallel compilation jobs must be tuned to available memory. The default MAX_JOBS=128 was appropriate for CPU throughput but disastrous for memory. The iterative reduction from 128 to 20, combined with a RAM expansion to 432GB, finally allowed the build to succeed. This is a practical lesson that applies to any large CUDA extension build.

3. Dependency ordering matters. The vLLM installation's tendency to downgrade PyTorch required careful version pinning and installation ordering. In ML infrastructure, package managers cannot always resolve complex dependency graphs correctly — manual intervention is often necessary.

4. Verification at each step prevents cascading failures. The assistant verified each component after installation — GPU detection, CUDA compiler, PyTorch CUDA support, flash-attn import — catching issues early before they could compound.

5. The boundary between infrastructure and application is porous. The flash-attn build issues were ultimately about serving performance, and the CUDA toolkit configuration affected model deployment. Infrastructure decisions have direct consequences for application behavior.

Conclusion

The ML environment setup described in this chunk is a microcosm of the broader challenges in production ML infrastructure. It demonstrates that before any model serving, any performance optimization, or any throughput benchmarking can occur, a solid foundation must be built — one that requires careful attention to driver versions, CUDA compatibility, memory management during builds, and dependency resolution. The assistant's methodical approach — installing and verifying each component, troubleshooting build failures through iterative parameter tuning, and maintaining a clear picture of the dependency graph — reflects the discipline required for reliable ML infrastructure engineering.

The successful stabilization of the environment, followed by the pivot to GLM-5-NVFP4 deployment with SGLang, set the stage for the performance tuning and load testing that would follow in subsequent sessions. The foundation was solid, and the team could now focus on the higher-level challenges of model serving, throughput optimization, and the subtle distributed systems bugs that would emerge under production load.## References

  1. Chunk summary — ML environment setup on Ubuntu 24.04 with NVIDIA drivers, CUDA Toolkit 13.1, flash-attn build resolution, and GLM-5-NVFP4 deployment with SGLang.
  2. Segment 72 summary — Root-caused bf16 corruption, investigated decode throughput scaling, resolved production PD bootstrap incident, diagnosed client-side session-bible tool hang.