【Python】获取 GPU 状态 及NV-Embed-v1大模型使用

GPUtil是一个 Python 模块，用于使用 从 NVIDA GPU 获取 GPU 状态。

nvidia-smi.GPUtil查找计算机上的所有 GPU，确定其可用性并返回可用 GPU 的有序列表。可用性取决于每个 GPU 的当前内存消耗和负载。该模块在编写时考虑了深度学习的 GPU 选择，但它不是特定于任务/库的，它可以应用于任何任务，在这些任务中，识别可用的 GPU 可能很有用。

此工具利用nvidia-smi来实时抓取gputil使用情况，提供丰富的GPU实时信息，包括但不限于：

gpu型号  || GPU temp. | GPU util. | Memory util. || Memory total | Memory used | Memory free

要求

安装了最新 NVIDIA 驱动程序的 NVIDIA GPU。GPUtil 使用nvidia-smi程序获取所有可用 NVIDIA GPU 的 GPU 状态。 nvidia-smi在安装 NVIDIA 驱动程序时自动安装。

支持 Python 2.X 和 3.X。

Python 库：

• subprocess （Python 标准库)
• distutils（Python 标准库）)
• math （Python 标准库)
• random（Python 标准库)
• time（Python 标准库)
• os（Python 标准库)
• sys（Python 标准库）)
• platform（Python 标准库)

在 CUDA 驱动程序版本 390.77、Python 2.7 和 3.5 上进行了测试。

安装

1. 打开终端 （Ctrl+Shift+T）
2. 类型pip install gputil
3. 测试安装

1. 在 GPUtil 文件夹以外的文件夹中打开终端
2. 通过在终端中键入来启动 python 控制台
3. 在新打开的 python 控制台中，键入：

import GPUtil
GPUtil.showUtilization()

1. 输出应如下所示，具体取决于 GPU 数量及其当前使用情况：

ID  GPU  MEM
--------------
  0    0%   0%

旧安装方式

1. 将存储库下载或克隆到您的计算机
2. 将 GPUtil 文件夹添加到 ~/.bashrc

1. 打开新终端 （按 Ctrl+Alt+T）
2. 打开 bashrc：

gedit ~/.bashrc

1. 已将 GPUtil 文件夹添加到环境变量PYTHONPATH（替换<path_to_gputil>为文件夹路径）：

export PYTHONPATH="$PYTHONPATH:<path_to_gputil>"

Example:
export PYTHONPATH="$PYTHONPATH:/home/anderskm/github/gputil"

1. 保存 ~/.bashrc 并关闭 gedit
2. 重新启动终端

1. 测试安装

1. 在 GPUtil 文件夹以外的文件夹中打开终端
2. 通过在终端中键入来启动 python 控制台
3. 在新打开的 python 控制台中，键入：

import GPUtil
GPUtil.showUtilization()

1. 输出应如下所示，具体取决于 GPU 数量及其当前使用情况：

ID  GPU  MEM
--------------
  0    0%   0%

用法

您所要做的就是将GPUtil包含在脚本的开头：

import GPUtil

一旦包含所有功能，所有功能都可用。这些函数以及输入、输出及其功能的简短描述可以在以下两节中找到。

主要功能

deviceIDs = GPUtil.getAvailable(order = 'first', limit = 1, maxLoad = 0.5, maxMemory = 0.5, includeNan=False, excludeID=[], excludeUUID=[])

返回可用 GPU 的列表 ID。可用性是根据当前内存使用情况和负载确定的。顺序、最大设备数量、最大负载和最大内存消耗由输入参数决定。

• 输入

• order- 确定返回可用 GPU 设备 ID 的顺序。order 应指定为以下字符串之一：

• 'first'- 按升序 id （defaut)
• 'last'- 按 id 降序排序可用的 GPU 设备 ID
• 'random'- 随机订购可用的 GPU 设备 ID
• 'load'- 按升序加载对可用的 GPU 设备 ID 进行排序
• 'memory'- 按内存使用率递增顺序对可用的 GPU 设备 ID 进行排序

