            
            Given:

N = 15
indexed_indices = [5, 6, 7, 8, 9]
E = 3

Then, for example:

Do_idx = [8, 8, 9, 6, 7, 5, 8, 5, 6, 6, 5, 9, 7, 7, 9]

样本索引 Sa_idx 是一个 2-D 数组，从 j 映射到形状为 [N + 1, 2] 的 (i, Do_idx[ i ] 偏移量) 对。行 j 和 j + 1 用作第 j 个样本的左右边界。

复制
已复制！

            
            Given:

S = 1024

Then, for example:

Sa_idx[0] = (0, 0)
Sa_idx[1] = (0, 1024)       => Do_idx[0] has length greater than S
Sa_idx[2] = (1, 512)        => Do_idx[0] has length 1536
Sa_idx[3] = (2, 0)          => Do_idx[1] has length 1536
Sa_idx[4] = (5, 300)        => Do_idx[2:5] are shorter documents relative to Do_idx[0:2]
Sa_idx[5] = (6, 24)         => Do_idx[5] has length 1300

shuffle 索引 Sh_idx 是一个 1-D 数组，从 k 映射到长度为 N 的 j。shuffle 索引根据 R 进行 shuffle。

复制
已复制！

            
            Given

N = 10

Then, for example:

Sh_idx = [4, 0, 2, 6, 1, 9, 5, 8, 7, 3]

要查询 GPTDataset 的第 k 个样本，我们执行以下操作

使用 shuffle 索引获取样本索引的索引 j。

复制
已复制！

            
            j = Sh_idx[k]

使用样本索引获取文档索引的左右样本边界索引以及每个文档的起始 token 偏移量。

复制
已复制！

            
            i, offset = Sa_idx[j]
i_next, offset_next = Sa_idx[j + 1]

使用文档索引从连续（在文档索引中）文档中检索 S 个 token。

复制
已复制！

            
            sample = []
sample += indexed_dataset[Do_idx[i]][offset:]
if i != i_next:
    sample += indexed_dataset[Do_idx[i + 1:i_next]]
sample += indexed_dataset[Do_idx[i_next]][:offset_next]

为了节省初始化期间的时间，每个索引都在一个进程 rank 上按顺序构建/缓存，然后在其他进程 rank 上并行加载。缓存的索引对于在 MegatronDataset.__init__ 函数中生成的哈希是唯一的。

BlendedDataset

BlendedDataset 由以下变量参数化：底层 MegatronDataset 实例 D、权重 W（每个数据集一个）和大小 S。BlendedDataset 将按权重比例从贡献数据集中抽取样本，直到达到所需大小的复合数据集。在每个采样步骤中，我们从具有最大采样误差的数据集中抽取单个样本。

BlendedDataset 创建两个“混合”索引以方便查找：（1）数据集索引和（2）数据集样本索引。

数据集索引 Da_idx 是一个 1-D 数组，从 i 映射到长度为 S 的数据集索引。

复制
已复制！

            
            Given

D = [d0, d1, d2]
W = [1/2, 1/4, 1/4]
S = 4

Then, for example:

Da_idx = [0, 1, 2, 0]

数据集样本索引 Sa_idx 是一个 1-D 映射，从 i 映射到长度为 S 的数据集 Da_idx[i] 的样本索引。

复制
已复制！

            
            Given

Da_idx = [0, 1, 2, 0]

Then, for example:

Sa_idx = [0, 0, 0, 1]

要查询 BlendedDataset 的第 k 个样本，我们执行以下操作

使用数据集索引从 D 中检索相应的数据集，并使用数据集样本索引从该数据集中检索相应的样本。
复制

已复制！
```
            
            sample = D[Da_idx[k]][Sa_idx[k]]
        
```

为了节省初始化期间的时间，每个索引都在一个进程 rank 上按顺序构建/缓存，然后在其他进程 rank 上并行加载。缓存的索引对于在 BlendedDataset.__init__ 函数中生成的哈希是唯一的。

datasets.blended_megatron_dataset_config 模块

class core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig(random_seed: int, sequence_length: int, blend: Optional[Tuple[List[str], Optional[List[float]]]] = None, blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]] = None, renormalize_blend_weights: bool = False, split: Optional[str] = None, num_dataset_builder_threads: int = 1, path_to_cache: Optional[str] = None, mmap_bin_files: bool = True, tokenizer: Optional[megatron.core.datasets.megatron_tokenizer.MegatronTokenizer] = None)

