Megatron Core User Guide

Microbatches Calculator

This api is used to calculate the number of microbatches required to fit a given model on a given batch size.

Megatron Core number of microbatches calculators.

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)

Bases: core.num_microbatches_calculator.NumMicroBatchesCalculator

Calculator of number of microbatches with constant global batch size.

Parameters
  • global_batch_size (int) – Global batch size.

  • micro_batch_size (int) – Micro batch size.

  • data_parallel_size (int) – Data parallel size.

  • decrease_batch_size_if_needed (bool) – If true, decrease batch size to ensure divisibility by DP size * microbatch size (if needed).

  • rank (int) – Rank (to determine whether logging should be performed).

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

Update number of microbatches depending on batch size rampup.

class core.num_microbatches_calculator.NumMicroBatchesCalculator

Bases: abc.ABC

Base class for number of microbatches calculator.

get() → int

Get number of microbatches.

get_current_global_batch_size() → int

Get current global batch size.

get_current_running_global_batch_size() → int

Get current running global batch size. If decrease_batch_size_if_needed is False, this just equals global batch size.

get_micro_batch_size() → int

Get current global batch size.

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

Update number of microbatches depending on batch size rampup.

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)

Bases: core.num_microbatches_calculator.NumMicroBatchesCalculator

Calculator of number of microbatches with batch size rampup. Over steps = (global-batch-size - start-batch-size) / batch_size_increment increment batch size from start-batch-size to global-batch-size using rampup-samples / steps samples.

Parameters
  • global_batch_size (int) – Global batch size post rampup.

  • micro_batch_size (int) – Micro batch size.

  • data_parallel_size (int) – Data parallel size.

  • decrease_batch_size_if_needed (bool) – If true, decrease batch size to ensure divisibility by DP size * microbatch size (if needed).

  • rank (int) – Rank (to determine whether logging should be performed).

  • start_global_batch_size (int) – Global batch size to start with.

  • batch_size_increment (int) – Global batch size increments.

  • ramup_samples (int) – Number of samples to use ramp up global batch size from start_global_batch_size to global_batch_size.

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

Update number of microbatches.

Parameters
  • consumed_samples (int) – Number of samples consumed.

  • consistency_check (bool) – Option to check current schedule’s consistency.

  • verbose (bool, optional) – Option to control logging. Defaults to False.

core.num_microbatches_calculator.destroy_num_microbatches_calculator()

Destroy number of microbatches calculator.

core.num_microbatches_calculator.get_current_global_batch_size() → int

Get current global batch size.

core.num_microbatches_calculator.get_current_running_global_batch_size() → int

Get current running global batch size, taking into account number of DP replicas might be incompatible with true global batch size if decrease_batch_size_if_needed is True.

core.num_microbatches_calculator.get_micro_batch_size() → int

Get micro batch size.

core.num_microbatches_calculator.get_num_microbatches() → int

Get number of 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

Initialize number of microbatches calculator. Supporting backward compatibility.

Parameters
  • rank (int) – Rank of the GPU, only rank 0 will log the information.

  • rampup_batch_size (Optional[List[int]]) – Rampup batch size, should be in format of [start_global_batch_size, batch_size_increment, ramup_samples].

  • global_batch_size (int) – Global batch size for the model.

  • micro_batch_size (int) – Micro batch size at initialization.

  • data_parallel_size (int) – Data parallel size.

  • decrease_batch_size_if_needed (bool, optional) – If true, scale down batch size to ensure divisibility by DP size * microbatch size. Defaults to 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

Reconfigure number of microbatches calculator. Supporting backward compatibility.

Parameters
  • rank (int) – Rank of the GPU, only rank 0 will log the information.

  • rampup_batch_size (Optional[List[int]]) – Rampup batch size, should be in format of [start_global_batch_size, batch_size_increment, ramup_samples].

  • global_batch_size (int) – Global batch size for the model.

  • micro_batch_size (int) – Micro batch size at initialization.

  • data_parallel_size (int) – Data parallel size.

  • decrease_batch_size_if_needed (bool, optional) – If true, scale down batch size to ensure divisibility by DP size * microbatch size. Defaults to False.

core.num_microbatches_calculator.unset_num_microbatches_calculator()

Unset microbatches calculator.

Useful for multiple runs. See tests/unit_tests/ckpt_converter/test_ckpt_converter.py for an example.

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

Update number of microbatches.

Parameters
  • consumed_samples (int) – Number of samples consumed.

  • consistency_check (bool, optional) – Option to check current schedule’s consistency. Defaults to True.

  • verbose (bool, optional) – Option to control logging. Defaults to False.

Previous datasets package
Next Optimizer Parameters Scheduler
© Copyright 2022-2025, NVIDIA. Last updated on Jan 14, 2025.