Microbatches 计算器

此 API 用于计算在给定批次大小下拟合给定模型所需的 microbatches 数量。

Megatron Core microbatches 计算器数量。

class core.num_microbatches_calculator.ConstantNumMicroBatchesCalculator(global_batch_size: int, micro_batch_size: int, data_parallel_size: int, decrease_batch_size_if_needed: bool, rank: int)

基类：core.num_microbatches_calculator.NumMicroBatchesCalculator

具有恒定全局批次大小的 microbatches 数量计算器。

参数

• global_batch_size (int) – 全局批次大小。

• micro_batch_size (int) – Micro 批次大小。

• data_parallel_size (int) – 数据并行大小。

• decrease_batch_size_if_needed (bool) – 如果为 true，则减小批次大小以确保可被 DP 大小 * microbatch 大小整除（如果需要）。

• rank (int) – 秩（用于确定是否应执行日志记录）。

update(consumed_samples, consistency_check, verbose=False) → None

根据批次大小 rampup 更新 microbatches 数量。

class core.num_microbatches_calculator.NumMicroBatchesCalculator

基类：abc.ABC

microbatches 计算器基类。

get() → int

获取 microbatches 数量。

get_current_global_batch_size() → int

获取当前全局批次大小。

get_current_running_global_batch_size() → int

获取当前运行的全局批次大小。如果 decrease_batch_size_if_needed 为 False，则此值等于全局批次大小。

get_micro_batch_size() → int

获取当前全局批次大小。

abstract update(consumed_samples, consistency_check, verbose=False) → None

根据批次大小 rampup 更新 microbatches 数量。

class core.num_microbatches_calculator.RampupBatchsizeNumMicroBatchesCalculator(global_batch_size: int, micro_batch_size: int, data_parallel_size: int, decrease_batch_size_if_needed: bool, rank: int, start_global_batch_size: int, batch_size_increment: int, ramup_samples: int)

基类：core.num_microbatches_calculator.NumMicroBatchesCalculator

具有批次大小 rampup 的 microbatches 数量计算器。在 steps = (global-batch-size - start-batch-size) / batch_size_increment 上，使用 rampup-samples / steps 样本将批次大小从 start-batch-size 增加到 global-batch-size。

参数

• global_batch_size (int) – Rampup 后的全局批次大小。

• micro_batch_size (int) – Micro 批次大小。

• data_parallel_size (int) – 数据并行大小。

• decrease_batch_size_if_needed (bool) – 如果为 true，则减小批次大小以确保可被 DP 大小 * microbatch 大小整除（如果需要）。

• rank (int) – 秩（用于确定是否应执行日志记录）。

• start_global_batch_size (int) – 要开始使用的全局批次大小。

• batch_size_increment (int) – 全局批次大小增量。

• ramup_samples (int) – 用于将全局批次大小从 start_global_batch_size rampup 到 global_batch_size 的样本数。

update(consumed_samples: int, consistency_check: bool, verbose: bool = False) → None

更新 microbatches 数量。

参数

• consumed_samples (int) – 已消耗的样本数。

• consistency_check (bool) – 用于检查当前计划一致性的选项。

• verbose (bool, optional) – 用于控制日志记录的选项。默认为 False。

core.num_microbatches_calculator.destroy_num_microbatches_calculator()

销毁 microbatches 计算器。

core.num_microbatches_calculator.get_current_global_batch_size() → int

获取当前全局批次大小。

core.num_microbatches_calculator.get_current_running_global_batch_size() → int

获取当前运行的全局批次大小，如果 decrease_batch_size_if_needed 为 True，则考虑到 DP 副本数量可能与真正的全局批次大小不兼容。

core.num_microbatches_calculator.get_micro_batch_size() → int

获取 micro 批次大小。

core.num_microbatches_calculator.get_num_microbatches() → int

获取 microbatches 数量。

core.num_microbatches_calculator.init_num_microbatches_calculator(rank: int, rampup_batch_size: Optional[List[int]], global_batch_size: int, micro_batch_size: int, data_parallel_size: int, decrease_batch_size_if_needed: bool = False) → None

初始化 microbatches 计算器。支持向后兼容性。

参数

• rank (int) – GPU 的秩，只有秩 0 会记录信息。

• rampup_batch_size (Optional[List[int]]) – Rampup 批次大小，应采用 [start_global_batch_size, batch_size_increment, ramup_samples] 格式。

• global_batch_size (int) – 模型的全局批次大小。

• micro_batch_size (int) – 初始化时的 Micro 批次大小。

• data_parallel_size (int) – 数据并行大小。

• decrease_batch_size_if_needed (bool, optional) – 如果为 true，则缩小批次大小以确保可被 DP 大小 * microbatch 大小整除。默认为 False。

core.num_microbatches_calculator.reconfigure_num_microbatches_calculator(rank: int, rampup_batch_size: Optional[List[int]], global_batch_size: int, micro_batch_size: int, data_parallel_size: int, decrease_batch_size_if_needed: bool = False) → None

重新配置 microbatches 计算器。支持向后兼容性。

参数

• rank (int) – GPU 的秩，只有秩 0 会记录信息。

• rampup_batch_size (Optional[List[int]]) – Rampup 批次大小，应采用 [start_global_batch_size, batch_size_increment, ramup_samples] 格式。

• global_batch_size (int) – 模型的全局批次大小。

• micro_batch_size (int) – 初始化时的 Micro 批次大小。

• data_parallel_size (int) – 数据并行大小。

• decrease_batch_size_if_needed (bool, optional) – 如果为 true，则缩小批次大小以确保可被 DP 大小 * microbatch 大小整除。默认为 False。

core.num_microbatches_calculator.unset_num_microbatches_calculator()

取消设置 microbatches 计算器。

用于多次运行。有关示例，请参见 tests/unit_tests/ckpt_converter/test_ckpt_converter.py。

core.num_microbatches_calculator.update_num_microbatches(consumed_samples: int, consistency_check: bool = True, verbose: bool = False) → None

更新 microbatches 数量。

参数

• consumed_samples (int) – 已消耗的样本数。

• consistency_check (bool, optional) – 用于检查当前计划一致性的选项。默认为 True。

• verbose (bool, optional) – 用于控制日志记录的选项。默认为 False。