基类：object

Megatron Core 数据集的配置对象

blend: Optional[Tuple[List[str], Optional[List[float]]]] = None: 混合，由数据集前缀列表和可选的数据集权重列表组成。例如，[[“dataset-path1”, “dataset-path2”], [0.3, 0.7]]。当权重为 None 时，它们从贡献数据集的长度推断而来。不与 ‘blend_per_split’ 一起使用。默认为 None。

blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]] = None: 一组混合，如上定义，每个拆分分布一个。不与 ‘blend’ 一起使用。默认为 None。

mmap_bin_files: bool = True: 是否 mmap .bin 文件或使用文件指针。

mock: bool = False: 是否绕过真实数据加载和验证，而支持模拟数据生成。从 ‘blend’ 和 ‘blend_per_split’ 自动创建。不传递给构造函数。

num_dataset_builder_threads: int = 1: 用于数据集构建的线程数。

path_to_cache: Optional[str] = None: 所有可重用数据集索引要缓存到的位置。

random_seed: int: 数据集创建期间所有 RNG 的种子。

renormalize_blend_weights: bool = False: 重新归一化混合权重，以考虑为确保满足请求的样本数而完成的中级数据集过采样。为了向后兼容数据样本顺序，默认为 False。

sequence_length: int: 序列长度。

split: Optional[str] = None: 拆分字符串，从单个分布中抽取样本时，数据集拆分的逗号分隔权重。不与 ‘blend_per_split’ 一起使用。默认为 None。

split_matrix: Optional[List[Tuple[float, float]]] = None: 拆分矩阵，由每个拆分的非重叠 book-end 按顺序组成。有关更多信息，请参阅 ‘convert_split_vector_to_split_matrix’。从 ‘split’ 自动创建。不传递给构造函数。

tokenizer: Optional[megatron.core.datasets.megatron_tokenizer.MegatronTokenizer] = None: MegatronTokenizer 实例或 None。对于执行在线 tokenization 的数据集是必需的。

core.datasets.blended_megatron_dataset_config.convert_split_vector_to_split_matrix(vector_a: List[float], vector_b: Optional[List[float]] = None) → List[Optional[Tuple[float, float]]]

从一个或可选的两个贡献拆分向量构建拆分矩阵。

例如，标准转换

[0.99, 0.01, 0.0] -> [(0, 0.99), (0.99, 1.0), None]

例如，当 Retro 预训练使用 [0.99, 0.01, 0.0] 拆分，而 Retro 预处理使用 [0.98, 0.02, 0.0] 拆分时，Retro 的转换

[0.99, 0.01, 0.0], [0.98, 0.02, 0.0] -> [(0, 0.98), (0.99, 1.0), None]

参数

vector_a (List[float]) – 主要拆分向量
vector_b (Optional[List[float]]) – 可选的辅助拆分向量，用于约束主要拆分向量。默认为 None。

返回值

由每个拆分的 book-end 按顺序组成的拆分矩阵

返回类型

List[Tuple[float, float]]

core.datasets.blended_megatron_dataset_config.parse_and_normalize_split(split: str) → List[float]

从字符串解析数据集拆分比率

参数: split (str) – 训练验证测试拆分字符串，例如 “99,1,0”
返回值: 训练验证测试拆分比率，例如 [0.99, 0.01, 0.0]
返回类型: List[float]

datasets.blended_megatron_dataset_builder 模块

class core.datasets.blended_megatron_dataset_builder.BlendedMegatronDatasetBuilder(cls: Type[megatron.core.datasets.megatron_dataset.MegatronDataset], sizes: List[int], is_built_on_rank: Callable, config: megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig)

基类：object

BlendedDataset 和 MegatronDataset 类的构建器类

参数

cls (Type[MegatronDataset]) – 要实例化的类，必须继承自 MegatronDataset
sizes (List[Optional[int]]) – 每个拆分要抽取的最小总样本数，或 None
is_built_on_rank (Callable) – 一个可调用对象，如果数据集应在当前 rank 上构建，则返回 True，否则返回 False。它应该是 Megatron Core 并行感知的，即全局 rank、本地组 rank 和虚拟 rank 可以告知其返回值。
config (BlendedMegatronDatasetConfig) – 通知数据集创建的配置对象

build() → List[Optional[Union[megatron.core.datasets.blended_dataset.BlendedDataset, megatron.core.datasets.megatron_dataset.MegatronDataset]]]

