AutomateEdge · ML Engineer · Tools

ML Engineer Tools: 2026 Install Path

This inventory covers the tools a 2026 ML Engineer will use daily or weekly. The list is intentionally deep on the training and serving stack — these are the tools that separate a productive ML Engineer from someone who is blocked by their own tooling. Every item here is either open source or has a usable free tier for learning. Hours-to-learn estimates represent the time until you can solve a novel problem unaided; production mastery takes 3-5x longer.

Core training framework — PyTorch

PyTorch 2.x is the lingua franca of ML engineering in 2026. The depth of knowledge required here exceeds what a generalist or AI Engineer needs — you are programming the framework, not just using it. torch.compile, custom autograd functions, and FSDP are all expected fluency for senior ML Engineers.

ToolPurposeHours to learnDocs
PyTorch 2.x Primary training framework — tensor ops, autograd, nn.Module, distributed primitives 60h Official docs →
PyTorch FSDP FullyShardedDataParallel for training models that don't fit on one GPU 20h Official docs →
torch.compile + Triton Graph compilation for 1.5-3x training throughput; write custom Triton kernels when needed 25h Official docs →
TorchMetrics GPU-accelerated, distributed-safe evaluation metrics library 4h Official docs →
einops Readable tensor reshape / rearrange operations — eliminates view/permute bugs 3h Official docs →

JAX ecosystem — research-leaning roles

JAX is the dominant compute paradigm in Google Brain / DeepMind lineage and is gaining ground at Anthropic and other frontier labs. Its functional transform model (jit, grad, vmap, pmap, scan) produces exceptionally clean research code. Not required for applied-product ML roles; high ROI for frontier-lab positions.

ToolPurposeHours to learnDocs
JAX Composable function transforms: jit, grad, vmap, pmap — the research ML compute engine 40h Official docs →
Flax / Equinox Neural network libraries for JAX; Flax is Google's; Equinox is the community favourite for clean code 20h Official docs →
Optax Gradient processing and optimiser library for JAX 8h Official docs →

Training infrastructure and experiment tracking

PyTorch Lightning, Hugging Face Accelerate, and Mosaic Composer sit above raw PyTorch and reduce boilerplate without hiding it. Experiment tracking with W&B or MLflow is non-negotiable — a training run that wasn't tracked didn't happen from a reproducibility standpoint.

ToolPurposeHours to learnDocs
Hugging Face Accelerate Minimal-overhead distributed training wrapper — one `accelerate launch` and your single-GPU script runs on N GPUs 10h Official docs →
PyTorch Lightning Training loop abstraction with built-in distributed, mixed-precision, and checkpoint support 15h Official docs →
MosaicML Composer Training library with built-in efficiency methods (progressive resizing, ALiBi, BlurPool) — popular for pretraining runs 15h Official docs →
Weights & Biases (W&B) Experiment tracking, hyperparameter sweeps, artifact versioning, dashboards — the de-facto standard 8h Official docs →
MLflow Open-source experiment tracking + model registry; preferred in enterprises with strict data-residency requirements 10h Official docs →
Optuna Hyperparameter optimisation — sampler-agnostic, integrates with W&B, supports pruning of bad runs 8h Official docs →

Hugging Face ecosystem

The Hugging Face platform is the central repository for pretrained models, datasets, and the tooling to work with them. For ML Engineers, Transformers, PEFT, TRL, and Datasets are daily tools. The fine-tuning and alignment workflows (LoRA, QLoRA, DPO, GRPO) all live here in 2026.

ToolPurposeHours to learnDocs
Hugging Face Transformers 5.x Pretrained model zoo — Llama 4, Qwen 3, Gemma 3, Mistral, Phi, BERT variants, ViT, Whisper 20h Official docs →
Hugging Face Datasets Versioned, streaming dataset access — the right way to handle multi-hundred-GB training corpora 8h Official docs →
PEFT (Parameter-Efficient Fine-Tuning) LoRA, QLoRA, IA3, prefix tuning — fine-tune 70B models on a single H100 without full weight updates 12h Official docs →
TRL (Transformer Reinforcement Learning) SFT, DPO, PPO, GRPO trainers for instruction-tuning and alignment pipelines 15h Official docs →
Hugging Face Hub Model + dataset + space hosting — your versioned model registry for open-source work 3h Official docs →