• limit- 将返回的 GPU 设备 ID 数量限制为指定数量。必须为正整数。（默认值 = 1)
• maxLoad- GPU 的最大当前相对负载被视为可用。负载maxLoad大于 GPU 的 GPU 不会返回。（默认值 = 0.5)
• maxMemory- GPU 的最大当前相对内存使用量被视为可用。当前内存使用量大于大于 GPU 的 GPU 不会返回。（默认值 = 0.5maxMemory)
• includeNan- True/false 标志，指示是否包含负载或内存使用率为 NaN 的 GPU（指示无法检索使用情况）。（默认值 = False)
• excludeID- ID 列表，应从可用 GPU 列表中排除。（默认值 = [])
• excludeUUID- 与excludeID相同，只是它使用 UUID。（默认值 = [])

• 输出

• deviceID - 所有可用 GPU 设备 ID 的列表。如果当前负载和内存使用率分别小于maxLoad和maxMemory，则认为 GPU 可用。该列表根据order排序。返回的设备 ID 的最大数量受limit的限制。

deviceID = GPUtil.getFirstAvailable(order = 'first', maxLoad=0.5, maxMemory=0.5, attempts=1, interval=900, verbose=False)

返回第一个可用的 GPU。可用性根据当前内存使用情况和负载确定，顺序由指定的顺序确定。如果未找到可用的 GPU，则会引发错误。使用默认值时，getAvailable(order = 'first', limit = 1, maxLoad = 0.5, maxMemory = 0.5)

• 输入

• order- 请参阅说明GPUtil.getAvailable(...)
• maxLoad- GPU 的最大当前相对负载被视为可用。负载大于maxLoadGPU 的 GPU 不会返回。（默认值 = 0.5)
• maxMemory- GPU 的最大当前相对内存使用量被视为可用。当前内存使用量大于大于 GPU 的 GPU 不会返回。（默认值 = 0.5maxMemory)
• attempts- 在放弃寻找可用 GPU 之前，函数应进行的尝试次数。（默认值 = 1)
• interval- 每次尝试查找可用 GPU 之间的间隔（以秒为单位）。（默认值 = 900 --> 15 分钟）
• verbose- 如果True ，则在每次尝试之前打印尝试次数，如果找到可用次数，则打印 GPU ID。
• includeNan- 请参阅GPUtil.getAvailable(...)的说明。（默认值 = False)
• excludeID- 请参阅GPUtil.getAvailable(...)的说明。（默认值 = [])
• excludeUUID- 请参阅 GPUtil.getAvailable(...)的说明。（默认值 = [])

• 输出

• deviceID - 包含 1 个元素的列表，其中包含第一个可用的 GPU 设备 ID。如果当前负载和内存使用率分别小于maxLoad和maxMemory，则认为 GPU 可用。阶数和限制分别固定为'first'和1
• 。

GPUtil.showUtilization(all=False, attrList=None, useOldCode=False)

打印所有 GPU 的当前状态（id、内存使用情况、uuid 负载）

• 输入

• all- 真/假标志，指示是否应显示 GPU 上的所有信息。覆盖attrList。
• attrList- 要显示的GPU属性列表列表。
• useOldCode- True/false 标志，指示是否应使用旧代码来显示 GPU 利用率。

• 输出

• 没有

例子

在 Caffe 中选择第一个可用的 GPU

在深度学习库 Caffe 中，用户可以通过其 Python 接口在使用 CPU 或 GPU 之间切换。这是通过分别调用方法caffe.set_mode_cpu()和caffe.set_mode_gpu()来完成的。下面是选择第一个带有 GPUtil 的可用 GPU 来运行 Caffe 网络的最小工作示例。

# Import caffe and GPUtil
import caffe
import GPUtil

# Set CUDA_DEVICE_ORDER so the IDs assigned by CUDA match those from nvidia-smi
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

# Get the first available GPU
DEVICE_ID_LIST = GPUtil.getFirstAvailable()
DEVICE_ID = DEVICE_ID_LIST[0] # grab first element from list