根据提供的混合构建所有数据集拆分

此方法是分布式感知的，必须在所有 rank 上调用。

返回的数据集拆分可能因配置而异。提供 config.blend 和 config.split 以从同一分布构建 BlendedDataset 和/或 MegatronDataset 拆分。提供 config.blend_per_split 以从单独的分布构建 BlendedDataset 和/或 MegatronDataset 拆分。在任何一种情况下，对于每个拆分，处理以下情况

拆分为 None - 不执行任何操作
拆分有一个贡献数据集，并且…
1. ‘size’ 不是 None - 构建中级数据集，其中低级数据集采样与大小成比例
2. ‘size’ 是 None - 构建中级数据集，没有过多的低级数据集采样
拆分有多个贡献数据集，并且…
1. ‘weights’ 不是 None 且 ‘size’ 不是 None - 构建中级数据集，其中低级数据集采样与其权重和大小成比例 - 构建一个长度略大于 ‘size’ 的顶级数据集，其中中级数据集采样与其权重和大小成比例
2. ‘weights’ 不是 None 且 ‘size’ 是 None - 错误
3. ‘weights’ 是 None 且 ‘size’ 不是 None - 构建中级数据集，没有过多的低级数据集采样 - 构建一个长度为 ‘size’ 的顶级数据集，其中中级数据集采样与其长度和大小成比例
  - 顶级数据集的 ‘size’ 上限为中级数据集长度的总和
4. ‘weights’ 是 None 且 ‘size’ 是 None - 构建中级数据集，没有过多的低级数据集采样 - 构建一个顶级数据集，没有过多的中级数据集采样

返回值: 包含每个拆分的数据集实例（或 None）的列表
返回类型: List[Optional[TopLevelDataset]]

static build_generic_dataset(cls: Union[Type[Union[megatron.core.datasets.blended_dataset.BlendedDataset, megatron.core.datasets.megatron_dataset.MegatronDataset, megatron.core.datasets.megatron_dataset.LowLevelDataset, torch.utils.data.Dataset]], Callable], is_built_on_rank: Callable, synchronize_ranks: bool, *args: Any) → Optional[Union[megatron.core.datasets.blended_dataset.BlendedDataset, megatron.core.datasets.megatron_dataset.MegatronDataset, megatron.core.datasets.megatron_dataset.LowLevelDataset, torch.utils.data.Dataset, Iterable]]

构建 DistributedDataset

当且仅当底层数据集类未在当前 rank 上构建且 torch.distributed 已初始化时，才返回 None。

参数

cls (Union[Type[DistributedDataset], Callable]) – 要构建的 DistributedDataset 类。在特殊情况下，例如当我们为 RawMegatronDataset 实例构建低级数据集时，我们可以接受返回 Iterable 的 Callable。
synchronize_ranks (bool) – 是否为 rank-0 / barrier / 其他-ranks 行为调用 barrier。当我们强制在更高级别执行此行为时，设置为 False。
args (Tuple[Any]) – 用于构建提供的 DistributedDataset 类的位置参数

异常

Exception – 当数据集构造函数引发 OSError 时

返回值

DistributedDataset 实例化、Iterable 实例化或 None

返回类型

Optional[Union[DistributedDataset, Iterable]]

datasets.megatron_tokenizer 模块

class core.datasets.megatron_tokenizer.MegatronTokenizer(*tokenizer_paths: str, **tokenizer_options: Any)

基类：abc.ABC

tokenizer 的抽象类

在缺少配置或类特定的唯一标识对象跟踪的情况下，我们必须包含所有关键字参数作为唯一标识符

参数

tokenizer_paths (Tuple[str]) – 所有 tokenizer 源路径或前缀
tokenizer_options (Dict[str, Any]) – 所有 tokenizer 选项

property bos

BOS token id

异常: NotImplementedError – 非抽象、可选属性

property cls

CLS token id

异常: NotImplementedError – 非抽象、可选属性

detokenize(ids: numpy.ndarray) → str

将 embedding id 转换为文本

参数: ids (numpy.ndarray) – 要转换的 id
返回值: 转换后的文本
返回类型: str
异常: NotImplementedError – 非抽象、可选方法

property eod

EOD token id

异常: NotImplementedError – 非抽象、可选属性

property eos

EOS token id

异常: NotImplementedError – 非抽象、可选属性

abstract property inv_vocab: 从 vocab id token 到文本 token 的字典

