重要提示

您正在查看 NeMo 2.0 文档。此版本引入了 API 的重大更改和一个新的库 NeMo Run。我们目前正在将 NeMo 1.0 中的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

Dask 集群函数#

nemo_curator.get_client(
cluster_type='cpu',
scheduler_address=None,
scheduler_file=None,
n_workers=4,
threads_per_worker=1,
nvlink_only=False,
protocol='tcp',
rmm_pool_size='1024M',
enable_spilling=True,
set_torch_to_use_rmm=False,
rmm_async=True,
rmm_maximum_pool_size=None,
rmm_managed_memory=False,
rmm_release_threshold=None,
**cluster_kwargs,
) → dask.distributed.Client#

初始化或连接到 Dask 集群。Dask 集群可以是基于 CPU 的或基于 GPU 的（如果 GPU 可用）。初始化确保 GPU 的最大内存效率，方法是

  1. 确保 PyTorch 内存池与 RAPIDS 内存池相同。（如果 set_torch_to_use_rmm 为 True）

  2. 为 cuDF 启用溢出。（如果 enable_spilling 为 True）

参数:

  • cluster_type – 如果 scheduler_address 和 scheduler_file 均为 None，则设置指定类型的本地（单节点）集群。可以是“cpu”或“gpu”。默认为“cpu”。get_client 中的许多选项仅适用于基于 CPU 或基于 GPU 的集群。请务必查看参数的描述。

  • scheduler_address – 要连接的现有 Dask 集群的地址。可以是调度器服务器的地址，如字符串“127.0.0.1:8786”或集群对象，如 LocalCluster()。如果指定，则忽略所有其他参数，并且客户端连接到现有集群。其他配置选项应在设置 Dask 集群时完成。

  • scheduler_file – 包含调度器信息的文件路径（如果可用）。如果指定，则忽略所有其他参数，并且客户端连接到现有集群。其他配置选项应在设置 Dask 集群时完成。

  • n_workers – 仅适用于基于 CPU 的集群。要启动的工作程序数。默认为 os.cpu_count()。对于基于 GPU 的集群，工作程序数锁定为 CUDA_VISIBLE_DEVICES 中的 GPU 数量。

  • threads_per_worker – 仅适用于基于 CPU 的集群。每个工作程序的线程数。默认为 1。在增加之前，请确保您的函数经常释放 GIL。

  • nvlink_only – 仅适用于基于 GPU 的集群。是否仅使用 nvlink。

  • protocol – 仅适用于基于 GPU 的集群。用于通信的协议。“tcp”或“ucx”。

  • rmm_pool_size – 仅适用于基于 GPU 的集群。RMM 池大小，用于初始化每个工作程序。可以是整数（字节）、浮点数（占设备总内存的百分比）、字符串（如“5GB”或“5000M”）或 None 以禁用 RMM 池。

  • enable_spilling – 仅适用于基于 GPU 的集群。启用从设备到主机的缓冲区的自动溢出（和“取消溢出”），以实现内存不足计算，即，对占用内存超过 GPU 可用内存的对象进行计算。

  • set_torch_to_use_rmm – 仅适用于基于 GPU 的集群。设置 PyTorch 内存池与 RAPIDS 内存池相同。这有助于避免在同一 GPU 上同时使用 PyTorch 和 RAPIDS 时出现 OOM 错误。

  • rmm_async – 仅适用于基于 GPU 的集群。使用 RAPIDS 内存管理器 (RMM) 初始化每个工作程序（有关更多信息，请参阅 RMM 文档：https://docs.rapids.ai/api/rmm/stable/）并将其设置为使用 RMM 的异步分配器。警告：异步分配器需要 CUDA Toolkit 11.2 或更高版本。它也与 RMM 池和托管内存不兼容。尝试同时启用两者将导致异常。

  • rmm_maximum_pool_size – 仅适用于基于 GPU 的集群。当设置了 rmm_pool_size 时，此参数指示最大池大小。可以是整数（字节）、浮点数（占设备总内存的百分比）、字符串（如“5GB”或“5000M”）或 None。默认情况下，使用 GPU 上的总可用内存。必须指定 rmm_pool_size 才能使用 RMM 池并设置最大池大小。注意：当与 –enable-rmm-async 配对使用时，由于碎片，无法保证最大大小。注意：此大小是每个工作程序的配置，而不是集群范围的配置。

  • rmm_managed_memory – 仅适用于基于 GPU 的集群。使用 RMM 初始化每个工作程序，并将其设置为使用托管内存。如果禁用，则仍可以通过指定 rmm_pool_size 来使用 RMM。警告：托管内存目前与 NVLink 不兼容。尝试同时启用两者将导致异常。

  • rmm_release_threshold – 仅适用于基于 GPU 的集群。当 rmm.async 为 True 并且池大小增长超过此值时，池持有的未使用内存将在下一个同步点释放。可以是整数（字节）、浮点数（占设备总内存的百分比）、字符串（如“5GB”或“5000M”）或 None。默认情况下，此功能已禁用。注意：此大小是每个工作程序的配置，而不是集群范围的配置。

  • cluster_kwargs – LocalCluster 或 LocalCUDACluster 配置的其他关键字参数。有关所有 LocalCluster 参数，请参阅 API 文档 https://docs.dask.org/en/stable/deploying-python.html#distributed.deploy.local.LocalCluster，或有关所有 LocalCUDACluster 参数，请参阅 https://docs.rapids.ai/api/dask-cuda/nightly/api/。

返回:

Dask 客户端对象。

nemo_curator.get_network_interfaces() → List[str]#

获取机器上所有有效网络接口的列表

返回:

机器上所有有效网络接口的列表

class nemo_curator.ToBackend(backend: Literal['pandas', 'cudf'])#

用于在后端之间移动数据帧的模块。

call(
dataset: DocumentDataset,
) → DocumentDataset#

对数据集执行任意操作

参数:

dataset (DocumentDataset) – 要操作的数据集

在本页中