# Select GPU mode
caffe.set_mode_gpu()
# Select GPU id
caffe.set_device(DEVICE_ID)

# Initialize your network here

注意：在编写此示例时，Caffe Python 包装器仅支持 1 个 GPU，尽管底层代码支持多个 GPU。从终端直接调用 Caffe 允许使用多个 GPU。

在 TensorFlow 中仅占用 1 个 GPU

默认情况下，当使用 gpu 作为设备时，TensorFlow 将占用所有可用的 GPU (例如 tf.device (‘ gpu: 0’))。通过设置环境变量 CUDA _ VISIBLE _ devICES，用户可以通过 CUDA 屏蔽 TensorFlow 应该可见的图形处理器(见 CUDA _ VISIBLE _ devICES-屏蔽图形处理器)。使用 GPUtil.py，可以基于可用的 GPU 以编程方式设置 CUDA _ VISIBLE _ DevICES。下面是如何使用 GPUtil 在 TensorFlow 中只占用1个 GPU 的最小工作示例。要运行代码，将其复制到一个新的 python 文件(例如 demo _ tensorflow _ gputil.py)并运行它(例如在终端中输入 python demo _ tensorflow _ gputil.py)。

注意: 即使您将运行代码的设备设置为 CPU，TensorFlow 也会占用所有可用的 GPU。为了避免这种情况，可以使用 os.environ [“ CUDA _ VISIBLE _ DevICES”] =”来隐藏所有 GPU。

# Import os to set the environment variable CUDA_VISIBLE_DEVICES
import os
import tensorflow as tf
import GPUtil

# Set CUDA_DEVICE_ORDER so the IDs assigned by CUDA match those from nvidia-smi
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

# Get the first available GPU
DEVICE_ID_LIST = GPUtil.getFirstAvailable()
DEVICE_ID = DEVICE_ID_LIST[0] # grab first element from list

# Set CUDA_VISIBLE_DEVICES to mask out all other GPUs than the first available device id
os.environ["CUDA_VISIBLE_DEVICES"] = str(DEVICE_ID)

# Since all other GPUs are masked out, the first available GPU will now be identified as GPU:0
device = '/gpu:0'
print('Device ID (unmasked): ' + str(DEVICE_ID))
print('Device ID (masked): ' + str(0))

# Run a minimum working example on the selected GPU
# Start a session
with tf.Session() as sess:
    # Select the device
    with tf.device(device):
        # Declare two numbers and add them together in TensorFlow
        a = tf.constant(12)
        b = tf.constant(30)
        result = sess.run(a+b)
        print('a+b=' + str(result))

输出应类似于下面的代码块。请注意，如何仅找到一个 GPU 并将其创建为 tensorflow 设备。

I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcublas.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcudnn.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcufft.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcuda.so.1 locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcurand.so locally
Device: /gpu:0
I tensorflow/core/common_runtime/gpu/gpu_device.cc:885] Found device 0 with properties: 
name: TITAN X (Pascal)
major: 6 minor: 1 memoryClockRate (GHz) 1.531
pciBusID 0000:02:00.0
Total memory: 11.90GiB
Free memory: 11.76GiB
I tensorflow/core/common_runtime/gpu/gpu_device.cc:906] DMA: 0 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:916] 0:   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:975] Creating TensorFlow device (/gpu:0) -> (device: 0, name: TITAN X (Pascal), pci bus id: 0000:02:00.0)
a+b=42

在单独的线程中监视 GPU

如果在训练期间使用 GPUtil 监控 GPU，则可能会显示 0% 的利用率。解决此问题的方法是使用单独的监视线程。

import GPUtil
from threading import Thread
import time

class Monitor(Thread):
    def __init__(self, delay):
        super(Monitor, self).__init__()
        self.stopped = False
        self.delay = delay # Time between calls to GPUtil
        self.start()

    def run(self):
        while not self.stopped:
            GPUtil.showUtilization()
            time.sleep(self.delay)

    def stop(self):
        self.stopped = True
        