property mask

MASK token id

异常: NotImplementedError – 非抽象、可选属性

offsets(ids: list[int], text: str) → list[int]

将 embedding id 转换为文本偏移量

参数

ids (list[int]) – 要转换的 id
text (str) – 要转换的文本

返回值

转换后的偏移量

返回类型

list[int]

异常

NotImplementedError – 非抽象、可选方法

property pad

PAD token id

异常: NotImplementedError – 非抽象、可选属性

property sep

SEP token id

异常: NotImplementedError – 非抽象、可选属性

abstract tokenize(text: str) → numpy.ndarray

将文本转换为 embedding id

参数: text (str) – 要转换的文本
返回值: 转换后的 embedding id
返回类型: numpy.ndarray

abstract property vocab: 从 vocab 文本 token 到 id token 的字典

abstract property vocab_size: 词汇表大小

datasets.indexed_dataset 模块

class core.datasets.indexed_dataset.DType(value)

基类：enum.Enum

用于写入/读取 IndexedDataset 索引的 NumPy 数据类型枚举

classmethod code_from_dtype(value: Type[numpy.number]) → int

从 dtype 获取代码

参数: value (Type[numpy.number]) – dtype
返回值: 代码
返回类型: int

classmethod dtype_from_code(value: int) → Type[numpy.number]

从代码获取 dtype

参数: value (int) – 代码
返回值: dtype
返回类型: Type[numpy.number]

float32 = 7

float64 = 6

int16 = 3

int32 = 4

int64 = 5

int8 = 2

static optimal_dtype(cardinality: Optional[int]) → Type[numpy.number]

获取用于特定基数的索引的 dtype

参数: cardinality (Optional[int]) – 要索引的元素数
返回值: 用于索引的 dtype
返回类型: Type[numpy.number]

static size(key: Union[int, Type[numpy.number]]) → int

以字节为单位获取 dtype/代码的大小

参数: key (Union[int, Type[numpy.number]]) – dtype 或代码
异常: ValueError – 如果 key 既不是 dtype 也不是整数代码
返回值: 以字节为单位的 dtype/代码的大小
返回类型: int

uint16 = 8

uint8 = 1

class core.datasets.indexed_dataset.IndexedDataset(*args: Any, **kwargs: Any)

基类：torch.utils.data.Dataset

低级接口数据集类

参数

path_prefix (str) – 索引 (.idx) 和数据 (.bin) 前缀
multimodal (bool) – 数据集是否为多模态。默认为 False。
mmap (bool) – 是否 mmap .bin 文件。默认为 True。
s3_config (Optional[S3Config]) – 仅为存储在 S3 上的数据提供。IndexedDataset 将索引 (.idx) 文件下载到 s3_config.path_to_idx_cache，并以 s3_config.bin_chunk_nbytes 块的形式从数据 (.bin) 文件流式传输数据。请注意，对于 S3 数据加载，必须禁用 mmap。默认为 None。

property document_indices: numpy.ndarray

获取文档索引

返回值: 文档索引
返回类型: numpy.ndarray

static exists(path_prefix: str) → bool

返回 IndexedDataset 是否存在于磁盘上的前缀处

参数: path_prefix (str) – 索引 (.idx) 和数据 (.bin) 文件的前缀
返回值: IndexedDataset 是否存在于磁盘上的前缀处
返回类型: bool

get(idx: int, offset: int = 0, length: Optional[int] = None) → numpy.ndarray

从数据集中检索单个项目，可以选择仅返回项目的一部分。

get(idx) 与 [idx] 相同，但 get() 不支持切片。

参数

idx (Union[int, numpy.integer]) – 数据集的索引
offset (int) – 序列中的整数 token 偏移量
length (int) – 要从序列中抓取的 token 数

返回值

索引处的序列 token 和模式

返回类型

Union[numpy.ndarray, Tuple[numpy.ndarray, numpy.ndarray]]

get_document_indices() → numpy.ndarray

获取文档索引

此方法计划弃用。

返回值: 文档索引
返回类型: numpy.ndarray

initialize(path_prefix: str, multimodal: bool, mmap: bool, s3_config: Optional[megatron.core.datasets.utils_s3.S3Config]) → None

初始化数据集

此方法由 IndexedDataset.__init__ 在对象创建期间调用，并由 IndexedDataset.__setstate__ 在取消 pickle 期间调用

参数

