AI Workstation Setup: Ubuntu, NVIDIA Driver, CUDA, Docker

Short answer: to set up an AI workstation on Ubuntu, install things in this order: (1) update Ubuntu and install the NVIDIA GPU driver with ubuntu-drivers, (2) check it with nvidia-smi, (3) create a Python environment and install PyTorch from the pip command on pytorch.org, (4) install Docker Engine and the NVIDIA Container Toolkit if you want containers, and (5) run one test at each layer before you start real work. Install the full CUDA Toolkit only if you need to compile CUDA code. The driver is the one piece every route depends on.

This guide is for people searching for how to set up a deep learning workstation, how to install the NVIDIA driver and CUDA on Ubuntu, how to install PyTorch with GPU support or how to use a GPU inside Docker. Every command and fact below was checked against the vendors' own documentation on 15 November 2026 (sources at the end). Where something is our own practical advice rather than a documented fact, we say so.

What is the correct order to set up an AI workstation?

Each layer sits on the one below it. If a layer fails, fix it before you move up, because errors in the upper layers usually point back to the lower ones.

Layer	What it is	Test that it works
1. Hardware	GPU seated, powered and cooled; ECC memory; fast NVMe	lspci | grep -i nvidia lists the GPU
2. Ubuntu LTS	Supported release, fully updated	hostnamectl shows the release
3. NVIDIA driver	Kernel module plus user-space libraries	nvidia-smi shows the GPU
4. CUDA Toolkit (optional)	Compiler and libraries for building CUDA code	nvcc --version
5. Python environment and PyTorch	Your ML framework	torch.cuda.is_available() returns True
6. Docker and NVIDIA Container Toolkit (optional)	Reproducible containers with GPU access	docker run --gpus all ... nvidia-smi

What should you check before installing any software?

NVIDIA's CUDA installation guide lists three pre-installation checks, and they take two minutes:

• Is the GPU visible to Linux? Run lspci | grep -i nvidia. If nothing shows, NVIDIA says to refresh the PCI hardware database with update-pciids and run the command again. If it still shows nothing, the problem is hardware (seating, power connectors, slot), not software.
• Is the Linux release supported? Run hostnamectl. The current CUDA guide lists Ubuntu 22.04 LTS, 24.04 LTS and 26.04 LTS. Docker Engine supports the same LTS releases.
• Is gcc present? Run gcc --version. NVIDIA notes gcc is required to develop with the CUDA Toolkit but not to run CUDA applications.

Then update the system. Ubuntu's NVIDIA driver documentation says to run sudo apt update && sudo apt upgrade and reboot before installing drivers, so that the kernel is current.

If you are still choosing the machine, read our entry-level ML workstation roadmap and best GPU for local LLMs (VRAM guide) first. Software setup cannot fix a GPU with too little VRAM for your model.

How do you install the NVIDIA driver on Ubuntu?

Use Ubuntu's own tool rather than downloading installers from random sources. Ubuntu's documentation recommends ubuntu-drivers, and warns that drivers installed from sources outside its guide could overwrite the packaged ones and may break Secure Boot.

1. List the drivers for your hardware, server-style: sudo ubuntu-drivers list --gpgpu. The --gpgpu option shows the compute-oriented drivers.
2. Install the recommended one: sudo ubuntu-drivers install --gpgpu. Or name a branch yourself from the list, for example sudo ubuntu-drivers install --gpgpu nvidia:<branch>-server, where the branch comes from step 1 on your machine.
3. Add the matching utilities if they were not pulled in: Ubuntu's page shows sudo apt install nvidia-utils-<branch>-server.
4. Reboot to load the driver.
5. Verify: run nvidia-smi. You should see your GPU, the driver version and the memory size. cat /proc/driver/nvidia/version shows the running driver version.

Which Ubuntu driver type should you pick: normal or -server?

Ubuntu packages two families. Unified Driver Architecture drivers suit general desktop use. Enterprise Ready Drivers, recognisable by the -server suffix, are the ones Ubuntu recommends for servers and for computing tasks. For a machine whose job is training and inference, the -server family is the documented fit. Ubuntu's page also mentions Fabric Manager and NSCQ packages, which you only need if your system has NVSwitch hardware.

What about Secure Boot?