# Instantiate monitor with a 10-second delay between updates
monitor = Monitor(10)

# Train, etc.

# Close monitor
monitor.stop()

什么是nvidia/NV-Embed-v1？

我们介绍了 NV-Embed，这是一个通用嵌入模型，在海量文本嵌入基准（MTEB 基准）中排名第一（截至 2024 年 5 月 24 日），有 56 个任务，包括检索、重新排序、分类、聚类和语义文本相似性任务。值得注意的是，我们的模型在该基准测试中的 15 个检索任务中也获得了 59.36 分的最高分。

NV-Embed提出了几种新的设计，包括让LLM关注潜在向量以获得更好的池化嵌入输出，并演示了一种两阶段指令调优方法，以提高检索和非检索任务的准确性。

型号详细信息

• 基本解码器 LLM：Mistral-7B-v0.1
• 嵌入维数：4096
• Vocab size：32k
• 最大输入Token长度：4096
• 模型参数：7.85B

如何使用

下面是如何对查询和段落进行编码的示例。

import torch
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel

# Each query needs to be accompanied by an corresponding instruction describing the task.
task_name_to_instruct = {"example": "Given a question, retrieve passages that answer the question",}

query_prefix = "Instruct: "+task_name_to_instruct["example"]+"\nQuery: "
queries = [
    'are judo throws allowed in wrestling?', 
    'how to become a radiology technician in michigan?'
    ]

# No instruction needed for retrieval passages
passage_prefix = ""
passages = [
    "Since you're reading this, you are probably someone from a judo background or someone who is just wondering how judo techniques can be applied under wrestling rules. So without further ado, let's get to the question. Are Judo throws allowed in wrestling? Yes, judo throws are allowed in freestyle and folkstyle wrestling. You only need to be careful to follow the slam rules when executing judo throws. In wrestling, a slam is lifting and returning an opponent to the mat with unnecessary force.",
    "Below are the basic steps to becoming a radiologic technologist in Michigan:Earn a high school diploma. As with most careers in health care, a high school education is the first step to finding entry-level employment. Taking classes in math and science, such as anatomy, biology, chemistry, physiology, and physics, can help prepare students for their college studies and future careers.Earn an associate degree. Entry-level radiologic positions typically require at least an Associate of Applied Science. Before enrolling in one of these degree programs, students should make sure it has been properly accredited by the Joint Review Committee on Education in Radiologic Technology (JRCERT).Get licensed or certified in the state of Michigan."
]

# load model with tokenizer
model = AutoModel.from_pretrained('nvidia/NV-Embed-v1', trust_remote_code=True)

# get the embeddings
max_length = 4096
query_embeddings = model.encode(queries, instruction=query_prefix, max_length=max_length)
passage_embeddings = model.encode(passages, instruction=passage_prefix, max_length=max_length)

# normalize embeddings
query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
passage_embeddings = F.normalize(passage_embeddings, p=2, dim=1)

# get the embeddings with DataLoader (spliting the datasets into multiple mini-batches)
# batch_size=2
# query_embeddings = model._do_encode(queries, batch_size=batch_size, instruction=query_prefix, max_length=max_length)
# passage_embeddings = model._do_encode(passages, batch_size=batch_size, instruction=passage_prefix, max_length=max_length)

scores = (query_embeddings @ passage_embeddings.T) * 100
print(scores.tolist())
#[[77.9402084350586, 0.4248958230018616], [3.757718086242676, 79.60113525390625]]

如何启用多 GPU

from transformers import AutoModel
from torch.nn import DataParallel

embedding_model = AutoModel.from_pretrained("nvidia/NV-Embed-v1")
for module_key, module in embedding_model._modules.items():
    embedding_model._modules[module_key] = DataParallel(module)

所需软件包

如果您遇到问题，请尝试安装 python 包，如下所示

pip uninstall -y transformer-engine
pip install torch==2.2.0
pip install transformers --upgrade
pip install flash-attn==2.2.0