Deep Learning Frameworks

Deep learning (DL) frameworks offer building blocks for designing, training and serving deep neural networks through a high-level programming interface. The widely-used DL frameworks, such as PyTorch, JAX, PyTorch Geometric (PyG), vLLM, SGLang and others, rely on GPU-accelerated libraries, such as cuDNN, cuBLAS, CUTLASS, NCCL, and DALI to deliver high-performance, multi-gpu, multi-node accelerated training, post-trainnig and serving.

NVIDIA-Optimized DL Frameworks

Developers, ML researchers, and data scientists can get easy access to NVIDIA optimized DL framework containers with DL examples that are performance-tuned and tested on NVIDIA GPUs against CUDA toolkit and libraries. This eliminates the need to manage packages and dependencies or build DL frameworks from source. Containerized DL frameworks, with all dependencies included, provide a frictionless place to start developing common applications, such as conversational AI, natural language understanding (NLU), recommenders, and computer vision. Visit the NVIDIA NGC™ catalog DL Frameworks collection to learn more.

PyTorch

PyTorch is a Python package that provides two high-level features:

Tensor computation (like numpy) with strong GPU acceleration.
Deep Neural Networks (DNNs) built on a tape-based autograd system.

Reuse your favorite Python packages, such as numpy, scipy and Cython, to extend PyTorch when needed.

PyTorch on NGC Sample models Automatic mixed precision

Model Deployment

For high performance inference deployment for PyTorch trained models:

Use the Torch-TensorRT integration to optimize and deploy models within PyTorch.
Export the PyTorch model to ONNX format, and import, optimize and deploy with NVIDIA TensorRT, an SDK for high performance deep learning inference.

Learning Resources

JAX

JAX is a Python library designed for high-performance numerical computing and machine learning research. JAX can automatically differentiate native Python and implement the NumPy API. With just a few lines of code change, JAX enables distributed training across multi-node, multi-GPU systems, with accelerated performance through XLA.

JAX on NGC JAX on GitHub XLA on GitHub JAX Toolbox

Learning Resources

vLLM

vLLM is a fast and easy-to-use library for LLM/VLM inference and serving. Originally developed in the Sky Computing Lab at UC Berkeley, vLLM has evolved into a community-driven project with contributions from both academia and industry. Built to make LLM inference on NVIDIA fast, scalable, and accessible to everyone — delivered through simple, OpenAI-compatible APIs and rapid day-0 support for the latest models and NVIDIA hardware.

vLLM on NGC

SGLang

SGLang is a fast serving framework for large language models and vision language models. Designed to deliver low-latency, high-throughput inference from a single GPU to large distributed clusters. With native support for popular models, Hugging Face and OpenAI-compatible APIs, it is designed to make production-grade inference accessible to anyone.

SGLang on NGC

PyG

PyG (PyTorch Geometric) is a library built upon PyTorch to easily write and train Graph Neural Networks (GNNs) for a wide range of applications related to structured data.

PyG on NGC