• Stars
    star
    288
  • Rank 143,818 (Top 3 %)
  • Language
    C
  • License
    Apache License 2.0
  • Created almost 4 years ago
  • Updated 4 months ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

Go Bindings for the NVIDIA Management Library (NVML)

Go Bindings for the NVIDIA Management Library (NVML)

Table of Contents

Overview

This repository provides Go bindings for the NVIDIA Management Library API (NVML).

At present, these bindings are only supported on Linux.

These bindings are not a reimplementation of NVML in Go, but rather a set of wrappers around the C API provided by libnvidia-ml.so. This library is part of the standard NVIDIA driver distribution, and should be available on any Linux system that has the NVIDIA driver installed. The API is designed to be backwards compatible, so the latest bindings should work with any version of libnvidia-ml.so installed on your system.

Note: A working NVIDIA driver with libnvidia-ml.so is not required to compile code that imports these bindings. However, you will get a runtime error if libnvidia-ml.so is not available in your library path at runtime.

Please see the following link for documentation on the full NVML Go API: http://godoc.org/github.com/NVIDIA/go-nvml/pkg/nvml

Quick Start

All you need is a simple import and a call to nvml.Init() to start using these bindings.

The code below shows an example of using these bindings to query all of the GPUs on your system and print out their UUIDs.

package main

import (
	"fmt"
	"log"

	"github.com/NVIDIA/go-nvml/pkg/nvml"
)

func main() {
	ret := nvml.Init()
	if ret != nvml.SUCCESS {
		log.Fatalf("Unable to initialize NVML: %v", nvml.ErrorString(ret))
	}
	defer func() {
		ret := nvml.Shutdown()
		if ret != nvml.SUCCESS {
			log.Fatalf("Unable to shutdown NVML: %v", nvml.ErrorString(ret))
		}
	}()

	count, ret := nvml.DeviceGetCount()
	if ret != nvml.SUCCESS {
		log.Fatalf("Unable to get device count: %v", nvml.ErrorString(ret))
	}

	for i := 0; i < count; i++ {
		device, ret := nvml.DeviceGetHandleByIndex(i)
		if ret != nvml.SUCCESS {
			log.Fatalf("Unable to get device at index %d: %v", i, nvml.ErrorString(ret))
		}

		uuid, ret := device.GetUUID()
		if ret != nvml.SUCCESS {
			log.Fatalf("Unable to get uuid of device at index %d: %v", i, nvml.ErrorString(ret))
		}

		fmt.Printf("%v\n", uuid)
	}
}

On my DGX workstation, this results in the following output:

$ go run main.go
GPU-edfee158-11c1-52b8-0517-92f30e7fac88
GPU-f22fb098-d1b3-3806-2655-ba25f02229c1
GPU-f613f823-1032-b3ec-a876-50f2e35e6f9e
GPU-3109fa37-4445-73c7-b695-1b5a4d13f58e
GPU-e28a6529-288c-7ddf-8fea-68c4833cda70
GPU-a27fb382-bad2-c02a-95ba-f6a1da38e76c
GPU-f5bb8d07-ee19-1787-4d9a-a84c4ac6b086
GPU-1ba0ca0e-6d1d-d9db-07d8-c1c5a8c32814

How the bindings are generated