path_prefix (str) – 索引 (.idx) 和数据 (.bin) 前缀
multimodal (bool) – 数据集是否为多模态
mmap (bool) – 是否 mmap .bin 文件
s3_config (Optional[S3Config]) – 有关详细信息，请参阅 IndexedDataset 文档字符串。

property sequence_lengths: numpy.ndarray

获取序列长度

返回值: 序列长度
返回类型: numpy.ndarray

property sequence_modes: numpy.ndarray

获取序列模式

返回值: 序列模式
返回类型: numpy.ndarray

set_document_indices(document_indices: numpy.ndarray) → None

设置文档索引

此方法计划弃用。

参数: document_indices (numpy.ndarray) – 文档索引

class core.datasets.indexed_dataset.IndexedDatasetBuilder(bin_path: str, dtype: Type[numpy.number] = numpy.int32, multimodal: bool = False)

基类：object

IndexedDataset 类的构建器类

参数

bin_path (str) – 数据 (.bin) 文件的路径
dtype (Type[numpy.number], optional) – 索引文件的 dtype。默认为 numpy.int32。
multimodal (bool, optional) – 数据集是否为多模态。默认为 False。

add_document(tensor: torch.Tensor, lengths: List[int], modes: Optional[List[int]] = None) → None

将整个文档添加到数据集

参数

tensor (torch.Tensor) – 要添加的文档
lengths (List[int]) – 文档中每个项目的长度
modes (Optional[List[int]], optional) – 文档中每个项目的模式。默认为 None。

add_index(path_prefix: str) → None

将整个 IndexedDataset 添加到数据集

参数: path_prefix (str) – 索引 (.idx) 和数据 (.bin) 前缀

add_item(tensor: torch.Tensor, mode: int = 0) → None

向数据集添加单个项目

参数

tensor (torch.Tensor) – 要添加到数据文件的项目
mode (int, optional) – 项目的模式。默认为 0。

end_document() → None: 完成文档，用于 IndexedDatasetBuilder.add_item

finalize(idx_path: str) → None

清理并写入索引 (.idx) 文件

参数: idx_path (str) – 索引文件的路径

core.datasets.indexed_dataset.get_bin_path(path_prefix: str) → str

从前缀获取数据文件的路径

参数: path_prefix (str) – 前缀
返回值: 数据文件的路径
返回类型: str

core.datasets.indexed_dataset.get_idx_path(path_prefix: str) → str

从前缀获取索引文件的路径

参数: path_prefix (str) – 前缀
返回值: 索引文件的路径
返回类型: str

datasets.megatron_dataset 模块

class core.datasets.megatron_dataset.MegatronDataset(*args: Any, **kwargs: Any)

基类： abc.ABC, torch.utils.data.Dataset

所有数据集类应继承的最高级别包装器类

参数

dataset (LowLevelDataset) – 用于构建 MegatronDataset 的数据集
dataset_path (Optional[str]) – 数据集在磁盘上的真实路径，用于记录
indices (numpy.ndarray) – 要公开的文档索引集
num_samples (Optional[int]) – 从索引数据集中构建的最小样本数。如果为 None，则构建与一个 epoch 对应的尽可能多的样本。
index_split (Split) – 索引 Split
config (BlendedMegatronDatasetConfig) – 配置

static build_low_level_dataset(dataset_path: str, config: megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig) → Union[megatron.core.datasets.indexed_dataset.IndexedDataset, Iterable]

通过从 BlendedMegatronDatasetBuilder.build_generic_dataset 中调用的函数构建低级别数据集

可能是低级别数据集跨越 train/valid/test 分割的任何子集，这就是为什么我们定义了一个静态的“build”函数，与中级别数据集类中的构造函数分开

参数

dataset_path (str) – 数据集在磁盘上的真实路径
config (BlendedMegatronDatasetConfig) – 数据集配置

返回值

低级别数据集

返回类型

LowLevelDataset

static numel_low_level_dataset(low_level_dataset: Union[megatron.core.datasets.indexed_dataset.IndexedDataset, Iterable]) → int

返回底层低级别数据集中的元素数量，以便分隔 train/valid/test 分割索引

可能是低级别数据集可以根据它支持的中级别数据集以任意数量的方式进行分割，这就是为什么我们在中级别数据集类中将“元素数量”函数与 __len__ 函数分开定义的原因

参数: low_level_dataset (LowLevelDataset) – 底层低级别数据集
返回值: 底层低级别数据集中的元素数量
返回类型: int