Distributed compute — DeepSpeed, Ray, Spark

Large-scale ML is distributed by default. DeepSpeed and Ray are the two most important additions beyond PyTorch's native distributed module. Spark handles the data-scale problems (terabyte-scale feature computation) that PyTorch doesn't touch.

ToolPurposeHours to learnDocs
DeepSpeed ZeRO stages 1/2/3, optimizer offload, mixed precision — doubles or quadruples effective model size per GPU 20h Official docs →
Ray + Ray Train Distributed Python + distributed ML training — scales from laptop to 1000-node cluster with the same API 20h Official docs →
Apache Spark (PySpark) Terabyte-scale batch feature computation, data deduplication, and ETL preprocessing 30h Official docs →
Dask Pandas-compatible distributed dataframes for medium-scale (100GB-10TB) data work without Spark's JVM overhead 12h Official docs →

Model serving

Inference is not just 'call model.forward()' at scale. Dynamic batching, KV cache management, quantisation, and GPU utilisation are the differences between a serving system that costs $2k/month and one that costs $12k/month for the same traffic. The tools below are the production-grade options ML Engineers are expected to know.

ToolPurposeHours to learnDocs
vLLM PagedAttention-based LLM serving — the state-of-the-art for throughput-efficient transformer inference 12h Official docs →
SGLang Structured generation language + serving runtime — faster than vLLM on constrained-generation workloads 10h Official docs →
NVIDIA Triton Inference Server Model serving for heterogeneous backends (PyTorch, ONNX, TensorRT, custom C++) with dynamic batching 20h Official docs →
TorchServe PyTorch-native model server with MAR packaging, handler lifecycle, and multi-model management 12h Official docs →
Ray Serve Python-native model serving with dynamic scaling, A/B traffic splitting, and multi-model composition 10h Official docs →
ONNX + ONNX Runtime Portable model format + optimised inference runtime for CPU, GPU, and edge devices 10h Official docs →

Data versioning and feature stores

Reproducible ML requires versioned data, not just versioned code. Feature stores solve the online/offline consistency problem that causes training-serving skew — the most pernicious silent bug in production ML systems.

ToolPurposeHours to learnDocs
DVC (Data Version Control) Git for ML datasets + model artifacts — version, share, and reproduce data-dependent experiments 12h Official docs →
Feast (open-source feature store) Point-in-time correct feature retrieval for training; low-latency online retrieval for serving 20h Official docs →
Delta Lake / Iceberg ACID-compliant table format for the data lake — time-travel queries, schema evolution, and upsert for ML pipelines 15h Official docs →
Great Expectations / Soda Data quality validation — assert schema, distributions, and nullability constraints before data enters the training pipeline 10h Official docs →

GPU infrastructure and low-level tooling

An ML Engineer who cannot use a profiler is flying blind on performance. CUDA and NVCC basics are not for writing kernels from scratch — they are for reading what the profiler is telling you and understanding why a kernel is slow. Mixed-precision and quantisation tooling is part of the daily inference-cost optimisation workflow.

ToolPurposeHours to learnDocs
NVIDIA Nsight Systems / Compute GPU profiling — timeline traces, kernel occupancy, memory bandwidth utilisation 15h Official docs →
bitsandbytes INT8 / INT4 quantisation for inference and QLoRA training — the default quantisation backend for Hugging Face PEFT 6h Official docs →
CUDA Toolkit basics (nvcc, device query) Understand thread/block/warp model, coalesced memory access, shared memory, warp divergence — enough to read Triton kernels 20h Official docs →
FlashAttention (flash-attn) IO-aware exact attention — 2-4x training speedup and linear memory for long sequences 5h Official docs →

MLOps platforms and pipeline orchestration

Training a model once is a research exercise. Training it reliably on a schedule, with data quality gates, drift alerting, and automated rollback, is ML engineering. Pick one orchestration framework to go deep on; be conversational about the others.