Most modern x86 systems have Secure Boot on. Ubuntu says ubuntu-drivers is the recommended route in that case because by default it installs only the pre-built, signed modules known to work with Secure Boot. The alternative, DKMS modules that are compiled on your machine, is not recommended by Ubuntu unless you run a custom kernel, because DKMS drivers are not signed with Canonical's key; you would have to enrol your own key. The practical lesson: if you want a trouble-free setup, stay on the signed pre-built route.

Do you need the CUDA Toolkit, and how do you install it?

It depends on what you will do:

You want to...	Install the CUDA Toolkit?
Train or run models with PyTorch from the pip wheels	Usually no. PyTorch's page lists installing CUDA under building from source only. You need the driver.
Compile CUDA C/C++ code with nvcc	Yes
Build a Python package from source that needs CUDA headers	Yes, the project will say so
Run everything in containers	No, on the host. The driver and Container Toolkit are the host prerequisites.

If you do need it, NVIDIA's guide offers four methods: package manager (RPM and DEB), runfile, Conda, and pip wheels. NVIDIA recommends the distribution-specific packages where possible, because the runfile does not integrate with the system's package manager. Two rules from the guide are worth following:

• Do not mix install methods. NVIDIA's compatibility matrix says to uninstall a conflicting installation before installing the other type. A leftover runfile install under a package-manager install is a classic source of "wrong CUDA version" confusion.
• Install the driver separately. NVIDIA's quick-start notes that starting with CUDA 13.4 the Linux driver is no longer bundled with the toolkit and must be installed separately. That matches the order in this guide.

After a package-manager install, NVIDIA says the PATH must include the toolkit's bin directory, for example export PATH=/usr/local/cuda-13.5/bin${PATH:+:${PATH}} for the 13.5 release in the current guide (use your installed version). The LD_LIBRARY_PATH variable is needed for the runfile method. Check with nvcc --version, and for a deeper test compile and run the deviceQuery sample from NVIDIA's cuda-samples repository.

Why does nvidia-smi show a CUDA version that differs from nvcc?

They answer different questions. NVIDIA's nvidia-smi documentation describes the "CUDA Version" field as the version of the CUDA user-mode driver, meaning the latest CUDA version the installed driver supports. It says this is usually, but not always, the version of the toolkit installed. nvcc --version reports the toolkit that is actually on your PATH. A driver that supports a newer CUDA than your toolkit is fine. NVIDIA's CUDA Compatibility guide also describes minor version compatibility (available since CUDA 11), which lets applications built within the same major CUDA release family run on a sufficiently new driver, with some feature limitations.

How do you install PyTorch with GPU support?

1. Check your Python version. PyTorch's install page says the latest stable release requires Python 3.11 or later. Ubuntu's package archive lists python3 3.10.6 for 22.04 and 3.12.3 for 24.04. So 24.04 works with its default Python, while 22.04 needs a newer Python installed separately.
2. Create an isolated environment for each project, for example python3 -m venv ~/envs/project1 and then source ~/envs/project1/bin/activate. (Our advice: one environment per project keeps framework versions from colliding.)
3. Get the install command from pytorch.org. Open pytorch.org/get-started/locally, choose Linux, Pip, Python and the CUDA version, and copy the command it prints. PyTorch notes that often the latest CUDA version is the better choice. Do not copy a pip command from an old blog post, because wheel names and CUDA versions change.
4. Verify the GPU is usable. PyTorch's page gives two checks: build a random tensor to confirm the install, then run import torch and torch.cuda.is_available(). It should return True.

Why no toolkit step? PyTorch's own page only asks you to install CUDA when building from source, and PyPI's metadata for the Linux PyTorch wheel lists NVIDIA's libraries (cuDNN, NCCL and others) as dependencies that pip installs for you. The pre-built route therefore needs the driver, not a system toolkit. Other frameworks differ, so read their install page the same way.

How do you use the GPU inside Docker?

Containers make AI setups repeatable: the same image runs on your workstation, a colleague's machine and a server. For GPU access you need three things on the host, in this order.

1. The NVIDIA driver (done above). NVIDIA's Container Toolkit guide lists it as the prerequisite and recommends installing it with your distribution's package manager.
2. Docker Engine. Docker's Ubuntu instructions require a 64-bit Ubuntu 22.04, 24.04 or a newer release, say to remove conflicting unofficial packages first, and install from Docker's apt repository: sudo apt install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin. Check it with sudo docker run hello-world.
3. The NVIDIA Container Toolkit. Add NVIDIA's apt repository and install the nvidia-container-toolkit packages exactly as NVIDIA's install page shows (it pins a version variable). Then configure Docker: sudo nvidia-ctk runtime configure --runtime=docker, which edits /etc/docker/daemon.json, and restart with sudo systemctl restart docker.