This project leverages two core technologies:

  1. Go's builtin support for cgo (https://golang.org/cmd/cgo/)
  2. A third-party tool called c-for-go (https://c.for-go.com/)

Using these tools, we are able to generate a set of Go bindings for NVML, given nothing more than a specific version of the nvml.h header file (which defines the full NVML API). Most of the process to generate these bindings is automated, but a few manual steps are required in order to make the generated bindings more useful from an end user's perspective.

The basic flow to generate the bindings is therefore to:

  1. Take the nvml.h file and pass it through c-for-go
  2. Take each of the low-level Go bindings generated by c-for-go and wrap them in a more user-friendly API

As an example, consider the Go bindings generated for the nvmlDeviceGetAccountingPids() API call below:

Original API in nvml.h:

nvmlReturn_t nvmlDeviceGetAccountingPids(nvmlDevice_t device, unsigned int *count, unsigned int *pids);

Auto-generated Go bindings from c-for-go:

func nvmlDeviceGetAccountingPids(Device Device, Count *uint32, Pids *uint32) Return {
	cDevice, _ := *(*C.nvmlDevice_t)(unsafe.Pointer(&Device)), cgoAllocsUnknown
	cCount, _ := (*C.uint)(unsafe.Pointer(Count)), cgoAllocsUnknown
	cPids, _ := (*C.uint)(unsafe.Pointer(Pids)), cgoAllocsUnknown
	__ret := C.nvmlDeviceGetAccountingPids(cDevice, cCount, cPids)
	__v := (Return)(__ret)
	return __v
}

Manual wrapper around the auto-generated bindings:

package nvml

func DeviceGetAccountingPids(Device Device) ([]int, Return) {
	var Count uint32 = 1 // Will be reduced upon returning
	for {
		Pids := make([]uint32, Count)
		ret := nvmlDeviceGetAccountingPids(Device, &Count, &Pids[0])
		if ret == SUCCESS {
			return uint32SliceToIntSlice(Pids[:Count]), ret
		}
		if ret != ERROR_INSUFFICIENT_SIZE {
			return nil, ret
		}
		Count *= 2
	}
}

func (Device Device) GetAccountingPids() ([]int, Return) {
	return DeviceGetAccountingPids(Device)
}

This manual wrapper makes it so that users don't have to write the boiler-plate code of figuring out the correct count to pass into the API while at the same time growing the Pids array and turning into a slice. It would be used as follows:

device, _ := nvml.DeviceGetHandleByIndex(0)
pids, _ := device.GetAccountingPids()
...

This is actually one of the more complicated examples. Most of the manual wrappers are very simple and look similar to the following:

Original API in nvml.h:

nvmlReturn_t nvmlDeviceGetUUID(nvmlDevice_t device, char *uuid, unsigned int length);

Auto-generated Go bindings from c-for-go:

func nvmlDeviceGetUUID(Device Device, Uuid *byte, Length uint32) Return {
	cDevice, _ := *(*C.nvmlDevice_t)(unsafe.Pointer(&Device)), cgoAllocsUnknown
	cUuid, _ := (*C.char)(unsafe.Pointer(Uuid)), cgoAllocsUnknown
	cLength, _ := (C.uint)(Length), cgoAllocsUnknown
	__ret := C.nvmlDeviceGetUUID(cDevice, cUuid, cLength)
	__v := (Return)(__ret)
	return __v
}

Manual wrapper around the auto-generated bindings:

package nvml

func DeviceGetUUID(Device Device) (string, Return) {
	Uuid := make([]byte, DEVICE_UUID_BUFFER_SIZE)
	ret := nvmlDeviceGetUUID(Device, &Uuid[0], DEVICE_UUID_BUFFER_SIZE)
	return string(Uuid[:clen(Uuid)]), ret
}

func (Device Device) GetUUID() (string, Return) {
	return DeviceGetUUID(Device)
}

While it does take some effort to take the auto-generated bindings and manually wrap them in the more user-friendly API, this only has to be done once per API call and then never touched again. As such, as new release of NVML come out, only the new API calls will need to be added.

The following section goes into the details of how the code is structured, and what each file's purpose is.

Code Structure

There are two top-level directories in this repository:

  • /gen
  • /pkg

The /gen directory is used to house any code used in the generation of the final Go bindings. The /pkg directory is used to house any static packages associated with this project as well as the actual Go bindings once they have been generated. The one exception is the code used to dynamically load the libnvidia-ml.so code from a host system and attach the go bindings to it. This package requires no generated code and is housed statically under pkg/dl. Once the code under gen/nvml has passed through c-for-go and any manual wrappers applied, the final generated bindings are placed under pkg/nvml.

In general, the code used to generate the NVML Go bindings can be broken into 4 logical parts:

  1. Code defining the NVML API and how any auto-generated bindings should be produced
  2. Code responsible for dynamically loading libnvidia-ml.so on a host system and hooking it up to the bindings
  3. Code bridging the gap between any auto-generated bindings and the manual wrappers around them
  4. The manual wrappers themselves
  5. Test code

Each of these parts is discussed in detail below, along with the files associated with them.

Code defining the NVML API

The following files aid in defining the NVML API and how any auto-generated bindings should be produced from it.

  • gen/nvml/nvml.h
  • gen/nvml/nvml.yml

The nvml.h file is a direct copy of nvml.h from the NVIDIA driver. Since the NVML API is guaranteed to be backwards compatible, we should strive to keep this always up to date with the latest.

Note: The make process modifies nvml.h in that it translates any opaque types defined by nvml.h into something more recognizable by cgo.

For example:

-typedef struct nvmlDevice_st* nvmlDevice_t;
+typedef struct
+{
+   struct nvmlDevice_st* handle;
+} nvmlDevice_t;

The two statements are semantically equivalent in terms of how they are laid out in memory, but cgo will only generate a unique type for nvmlDevice_t when expressed as the latter. When building the bindings we first update nvml.h using sed, and then run c-for-go over it.

Finally, the nvml.yml file is the input file to c-for-go that tells it how to parse nvml.h and auto-generate bindings for it. Please see the c-for-go wiki for more information about the contents of this file and how it works.

Code to load libnvidia-ml.so

The code under pkg/dl is responsible for dynamically loading the libnvidia-ml.so binary from a host system and connecting the go bindings to it. This happens under the hood whenever a user makes an nvml.Init() call. It is transparent to the end user, and should work without any further user-intervention.

Depending on the version of libnvidia-ml.so that is found, certain versioned symbols need to be updated. At the time of this writing, these symbols include the following (as defined in nvml.h):

#ifndef NVML_NO_UNVERSIONED_FUNC_DEFS
    #define nvmlInit                                nvmlInit_v2
    #define nvmlDeviceGetPciInfo                    nvmlDeviceGetPciInfo_v3
    #define nvmlDeviceGetCount                      nvmlDeviceGetCount_v2
    #define nvmlDeviceGetHandleByIndex              nvmlDeviceGetHandleByIndex_v2
    #define nvmlDeviceGetHandleByPciBusId           nvmlDeviceGetHandleByPciBusId_v2
    #define nvmlDeviceGetNvLinkRemotePciInfo        nvmlDeviceGetNvLinkRemotePciInfo_v2
    #define nvmlDeviceRemoveGpu                     nvmlDeviceRemoveGpu_v2
    #define nvmlDeviceGetGridLicensableFeatures     nvmlDeviceGetGridLicensableFeatures_v3
    #define nvmlEventSetWait                        nvmlEventSetWait_v2
    #define nvmlDeviceGetAttributes                 nvmlDeviceGetAttributes_v2
    #define nvmlDeviceGetComputeRunningProcesses    nvmlDeviceGetComputeRunningProcesses_v2
    #define nvmlDeviceGetGraphicsRunningProcesses   nvmlDeviceGetGraphicsRunningProcesses_v2
#endif // #ifndef NVML_NO_UNVERSIONED_FUNC_DEFS

The actual versions that these API calls are assigned to will depend on the version of the NVIDIA driver (and hence the version of libnvidia-ml.so that you have linked in). These updates happen in the updateVersionedSymbols() function of gen/nvml/init.go as seen below.

// Default all versioned APIs to v1 (to infer the types)
var nvmlInit = nvmlInit_v1
var nvmlDeviceGetPciInfo = nvmlDeviceGetPciInfo_v1
var nvmlDeviceGetCount = nvmlDeviceGetCount_v1
...

// updateVersionedSymbols()
func updateVersionedSymbols() {
	ret := nvml.Lookup("nvmlInit_v2")
	if ret == SUCCESS {
		nvmlInit = nvmlInit_v2
	}
	ret = nvml.Lookup("nvmlDeviceGetPciInfo_v2")
	if ret == SUCCESS {
		nvmlDeviceGetPciInfo = nvmlDeviceGetPciInfo_v2
	}
	ret = nvml.Lookup("nvmlDeviceGetPciInfo_v3")
	if ret == SUCCESS {
		nvmlDeviceGetPciInfo = nvmlDeviceGetPciInfo_v3
	}
	ret = nvml.Lookup("nvmlDeviceGetCount_v2")
	if ret == SUCCESS {
		nvmlDeviceGetCount = nvmlDeviceGetCount_v2
	}
	...
}

Whenever a new version of NVML comes out that either (1) adds a new versioned API call, or (2) bumps the version of an existing API call -- we need to make sure and update this function appropriately (as well as make the necessary changes to nvidia.yml to ensure all v1 symbols are imported appropriately).

Code to bridge the auto-generated and manual bindings

The files below define a set of "glue" code between the auto-generated bindings from c-for-go and the manual wrappers providing a more user-friendly API to the end user.

  • gen/nvml/cgo_helpers.go
  • gen/nvml/return.go

The cgo_helpers.go file defines functions that help in dealing with the types coming out of the C API and turning them into more usable Go types. It is actually a stripped down version of the auto-generated cgo_helpers.go file from c-for-go that we have whittled down to the bare essentials. We also define a few of our own functions in here as well. For example, doing things like finding the length of a NULL terminated string inside a byte slice (clen()), and converting a uint32 slice into an int slice (uint32SliceToIntSlice()), etc.

The return.go file simply wraps the Return type created by c-for-go (which is a go-ified version of the nvmlReturn_t type from C) and has it implement the Error interface so it can be returned as a normal Go error type if desired. The string returned as part of the error is the result of calling nvmlErrorString() under the hood.

Manual wrappers around the auto-generated bindings from c-for-go

The following files add manual wrappers around all of the auto-generated bindings from c-for-go. Only these manual wrappers are expected as part of the API for the package -- the auto-generated bindings are only available for internal use.

  • gen/nvml/init.go
  • gen/nvml/system.go
  • gen/nvml/event_set.go
  • gen/nvml/vgpu.go
  • gen/nvml/unit.go
  • gen/nvml/device.go

These wrappers add boiler-plate code around the auto-generated bindings so that the end-user doesn't have to do this themselves every time a call is made.

When appropriate, they also bind functions to the top-level types that are defined (e.g. Unit, Device, EventSet, Vgpu, etc.) so that functions can be called directly on instances of these types instead of using a call at the package scope.

A few examples of a this can be seen below:

// nvml.UnitGetDevices()
func UnitGetDevices(Unit Unit) ([]Device, Return) {
	var DeviceCount uint32 = 1 // Will be reduced upon returning
	for {
		Devices := make([]Device, DeviceCount)
		ret := nvmlUnitGetDevices(Unit, &DeviceCount, &Devices[0])
		if ret == SUCCESS {
			return Devices[:DeviceCount], ret
		}
		if ret != ERROR_INSUFFICIENT_SIZE {
			return nil, ret
		}
		DeviceCount *= 2
	}
}

func (Unit Unit) GetDevices() ([]Device, Return) {
	return UnitGetDevices(Unit)
}


// nvml.DeviceGetUUID()
func DeviceGetUUID(Device Device) (string, Return) {
	Uuid := make([]byte, DEVICE_UUID_BUFFER_SIZE)
	ret := nvmlDeviceGetUUID(Device, &Uuid[0], DEVICE_UUID_BUFFER_SIZE)
	return string(Uuid[:clen(Uuid)]), ret
}

func (Device Device) GetUUID() (string, Return) {
	return DeviceGetUUID(Device)
}


// nvml.EventSetWait()
func EventSetWait(Set EventSet, Timeoutms uint32) (EventData, Return) {
	var Data EventData
	ret := nvmlEventSetWait(Set, &Data, Timeoutms)
	return Data, ret
}

func (Set EventSet) Wait(Timeoutms uint32) (EventData, Return) {
	return EventSetWait(Set, Timeoutms)
}


// nvml.VgpuInstanceGetUUID()
func VgpuInstanceGetUUID(VgpuInstance VgpuInstance) (string, Return) {
	Uuid := make([]byte, DEVICE_UUID_BUFFER_SIZE)
	ret := nvmlVgpuInstanceGetUUID(VgpuInstance, &Uuid[0], DEVICE_UUID_BUFFER_SIZE)
	return string(Uuid[:clen(Uuid)]), ret
}

func (VgpuInstance VgpuInstance) GetUUID() (string, Return) {
	return VgpuInstanceGetUUID(VgpuInstance)
}

Whenever a new version of NVML comes out that adds new API calls, a new set of manual wrappers will need to be added to keep the API up-to-date. Adding the initial set of wrappers was very time consuming, but adding additional wrappers should be straightforward so long as we keep good pace with each new release.

Test code

At present, all test code is under the following file:

  • gen/nvml/nvml_test.go

The test coverage is fairly sparse and could be greatly improved.

Building and Testing

Building and testing the bindings is fairly straight-forward. The only prerequisite is a working installation of c-for-go from https://github.com/xlab/c-for-go.

Note: Please check the Makefile for the specific version of c-for-go used.

Once this is available, just run the sequence below to build and test these NVML Go bindings. The generated bindings will be placed under go-nvml/pkg/nvml.

$ make
c-for-go -out pkg gen/nvml/nvml.yml
  processing gen/nvml/nvml.yml done.
cp gen/nvml/*.go pkg/nvml
cd pkg/nvml; \
    go tool cgo -godefs types.go > types_gen.go; \
    go fmt types_gen.go; \
cd -> /dev/null
types_gen.go
rm -rf pkg/nvml/types.go pkg/nvml/_obj

$ make test
cd pkg/nvml; \
    go test -v .; \
cd -> /dev/null
=== RUN   TestInit
    TestInit: nvml_test.go:26: Init: Success
    TestInit: nvml_test.go:33: Shutdown: Success
--- PASS: TestInit (0.06s)
=== RUN   TestSystem
    TestSystem: nvml_test.go:45: SystemGetDriverVersion: Success
    TestSystem: nvml_test.go:46:   version: 410.104
    TestSystem: nvml_test.go:53: SystemGetNVMLVersion: Success
    TestSystem: nvml_test.go:54:   version: 10.410.104
    TestSystem: nvml_test.go:61: SystemGetCudaDriverVersion: Success
    TestSystem: nvml_test.go:62:   version: 10000
    TestSystem: nvml_test.go:69: SystemGetCudaDriverVersion_v2: Success
    TestSystem: nvml_test.go:70:   version: 10000
    TestSystem: nvml_test.go:77: SystemGetProcessName: Success
    TestSystem: nvml_test.go:78:   name: /lib/systemd/s
    TestSystem: nvml_test.go:85: SystemGetHicVersion: Success
    TestSystem: nvml_test.go:86:   count: 0
    TestSystem: nvml_test.go:96: SystemGetTopologyGpuSet: Success
    TestSystem: nvml_test.go:97:   count: 4
    TestSystem: nvml_test.go:99:   device[0]: {0x7f5875284408}
    TestSystem: nvml_test.go:99:   device[1]: {0x7f5875298e90}
    TestSystem: nvml_test.go:99:   device[2]: {0x7f58752ad918}
    TestSystem: nvml_test.go:99:   device[3]: {0x7f58752c23a0}
--- PASS: TestSystem (0.10s)
=== RUN   TestUnit
    TestUnit: nvml_test.go:112: UnitGetCount: Success
    TestUnit: nvml_test.go:113:   count: 0
    TestUnit: nvml_test.go:117: Skipping test with no Units.
--- SKIP: TestUnit (0.06s)
=== RUN   TestEventSet
    TestEventSet: nvml_test.go:253: EventSetCreate: Success
    TestEventSet: nvml_test.go:254:   set: {0x2122f10}
    TestEventSet: nvml_test.go:261: EventSetWait: Timeout
    TestEventSet: nvml_test.go:262:   data: {{<nil>} 0 0 0 0}
    TestEventSet: nvml_test.go:269: EventSet.Wait: Timeout
    TestEventSet: nvml_test.go:270:   data: {{<nil>} 0 0 0 0}
    TestEventSet: nvml_test.go:277: EventSetFree: Success
    TestEventSet: nvml_test.go:285: EventSet.Free: Success
--- PASS: TestEventSet (0.06s)
PASS
ok  github.com/NVIDIA/go-nvml/pkg/nvml 0.283s

Note: A working NVIDIA driver with libnvidia-ml.so is not required to compile code that imports these bindings. However, you will get a runtime error if libnvidia-ml.so is not available in your library path at runtime.

Updating the Code

The general steps to update the bindings to a newer version of the NVML API are as follows:

Update nvml.h

Pull down the nvml.h containing the updated API and commit it back to gen/nvml/nvml.h. The Makefile contains a command:

$ make update-nvml-h
Found 5 NVML packages:

No.  Version   Upload Time          Package
  1  11.5.50   2021-11-23-22:46:02  nvidia/cuda-nvml-dev/11.5.50/linux-64/cuda-nvml-dev-11.5.50-h511b398_0.tar.bz2
  2  11.4.120  2021-11-03-22:08:33  nvidia/cuda-nvml-dev/11.4.120/linux-64/cuda-nvml-dev-11.4.120-hb8c74d6_0.tar.bz2
  3  11.4.43   2021-09-08-00:10:30  nvidia/cuda-nvml-dev/11.4.43/linux-64/cuda-nvml-dev-11.4.43-he36855d_0.tar.bz2
  4  11.3.58   2021-09-08-00:36:34  nvidia/cuda-nvml-dev/11.3.58/linux-64/cuda-nvml-dev-11.3.58-hc25e488_0.tar.bz2
  5  11.3.58   2021-09-08-00:36:31  nvidia/cuda-nvml-dev/11.3.58/linux-64/cuda-nvml-dev-11.3.58-h70090ce_0.tar.bz2

Pick an NVML package to update ([1]-5): 1

NVML version: 11.5.50
Package: nvidia/cuda-nvml-dev/11.5.50/linux-64/cuda-nvml-dev-11.5.50-h511b398_0.tar.bz2

Updating nvml.h to 11.5.50 from https://api.anaconda.org/download/nvidia/cuda-nvml-dev/11.5.50/linux-64/cuda-nvml-dev-11.5.50-h511b398_0.tar.bz2 ...
Successfully updated nvml.h to 11.5.50.

that copies the file from the Anaconda package anaconda.org/nvidia/cuda-nvml-dev. Available files can be found at https://anaconda.org/nvidia/cuda-nvml-dev/files (platform: linux-64).

Since gen/nvml/nvml.h is under version control, running:

git diff -w

(ignoring whitespace) will show us which new API calls there are.

Add new versioned APIs

If there are changes to the versioned APIs (defined as in the #ifndef NVML_NO_UNVERSIONED_FUNC_DEFS block in gen/nvml/nvml.h) nvml.yml and init.go must be updated accordingly.

The modified versioned calls can be found bu running:

git diff -w gen/nvml/nvml.h | grep -E "^\+\s*#define.*?_v[^1]"

Add manual wrappers

Write a set of manual wrappers around any new calls as described in one of the previous sections above.

The following command should show the API calls added in the update:

git diff -w gen/nvml/nvml.h | grep "+nvmlReturn_t DECLDIR nvml"

Note that these includes the new versions of existing calls -- which should already have been handled in the previous section. To exclude these run:

git diff -w gen/nvml/nvml.h | grep "+nvmlReturn_t DECLDIR nvml" | grep -vE "_v\d+\("

Of course this is just the general flow, and there may be more work to do if new types are added, or a new API is created that does something outside the scope of what has been done so far. These guidelines should be a good starting point though.

Keep in mind, that all updates to the NVML bindings code should be made in the gen/ directory of the repository. Only when releasing new bindings will this code be processed and pushed to the pkg/ directory for release.

Releasing

Once the code in gen/ has been fully updated to support a particular version of NVML, a new release should be created.

As part of the release, two things need to happen:

  1. A new set of bindings needs to be generated from the code under gen/ and committed into pkg/
  2. A tag with the appropriate NVML release needs to be added to the repo and pushed upstream.

An example of this workflow for the 11.0 release of NVML can be seen below:

# Commit the generated bindings back to main
git checkout main
make
git add -f pkg/
git commit -m "Add bindings for v11.0 of the NVML API"
git push origin

# Tag the repo with the version number and push it upstream
git tag v11.0
git push origin v11.0

If updates need to be made against a particular version (due to bugs in the bindings code, for example), then we append a -<revision> number to the version tag we push.

For example:

git checkout v11.0
git checkout -b bug-fixes-for-v11.0
... fix bugs and commit
git tag v11.0-1
git push v11.0-1

Since the NVML API is designed to be backwards compatible, we envision it being rare to require such backports (because people can just use the latest bindings instead of relying on a particular version). However, we may perform such backports from time-to-time as deemed necessary (or upon request).

Contributing

Please see the file CONTRIBUTING.md for details on how to contribute to this project.

More Repositories

1

nvidia-docker

Build and run Docker containers leveraging NVIDIA GPUs
16,896
star
2

open-gpu-kernel-modules

NVIDIA Linux open GPU kernel module source
C
14,997
star
3

DeepLearningExamples

State-of-the-Art Deep Learning scripts organized by models - easy to train and deploy with reproducible accuracy and performance on enterprise-grade infrastructure.
Jupyter Notebook
13,339
star
4

NeMo

A scalable generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (Automatic Speech Recognition and Text-to-Speech)
Python
12,016
star
5

FastPhotoStyle

Style transfer, deep learning, feature transform
Python
11,020
star
6

TensorRT

NVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.
C++
10,618
star
7

Megatron-LM

Ongoing research training transformer models at scale
Python
10,332
star
8

TensorRT-LLM

TensorRT-LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and build TensorRT engines that contain state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT-LLM also contains components to create Python and C++ runtimes that execute those TensorRT engines.
C++
8,542
star
9

vid2vid

Pytorch implementation of our method for high-resolution (e.g. 2048x1024) photorealistic video-to-video translation.
Python
8,482
star
10

apex

A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch
Python
8,239
star
11

pix2pixHD

Synthesizing and manipulating 2048x1024 images with conditional GANs
Python
6,488
star
12

cuda-samples

Samples for CUDA Developers which demonstrates features in CUDA Toolkit
C
6,119
star
13

cutlass

CUDA Templates for Linear Algebra Subroutines
C++
5,519
star
14

FasterTransformer

Transformer related optimization, including BERT, GPT
C++
5,313
star
15

DALI

A GPU-accelerated library containing highly optimized building blocks and an execution engine for data processing to accelerate deep learning training and inference applications.
C++
5,048
star
16

thrust

[ARCHIVED] The C++ parallel algorithms library. See https://github.com/NVIDIA/cccl
C++
4,914
star
17

tacotron2

Tacotron 2 - PyTorch implementation with faster-than-realtime inference
Jupyter Notebook
4,562
star
18

warp

A Python framework for high performance GPU simulation and graphics
Python
4,206
star
19

DIGITS

Deep Learning GPU Training System
HTML
4,105
star
20

NeMo-Guardrails

NeMo Guardrails is an open-source toolkit for easily adding programmable guardrails to LLM-based conversational systems.
Python
4,064
star
21

nccl

Optimized primitives for collective multi-GPU communication
C++
3,187
star
22

flownet2-pytorch

Pytorch implementation of FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks
Python
2,938
star
23

ChatRTX

A developer reference project for creating Retrieval Augmented Generation (RAG) chatbots on Windows using TensorRT-LLM
TypeScript
2,635
star
24

k8s-device-plugin

NVIDIA device plugin for Kubernetes
Go
2,481
star
25

libcudacxx

[ARCHIVED] The C++ Standard Library for your entire system. See https://github.com/NVIDIA/cccl
C++
2,294
star
26

GenerativeAIExamples

Generative AI reference workflows optimized for accelerated infrastructure and microservice architecture.
Python
2,192
star
27

nvidia-container-toolkit

Build and run containers leveraging NVIDIA GPUs
Go
2,171
star
28

waveglow

A Flow-based Generative Network for Speech Synthesis
Python
2,133
star
29

MinkowskiEngine

Minkowski Engine is an auto-diff neural network library for high-dimensional sparse tensors
Python
2,007
star
30

TransformerEngine

A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper and Ada GPUs, to provide better performance with lower memory utilization in both training and inference.
Python
1,917
star
31

Stable-Diffusion-WebUI-TensorRT

TensorRT Extension for Stable Diffusion Web UI
Python
1,886
star
32

semantic-segmentation

Nvidia Semantic Segmentation monorepo
Python
1,763
star
33

gpu-operator

NVIDIA GPU Operator creates/configures/manages GPUs atop Kubernetes
Go
1,735
star
34

cub

[ARCHIVED] Cooperative primitives for CUDA C++. See https://github.com/NVIDIA/cccl
Cuda
1,679
star
35

DeepRecommender

Deep learning for recommender systems
Python
1,662
star
36

stdexec

`std::execution`, the proposed C++ framework for asynchronous and parallel programming.
C++
1,554
star
37

OpenSeq2Seq

Toolkit for efficient experimentation with Speech Recognition, Text2Speech and NLP
Python
1,511
star
38

CUDALibrarySamples

CUDA Library Samples
Cuda
1,468
star
39

VideoProcessingFramework

Set of Python bindings to C++ libraries which provides full HW acceleration for video decoding, encoding and GPU-accelerated color space and pixel format conversions
C++
1,303
star
40

deepops

Tools for building GPU clusters
Shell
1,252
star
41

open-gpu-doc

Documentation of NVIDIA chip/hardware interfaces
C
1,243
star
42

aistore

AIStore: scalable storage for AI applications
Go
1,233
star
43

Q2RTX

NVIDIA’s implementation of RTX ray-tracing in Quake II
C
1,217
star
44

trt-samples-for-hackathon-cn

Simple samples for TensorRT programming
Python
1,211
star
45

cccl

CUDA Core Compute Libraries
C++
1,200
star
46

MatX

An efficient C++17 GPU numerical computing library with Python-like syntax
C++
1,187
star
47

partialconv

A New Padding Scheme: Partial Convolution based Padding
Python
1,145
star
48

sentiment-discovery

Unsupervised Language Modeling at scale for robust sentiment classification
Python
1,055
star
49

nvidia-container-runtime

NVIDIA container runtime
Makefile
1,035
star
50

modulus

Open-source deep-learning framework for building, training, and fine-tuning deep learning models using state-of-the-art Physics-ML methods
Python
991
star
51

gpu-monitoring-tools

Tools for monitoring NVIDIA GPUs on Linux
C
974
star
52

jetson-gpio

A Python library that enables the use of Jetson's GPIOs
Python
898
star
53

dcgm-exporter

NVIDIA GPU metrics exporter for Prometheus leveraging DCGM
Go
886
star
54

retinanet-examples

Fast and accurate object detection with end-to-end GPU optimization
Python
885
star
55

flowtron

Flowtron is an auto-regressive flow-based generative network for text to speech synthesis with control over speech variation and style transfer
Jupyter Notebook
867
star
56

nccl-tests

NCCL Tests
Cuda
864
star
57

cuda-python

CUDA Python Low-level Bindings
Python
859
star
58

mellotron

Mellotron: a multispeaker voice synthesis model based on Tacotron 2 GST that can make a voice emote and sing without emotive or singing training data
Jupyter Notebook
852
star
59

gdrcopy

A fast GPU memory copy library based on NVIDIA GPUDirect RDMA technology
C++
832
star
60

libnvidia-container

NVIDIA container runtime library
C
818
star
61

BigVGAN

Official PyTorch implementation of BigVGAN (ICLR 2023)
Python
806
star
62

spark-rapids

Spark RAPIDS plugin - accelerate Apache Spark with GPUs
Scala
800
star
63

nv-wavenet

Reference implementation of real-time autoregressive wavenet inference
Cuda
728
star
64

DLSS

NVIDIA DLSS is a new and improved deep learning neural network that boosts frame rates and generates beautiful, sharp images for your games
C
727
star
65

tensorflow

An Open Source Machine Learning Framework for Everyone
C++
719
star
66

gvdb-voxels

Sparse volume compute and rendering on NVIDIA GPUs
C
674
star
67

MAXINE-AR-SDK

NVIDIA AR SDK - API headers and sample applications
C
671
star
68

nvvl

A library that uses hardware acceleration to load sequences of video frames to facilitate machine learning training
C++
665
star
69

runx

Deep Learning Experiment Management
Python
633
star
70

NVFlare

NVIDIA Federated Learning Application Runtime Environment
Python
630
star
71

NeMo-Aligner

Scalable toolkit for efficient model alignment
Python
564
star
72

nvcomp

Repository for nvCOMP docs and examples. nvCOMP is a library for fast lossless compression/decompression on the GPU that can be downloaded from https://developer.nvidia.com/nvcomp.
C++
545
star
73

multi-gpu-programming-models

Examples demonstrating available options to program multiple GPUs in a single node or a cluster
Cuda
535
star
74

Dataset_Synthesizer

NVIDIA Deep learning Dataset Synthesizer (NDDS)
C++
530
star
75

TensorRT-Model-Optimizer

TensorRT Model Optimizer is a unified library of state-of-the-art model optimization techniques such as quantization, pruning, distillation, etc. It compresses deep learning models for downstream deployment frameworks like TensorRT-LLM or TensorRT to optimize inference speed on NVIDIA GPUs.
Python
513
star
76

jitify

A single-header C++ library for simplifying the use of CUDA Runtime Compilation (NVRTC).
C++
512
star
77

nvbench

CUDA Kernel Benchmarking Library
Cuda
501
star
78

libglvnd

The GL Vendor-Neutral Dispatch library
C
501
star
79

NeMo-Curator

Scalable data pre processing and curation toolkit for LLMs
Jupyter Notebook
500
star
80

cuda-quantum

C++ and Python support for the CUDA Quantum programming model for heterogeneous quantum-classical workflows
C++
496
star
81

AMGX

Distributed multigrid linear solver library on GPU
Cuda
474
star
82

cuCollections

C++
470
star
83

enroot

A simple yet powerful tool to turn traditional container/OS images into unprivileged sandboxes.
Shell
459
star
84

NeMo-Framework-Launcher

Provides end-to-end model development pipelines for LLMs and Multimodal models that can be launched on-prem or cloud-native.
Python
459
star
85

hpc-container-maker

HPC Container Maker
Python
442
star
86

MDL-SDK

NVIDIA Material Definition Language SDK
C++
438
star
87

PyProf

A GPU performance profiling tool for PyTorch models
Python
437
star
88

framework-reproducibility

Providing reproducibility in deep learning frameworks
Python
424
star
89

gpu-rest-engine

A REST API for Caffe using Docker and Go
C++
421
star
90

DCGM

NVIDIA Data Center GPU Manager (DCGM) is a project for gathering telemetry and measuring the health of NVIDIA GPUs
C++
394
star
91

NvPipe

NVIDIA-accelerated zero latency video compression library for interactive remoting applications
Cuda
390
star
92

torch-harmonics

Differentiable signal processing on the sphere for PyTorch
Jupyter Notebook
386
star
93

cuQuantum

Home for cuQuantum Python & NVIDIA cuQuantum SDK C++ samples
Jupyter Notebook
344
star
94

data-science-stack

NVIDIA Data Science stack tools
Shell
317
star
95

ai-assisted-annotation-client

Client side integration example source code and libraries for AI-Assisted Annotation SDK
C++
308
star
96

video-sdk-samples

Samples demonstrating how to use various APIs of NVIDIA Video Codec SDK
C++
301
star
97

egl-wayland

The EGLStream-based Wayland external platform
C
299
star
98

nvidia-settings

NVIDIA driver control panel
C
292
star
99

NVTX

The NVIDIA® Tools Extension SDK (NVTX) is a C-based Application Programming Interface (API) for annotating events, code ranges, and resources in your applications.
C
290
star
100

gpu-feature-discovery

GPU plugin to the node feature discovery for Kubernetes
Go
286
star