ToolPurposeHours to learnDocs
Kubeflow Pipelines Kubernetes-native ML pipeline orchestration — the standard at large engineering orgs with Kubernetes clusters 25h Official docs →
ZenML Python-first MLOps framework — pipeline steps as decorated functions, stack-agnostic, integrates with W&B + MLflow + Kubeflow 15h Official docs →
Metaflow (Netflix) Data-science-friendly pipeline framework with automatic versioning and local-to-cloud scaling 12h Official docs →
Prefect / Airflow General-purpose workflow orchestration; useful for ML pipelines that mix data-engineering and ML steps 15h Official docs →
Evidently AI Data drift + model monitoring reports — the open-source monitoring layer for production ML 8h Official docs →

Quickstart install path

A step-by-step first-day install that gets you from zero to producing work. Follow in order — later steps assume earlier ones.

  1. Install Python 3.11+ via pyenv or uv. Create a project environment: `uv init ml-engineer-track && cd ml-engineer-track`.
  2. Install the core training stack: `uv add torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124` (adjust CUDA version to your GPU). Add `einops transformers datasets accelerate peft trl`.
  3. Install experiment tracking: `uv add wandb mlflow optuna`. Sign up for a free W&B account and run `wandb login`.
  4. Install distributed and serving dependencies: `uv add deepspeed ray[train,serve] vllm`. Note: DeepSpeed requires a C++ compiler; on Ubuntu `apt install build-essential`.
  5. Install data tooling: `uv add dvc great_expectations pyspark dask`.
  6. Install profiling and quantisation: `uv add bitsandbytes flash-attn` (flash-attn requires a CUDA-capable GPU at install time).
  7. Verify GPU access: `python -c 'import torch; print(torch.cuda.get_device_name(0))'`. If not available, set up a cloud instance (Lambda Labs, RunPod) with CUDA 12.4+.
  8. Clone the Hugging Face PEFT examples repo and run the LoRA fine-tuning example end-to-end as a sanity check on your environment.

Frequently asked questions

Do I need a local GPU to get started?

No. Weeks 1-10 require no GPU — math, PyTorch on CPU, and small single-GPU experiments on Colab Free. From week 11 onward, rent on Lambda Labs, Vast.ai, or RunPod. A 1x H100 node for 3 hours costs about $10-15 and covers most distributed training exercises.

Should I use Lightning or raw PyTorch?

Write raw training loops first — every PyTorch Lightning abstraction makes more sense after you have written what it is replacing. Switch to Lightning or Accelerate when the boilerplate is costing you more time than the abstractions save. Most interview environments want you to debug raw PyTorch, not Lightning internals.

Is vLLM the right serving stack for all use cases?

For autoregressive LLM serving, vLLM is currently the highest-throughput open-source option and the right default. For non-LLM models (CNNs, BERT-class encoders, custom architectures), Triton Inference Server with TensorRT is often faster. For multi-model ensemble pipelines, Triton's ensemble model type is cleaner than vLLM's. Pick the tool that matches the workload.

Is DeepSpeed or FSDP better?

For PyTorch-native codebases, FSDP is now the recommended path — it is maintained by the PyTorch team, has fewer installation dependencies, and integrates cleanly with torch.compile. DeepSpeed has a richer optimizer-offload story (ZeRO-3 CPU/NVMe) and more deployment history at large scale. Many production systems use both: FSDP for training code, DeepSpeed for the optimizer offload extension.

How important is W&B vs MLflow for getting a job?

W&B is the more common tool in ML-native companies and research labs. MLflow is more common in enterprise and data-platform-heavy organisations. Learn W&B to the point you can demonstrate live experiment tracking in an interview; know MLflow's architecture conceptually. Both appear on job descriptions; neither is a hard gate in isolation.

What is the minimum viable MLOps setup for a portfolio project?

Three things: (1) W&B run tracking with hyperparameter logging, loss curves, and artifact versioning, (2) a DVC-tracked dataset so the training run is reproducible from a fresh clone, (3) a model card documenting training config, eval results, and limitations. That is more MLOps rigour than most portfolio projects show and will get positive attention in reviews.

Should I build a RAG system as an ML Engineer?

RAG is primarily an AI Engineer workload — wiring a retrieval layer to a hosted model. As an ML Engineer, your portfolio should emphasise training from scratch, fine-tuning, or serving optimisation. One RAG project is fine for breadth; five is scope creep into a different specialisation.

Continue on the ML Engineer track