datasets.gpt_dataset 模块

class core.datasets.gpt_dataset.GPTDataset(*args: Any, **kwargs: Any)

基类： megatron.core.datasets.megatron_dataset.MegatronDataset

基本 GPT 数据集

参数

indexed_dataset (IndexedDataset) – 用于构建 GPTDataset 的 IndexedDataset
dataset_path (Optional[str]) – 数据集在磁盘上的真实路径，用于记录
indexed_indices (numpy.ndarray) – 要公开的文档索引集
num_samples (Optional[int]) – 从索引数据集中抽取的样本数。如果为 None，则构建与一个 epoch 对应的尽可能多的样本。
index_split (Split) – indexed_indices Split
config (GPTDatasetConfig) – 配置

static build_low_level_dataset(dataset_path: str, config: core.datasets.gpt_dataset.GPTDatasetConfig) → megatron.core.datasets.indexed_dataset.IndexedDataset

抽象方法实现

参数

dataset_path (str) – IndexedDataset .bin 和 .idx 文件的真实路径前缀
config (GPTDatasetConfig) – 配置

返回值

底层 IndexedDataset

返回类型

IndexedDataset

static numel_low_level_dataset(low_level_dataset: megatron.core.datasets.indexed_dataset.IndexedDataset) → int

抽象方法实现

对于 GPT，底层 IndexedDataset 应按序列分割，而不是像 BERT 那样按文档分割

参数: low_level_dataset (IndexedDataset) – 底层 IndexedDataset
返回值: 底层 IndexedDataset 中唯一元素的数量
返回类型: int

class core.datasets.gpt_dataset.GPTDatasetConfig(*args: Any, **kwargs: Any)

基类： megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig

Megatron Core GPT 数据集的配置对象

add_extra_token_to_sequence: bool = True: 选项：绘制具有一个额外 token 的序列，以确保样本输入 token 和样本输出 token 都具有所需的序列长度

create_attention_mask: bool = True: 选项：启用 attention mask 生成。如果 attention kernel 本身生成 mask，则可以禁用。

drop_last_partial_validation_sequence: bool = True: 选项：删除最后一个部分验证序列

eod_mask_loss: bool = None: 选项：启用 EOD mask loss

reset_attention_mask: bool = None: 选项：从数据集重置 attention mask

reset_position_ids: bool = None: 选项：以一定间隔重置数据集中的位置 ID

s3_cache_path: str = None: 用于缓存 s3 数据加载的索引的路径。

class core.datasets.gpt_dataset.MockGPTDataset(*args: Any, **kwargs: Any)

基类： core.datasets.gpt_dataset.GPTDataset

模拟 GPT 数据集

参数

indexed_dataset (MockGPTLowLevelDataset) – 用于构建 MockGPTDataset 的 MockGPTLowLevelDataset
dataset_path (Optional[str]) – 此参数对于 MockGPTDataset 无关紧要
indices (numpy.ndarray) – 要公开的数据集索引集
num_samples (int) – 从数据集抽取的样本数
index_split (Split) – 索引 Split
config (GPTDatasetConfig) – 配置

static build_low_level_dataset(dataset_path: Optional[str], config: core.datasets.gpt_dataset.GPTDatasetConfig) → core.datasets.gpt_dataset.MockGPTLowLevelDataset

抽象方法实现

参数

dataset_path (Optional[str]) – 此参数对于 MockGPTLowLevelDataset 无关紧要
config (GPTDatasetConfig) – 配置

返回值

底层 MockGPTLowLevelDataset

返回类型

MockGPTLowLevelDataset

static numel_low_level_dataset(low_level_dataset: core.datasets.gpt_dataset.MockGPTLowLevelDataset) → int

抽象方法实现

参数: low_level_dataset (MockGPTLowLevelDataset) – 底层 MockGPTLowLevelDataset
返回值: 底层 MockGPTLowLevelDataset 中唯一元素的数量
返回类型: int

class core.datasets.gpt_dataset.MockGPTLowLevelDataset(tokenizer: megatron.core.datasets.megatron_tokenizer.MegatronTokenizer)

基类：object

模拟 GPT 低级别数据集

此类旨在生成经典 “Megatron-LM” GPT 风格的 tokenized 数据。值得注意的是，我们将文档结束 token 添加到 __getitem__ 中索引的每个元素