Test it with NVIDIA's sample workload: sudo docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi. If the container prints the same GPU table you saw on the host, GPU access works.

Two cautions from the vendors' own docs. First, Docker states that the docker group grants root-level privileges to its members, so add only trusted users to it. Second, NVIDIA documents a known issue on systems using systemd cgroup drivers: containers can lose access to their GPUs when systemctl daemon-reload is run. If a running container suddenly cannot see the GPU, check NVIDIA's troubleshooting page before assuming a hardware fault.

For multi-service setups that go beyond one machine, our Docker and Kubernetes home-server guide covers the container side, and the virtualization use-case page covers running several workloads on one host.

What goes wrong most often, and where to look first?

Symptom	Most likely layer	What to check
lspci shows no NVIDIA device	Hardware	Seating, PCIe power cables, slot. NVIDIA suggests update-pciids first.
nvidia-smi fails after a driver install	Driver or Secure Boot	Reboot first. If you used DKMS modules on a Secure Boot machine, Ubuntu notes they are not signed with Canonical's key. Prefer the ubuntu-drivers pre-built route.
Driver conflicts after mixing installers	Driver	Ubuntu warns that drivers from outside sources can overwrite packaged ones. Stay with one source.
deviceQuery says no CUDA-capable device	Driver or device files	NVIDIA says this likely means the /dev/nvidia* files are missing or have wrong permissions.
nvcc: command not found	CUDA Toolkit PATH	Add the toolkit bin directory to PATH, or install the toolkit if you never did.
torch.cuda.is_available() returns False	Driver or wrong build	Run nvidia-smi first. Then confirm you installed the CUDA build of PyTorch from the selector, not the CPU build.
Container cannot see the GPU	Container runtime	Did you run nvidia-ctk runtime configure and restart Docker? Did you pass --gpus all? Check the systemd cgroup issue above.

How should you keep an AI workstation stable after setup?

These habits are our practical advice, not vendor requirements, but they prevent the most common self-inflicted outages:

• Write down the versions that work: Ubuntu release, driver branch, PyTorch version, CUDA build. Put them in a README next to your project.
• Do not upgrade the driver in the middle of a long training run. Schedule updates between jobs and reboot afterwards.
• Use one environment or container per project and keep a requirements file, so one project's upgrade cannot break another.
• Keep checkpoints on separate storage from the OS drive, and back them up. A reinstall should cost you an afternoon, not your results. Our RAID levels explained guide covers drive redundancy, which is not a backup.
• Watch temperatures and power during the first long run. Sustained load exposes cooling and power-supply limits that short tests hide. See our guides on why workstations thermal throttle and UPS sizing for servers and workstations.

Which hardware should the software run on?

A clean software stack still needs the right machine under it. The decisions that matter most, with our guides for each:

• GPU and VRAM: VRAM decides which models fit. See how much VRAM local LLMs need, fine-tuning VRAM and RAM, and RTX vs RTX A-series vs Tesla for the driver-licence and cooling differences between GPU families.
• System memory: our AI and ML use-case page suggests system RAM of about twice total VRAM as a floor. See how much RAM by profession and ECC vs non-ECC RAM for why long runs benefit from error-correcting memory.
• Storage: NVMe for datasets and checkpoints so the GPU is not left waiting.
• Cloud or on-premise: if you are weighing rented GPUs against owning a machine, read on-prem AI server vs cloud GPU and custom AI training server cost.
• Teams and sharing: several people on one box is a different problem. See shared GPU workstation for a data science team and private ChatGPT-style LLM server hardware.

Where does ProStation Systems fit?

ProStation Systems builds brand-new custom tower servers and workstations in India, and the build is shaped around your workload, not a fixed catalogue box. For AI work our server tiers map like this, using the specifications published on our site:

Tier	CPU	Memory	Storage	GPU option
Starter	Intel Xeon E (entry)	16-64 GB ECC DDR4	1 TB NVMe SSD	Optional
Pro	Intel Xeon Scalable	64-256 GB ECC DDR4	2 TB NVMe SSD plus HDD options	Optional NVIDIA RTX
Ultra	Dual Intel Xeon Scalable	256-512 GB ECC DDR4	4 TB+ NVMe SSD (RAID)	NVIDIA RTX / A-series / Tesla

Three points are relevant to the setup work in this guide. Our FAQ lists Ubuntu Server among the operating systems we pre-install, and asks you to mention any other OS in your quote request, so you can name the Ubuntu LTS release you want. Our free consulting includes an OS and software compatibility check, which is the right place to tell us your frameworks (PyTorch, TensorFlow, Docker and so on) before the configuration is fixed. Every build carries a 1-year warranty, with extended cover available (see our warranty page). We have not described the GPU driver stack as pre-installed in this guide, because that is something to confirm in your own quote.

If your work is research or lab based, our AI/ML labs page describes the kinds of teams we build for, and the startup AI lab hardware roadmap shows how to grow from one workstation to a shared setup.

Frequently Asked Questions

Q1. Do I need to install the CUDA Toolkit to use PyTorch?
Not for the standard route. PyTorch's install page tells you to pick a pre-built binary with pip and choose the CUDA version that suits your machine. It lists installing CUDA yourself only for building PyTorch from source. What you always need is a working NVIDIA driver. Install the full CUDA Toolkit when you want to compile your own CUDA code with nvcc or when a project's build instructions ask for it.

Q2. What does the CUDA Version shown by nvidia-smi mean?
NVIDIA's nvidia-smi documentation describes it as the CUDA user-mode driver version, which is the latest CUDA version the installed driver supports. It is usually, but not always, the version of the CUDA Toolkit installed on the machine. A machine can show CUDA Version 13.x in nvidia-smi and still have no toolkit installed, or a different toolkit version. Check the toolkit separately with nvcc --version.

Q3. Which NVIDIA driver should I install on Ubuntu for AI work?
Ubuntu's server documentation says it packages two kinds of drivers: Unified Driver Architecture drivers for general desktop use, and Enterprise Ready Drivers (the -server packages) which it recommends for servers and computing tasks. Run sudo ubuntu-drivers list --gpgpu to see what is offered for your GPU and sudo ubuntu-drivers install --gpgpu to let Ubuntu pick the best match. If your machine uses Secure Boot, this tool installs pre-built, signed modules by default.

Q4. Can I run AI workloads in Docker with a GPU?
Yes. Install the NVIDIA GPU driver on the host, install Docker Engine, then install and configure the NVIDIA Container Toolkit. NVIDIA's sample workload, sudo docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi, confirms that a container can see the GPU. The driver lives on the host; the container does not need its own driver.

Q5. Which Ubuntu version should an AI workstation run?
NVIDIA's CUDA installation guide lists Ubuntu 22.04 LTS, 24.04 LTS and 26.04 LTS as supported distributions for the current toolkit, and Docker Engine supports the same LTS releases. A practical point on Python: Ubuntu's package archive lists python3 3.10.6 for 22.04 and 3.12.3 for 24.04, while PyTorch's install page says the latest stable release needs Python 3.11 or later. On 22.04 you will therefore need a newer Python than the system default, for example in a separate environment.

Q6. How much RAM and storage does an AI workstation need before I set up the software?
It depends on your models and datasets. As a starting rule, our AI/ML use-case page suggests system RAM of about twice your total GPU VRAM as a floor, with fast NVMe storage so the GPU is not left waiting for data. Our guides on local LLM VRAM and on RAM by profession give the working-set method in detail. Decide the hardware from your workload first, then set up the software.

Planning an AI workstation?

Tell us the models or frameworks you run, your dataset sizes and how many people will share the machine. ProStation Systems will recommend a tier and a GPU, and send an itemised quote. Use the customize form or book a free consultation.

Sources (read 15 November 2026): NVIDIA CUDA Installation Guide for Linux, NVIDIA CUDA Quick Start Guide, NVIDIA CUDA Compatibility, NVIDIA nvidia-smi documentation, Ubuntu Server: NVIDIA drivers installation, PyTorch: Get Started Locally, Docker Engine on Ubuntu, Docker Linux post-installation steps, NVIDIA Container Toolkit installation and sample workload. Version numbers change quickly, so always take driver branches and PyTorch commands from these pages at the time you install.