参数: tokenizer (MegatronTokenizer) – tokenizer，我们使用其特殊 token 信息来扩充模拟数据。

get(idx: int, offset: int = 0, length: Optional[int] = None) → numpy.ndarray

此函数是对 __getitem__ 的抽象，支持切片

参数

idx (int) – 数据集中的索引
offset (int) – 序列中的整数 token 偏移量
length (Optional[int]) – 从序列中获取的 token 数量

返回值

索引处的序列 token

返回类型

numpy.ndarray

max_sequence_length: int = 4096: 要生成的硬编码最大序列长度

seed: int = 0: 用于设置 NumPy RNG 的硬编码随机种子

size: int = 100000: 要生成的硬编码样本数

datasets.masked_dataset 模块

class core.datasets.masked_dataset.MaskedWordPieceDataset(*args: Any, **kwargs: Any)

基类： megatron.core.datasets.megatron_dataset.MegatronDataset

masked WordPiece 数据集的半抽象基类

此实现严格假设所有继承者数据集都构建在 IndexedDataset 类之上。如果需要，此假设可能会在未来被下推到继承者。

注意：WordPiece tokenization 在单词中的所有 token/piece 前添加双哈希 “##”，除了第一个 token/piece。

参数

indexed_dataset (IndexedDataset) – 用于构建 MegatronDataset 的 IndexedDataset
dataset_path (str) – 数据集在磁盘上的真实路径，用于记录
indexed_indices (numpy.ndarray) – 要公开的文档索引集
num_samples (Optional[int]) – 从索引数据集中抽取的样本数。如果为 None，则构建与一个 epoch 对应的尽可能多的样本。
index_split (Split) – indexed_indices Split
config (MaskedWordPieceDatasetConfig) – 配置

static build_low_level_dataset(dataset_path: str, config: core.datasets.masked_dataset.MaskedWordPieceDatasetConfig) → megatron.core.datasets.indexed_dataset.IndexedDataset

static numel_low_level_dataset(low_level_dataset: megatron.core.datasets.indexed_dataset.IndexedDataset) → int

class core.datasets.masked_dataset.MaskedWordPieceDatasetConfig(*args: Any, **kwargs: Any)

基类： megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig

Megatron Core Masked WordPiece 数据集的配置对象

masking_do_full_word: bool = None: 是否 mask 整个单词或其组成部分

masking_do_permutation: bool = None: 是否额外 shuffle 候选 N-gram 的子集

masking_max_ngram: int = None: 要考虑 mask 或 permute 的最大长度 N-gram

masking_probability: float = None: 我们 mask 候选 N-gram 的概率

masking_use_geometric_distribution: bool = None: 是否根据 SpanBERT 从几何分布中抽取 N-gram 的大小 https://arxiv.org/abs/1907.10529 (第 3.1 节)

masking_use_longer_ngrams: bool = None: 是否优先选择较长的 N-gram 而不是较短的 N-gram

short_sequence_probability: float = None: 我们返回比目标序列长度更短的序列的概率

datasets.bert_dataset 模块

class core.datasets.bert_dataset.BERTMaskedWordPieceDataset(*args: Any, **kwargs: Any)

基类： megatron.core.datasets.masked_dataset.MaskedWordPieceDataset

假设 WordPiece tokenization 的 BERT 数据集

参数

indexed_dataset (IndexedDataset) – 用于构建 MegatronDataset 的 IndexedDataset
dataset_path (str) – 数据集在磁盘上的真实路径，用于记录
indexed_indices (numpy.ndarray) – 要公开的文档索引集
num_samples (Optional[int]) – 从索引数据集中抽取的样本数。如果为 None，则构建与一个 epoch 对应的尽可能多的样本。
index_split (Split) – indexed_indices Split
config (BERTMaskedWordPieceDatasetConfig) – 配置

class core.datasets.bert_dataset.BERTMaskedWordPieceDatasetConfig(*args: Any, **kwargs: Any)

基类： megatron.core.datasets.masked_dataset.MaskedWordPieceDatasetConfig

Megatron Core BERT WordPiece 数据集的配置对象

classification_head: bool = None: 选项：在采样期间执行下一个序列预测

datasets.t5_dataset 模块

class core.datasets.t5_dataset.T5MaskedWordPieceDataset(*args: Any, **kwargs: Any)

基类： megatron.core.datasets.masked_dataset.MaskedWordPieceDataset

假设 WordPiece tokenization 的 T5 数据集

参数

indexed_dataset (IndexedDataset) – 用于构建 MegatronDataset 的 IndexedDataset
dataset_path (str) – 数据集在磁盘上的真实路径，用于记录
indexed_indices (numpy.ndarray) – 要公开的文档索引集
num_samples (Optional[int]) – 从索引数据集中抽取的样本数。如果为 None，则构建与一个 epoch 对应的尽可能多的样本。
index_split (Split) – indexed_indices Split
config (T5MaskedWordPieceDatasetConfig) – 配置

static config_attention_mask(encoder_tokens: torch.tensor, decoder_tokens: torch.tensor, encoder_mask: torch.tensor, decoder_mask: torch.tensor, use_local: bool = False, test_te_version: Optional[str] = None) → torch.tensor

根据 transformer 实现（例如 TE 与 local）、TE 版本和 TE 后端，配置 encoder_mask、decoder_mask、encoder_decoder_mask 的 attention-mask

参数

encoder_tokens (torch.tensor) – token 的 2-D 数组 (bs, kv_len)
decoder_tokens (torch.tensor) – token 的 2-D 数组 (bs, q_len)
encoder_mask (torch.tensor) – token 的 2-D 数组 (bs, kv_len)
decoder_mask (torch.tensor) – token 的 2-D 数组 (bs, q_len)
use_local (bool) – 当前 T5 模型是否使用 local (vs TE) transformer 实现

返回值

已配置的 encoder_mask、decoder_mask、encoder_decoder_mask torch.tensor: 已配置的 encoder attention mask torch.tensor: 已配置的 decoder attention mask torch.tensor: 已配置的 encoder-decoder attention mask

class core.datasets.t5_dataset.T5MaskedWordPieceDatasetConfig(*args: Any, **kwargs: Any)

基类： megatron.core.datasets.masked_dataset.MaskedWordPieceDatasetConfig

Megatron Core T5 WordPiece 数据集的配置对象

注意：作为 Megatron-LM 的临时遗留。T5 tokenizer 具有一个属性，该属性定义了采样期间使用的许多特殊 sentinel token。__post_init__ 中的断言用于保持与 Megatron-LM 的兼容性，直到 T5 tokenizer 进入 Megatron Core。

sequence_length_decoder: int = None: decoder 的序列长度

sequence_length_encoder: Optional[int] = None: sequence_length 别名和 encoder 的序列长度

datasets.blended_dataset 模块

class core.datasets.blended_dataset.BlendedDataset(*args: Any, **kwargs: Any)

基类：torch.utils.data.Dataset

MegatronDataset 实例集合的结合类

参数

datasets (List[MegatronDataset]) – 要混合的 MegatronDataset 实例
weights (List[Union[int, float]]) – 确定数据集混合比率的权重
size (Optional[int]) – 从混合中抽取的样本数。如果为 None，则对于每个数据集索引 idx，从 datasets[idx] 中精确抽取 weights[idx] 个样本。
config (BlendedMegatronDatasetConfig) – 配置

异常

RuntimeError – 当数据集在初始化后的样本数少于或多于 ‘size’ 时

datasets.utils 模块

class core.datasets.utils.Split(value)

基类：enum.Enum

一个枚举。

test = 2

train = 0

valid = 1

core.datasets.utils.compile_helpers(): 在运行时编译 C++ 辅助函数。确保在单个进程上调用此函数。

core.datasets.utils.get_blend_from_list(blend: Optional[List[str]]) → Optional[Tuple[List[str], Optional[List[float]]]]

从 blend 列表获取 megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig blend

参数: blend (Optional[List[str]]) – blend 列表，可以是 (1) 前缀列表，例如 [“path/to/dataset_1_prefix”, “path/to/dataset_2_prefix”]，或 (2) 扁平化的、压缩的权重和前缀列表，例如 [“30”, “path/to/dataset_1_prefix”, “70”, “path/to/dataset_2_prefix”]
返回值: blend，由数据集前缀列表和可选的数据集权重列表组成，例如 [[“path/to/dataset_1_prefix”, “path/to/dataset_2_prefix”], [30.0, 70.0]]。
返回类型: Optional[Tuple[List[str], Optional[List[float]]]]

core.datasets.utils.normalize(weights: List[float]) → List[float]

执行非指数归一化

参数: weights (List[float]) – 权重
返回值: 归一化后的权重
返回类型: List[float]

模块内容

上一个 distributed package

下一个 Microbatches Calculator