Accelerate documentation
Megatron-LM utilities
Megatron-LM utilities
MegatronLMPlugin
class accelerate.utils.MegatronLMPlugin
< source >( tp_degree: int = Nonepp_degree: int = Noneuse_custom_fsdp: bool = Noneoverlap_cpu_optimizer_d2h_h2d: bool = Noneno_load_optim: bool = Noneeod_mask_loss: bool = Noneno_save_optim: bool = Noneoptimizer_cpu_offload: bool = Noneuse_precision_aware_optimizer: bool = Nonedecoder_last_pipeline_num_layers: int = Nonerecompute_granularity: str = Nonerecompute_method: str = Nonerecompute_num_layers: int = Noneattention_backend: bool = Noneexpert_model_parallel_size: int = Nonecontext_parallel_size: int = Noneattention_dropout: float = Nonehidden_dropout: float = Noneattention_softmax_in_fp32: bool = Noneexpert_tensor_parallel_size: int = Nonecalculate_per_token_loss: bool = Noneuse_rotary_position_embeddings: bool = Nonenum_micro_batches: int = Nonegradient_clipping: float = Nonesequence_parallelism: bool = Nonerecompute_activations: bool = Noneuse_distributed_optimizer: bool = Nonepipeline_model_parallel_split_rank: int = Nonenum_layers_per_virtual_pipeline_stage: int = Noneis_train_batch_min: str = Truetrain_iters: int = Nonetrain_samples: int = Noneweight_decay_incr_style: str = 'constant'start_weight_decay: float = Noneend_weight_decay: float = Nonelr_decay_style: str = 'linear'lr_decay_iters: int = Nonelr_decay_samples: int = Nonelr_warmup_iters: int = Nonelr_warmup_samples: int = Nonelr_warmup_fraction: float = Nonemin_lr: float = 0consumed_samples: list = Noneno_wd_decay_cond: typing.Optional[typing.Callable] = Nonescale_lr_cond: typing.Optional[typing.Callable] = Nonelr_mult: float = 1.0megatron_dataset_flag: bool = Falseseq_length: int = Noneencoder_seq_length: int = Nonedecoder_seq_length: int = Nonetensorboard_dir: str = Noneset_all_logging_options: bool = Falseeval_iters: int = 100eval_interval: int = 1000return_logits: bool = Falsecustom_train_step_class: typing.Optional[typing.Any] = Nonecustom_train_step_kwargs: typing.Optional[dict[str, typing.Any]] = Nonecustom_model_provider_function: typing.Optional[typing.Callable] = Nonecustom_prepare_model_function: typing.Optional[typing.Callable] = Nonecustom_megatron_datasets_provider_function: typing.Optional[typing.Callable] = Nonecustom_get_batch_function: typing.Optional[typing.Callable] = Nonecustom_loss_function: typing.Optional[typing.Callable] = Noneother_megatron_args: typing.Optional[dict[str, typing.Any]] = None )
Parameters
- tp_degree (
int, defaults toNone) — Tensor parallelism degree. - pp_degree (
int, defaults toNone) — Pipeline parallelism degree. - num_micro_batches (
int, defaults toNone) — Number of micro-batches. - gradient_clipping (
float, defaults toNone) — Gradient clipping value based on global L2 Norm (0 to disable). - sequence_parallelism (
bool, defaults toNone) — Enable sequence parallelism. - recompute_activations (
bool, defaults toNone) — Enable selective activation recomputation. - use_distributed_optimizer (
bool, defaults toNone) — Enable distributed optimizer. - pipeline_model_parallel_split_rank (
int, defaults toNone) — Rank where encoder and decoder should be split. - num_layers_per_virtual_pipeline_stage (
int, defaults toNone) — Number of layers per virtual pipeline stage. - is_train_batch_min (
str, defaults toTrue) — If both tran & eval dataloaders are specified, this will decide themicro_batch_size. - train_iters (
int, defaults toNone) — Total number of samples to train over all training runs. Note that either train-iters or train-samples should be provided when usingMegatronLMDummyScheduler. - train_samples (
int, defaults toNone) — Total number of samples to train over all training runs. Note that either train-iters or train-samples should be provided when usingMegatronLMDummyScheduler. - weight_decay_incr_style (
str, defaults to'constant') — Weight decay increment function. choices=[“constant”, “linear”, “cosine”]. - start_weight_decay (
float, defaults toNone) — Initial weight decay coefficient for L2 regularization. - end_weight_decay (
float, defaults toNone) — End of run weight decay coefficient for L2 regularization. - lr_decay_style (
str, defaults to'linear') — Learning rate decay function. choices=[‘constant’, ‘linear’, ‘cosine’]. - lr_decay_iters (
int, defaults toNone) — Number of iterations for learning rate decay. If None defaults totrain_iters. - lr_decay_samples (
int, defaults toNone) — Number of samples for learning rate decay. If None defaults totrain_samples. - lr_warmup_iters (
int, defaults toNone) — Number of iterations to linearly warmup learning rate over. - lr_warmup_samples (
int, defaults toNone) — Number of samples to linearly warmup learning rate over. - lr_warmup_fraction (
float, defaults toNone) — Fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over. - min_lr (
float, defaults to0) — Minimum value for learning rate. The scheduler clip values below this threshold. - consumed_samples (
List, defaults toNone) — Number of samples consumed in the same order as the dataloaders toaccelerator.preparecall. - no_wd_decay_cond (
Optional, defaults toNone) — Condition to disable weight decay. - scale_lr_cond (
Optional, defaults toNone) — Condition to scale learning rate. - lr_mult (
float, defaults to1.0) — Learning rate multiplier. - megatron_dataset_flag (
bool, defaults toFalse) — Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format. - seq_length (
int, defaults toNone) — Maximum sequence length to process. - encoder_seq_length (
int, defaults toNone) — Maximum sequence length to process for the encoder. - decoder_seq_length (
int, defaults toNone) — Maximum sequence length to process for the decoder. - tensorboard_dir (
str, defaults toNone) — Path to save tensorboard logs. - set_all_logging_options (
bool, defaults toFalse) — Whether to set all logging options. - eval_iters (
int, defaults to100) — Number of iterations to run for evaluation validation/test for. - eval_interval (
int, defaults to1000) — Interval between running evaluation on validation set. - return_logits (
bool, defaults toFalse) — Whether to return logits from the model. - custom_train_step_class (
Optional, defaults toNone) — Custom train step class. - custom_train_step_kwargs (
Optional, defaults toNone) — Custom train step kwargs. - custom_model_provider_function (
Optional, defaults toNone) — Custom model provider function. - custom_prepare_model_function (
Optional, defaults toNone) — Custom prepare model function. - custom_megatron_datasets_provider_function (
Optional, defaults toNone) — Custom megatron train_valid_test datasets provider function. - custom_get_batch_function (
Optional, defaults toNone) — Custom get batch function. - custom_loss_function (
Optional, defaults toNone) — Custom loss function. - other_megatron_args (
Optional, defaults toNone) — Other Megatron-LM arguments. Please refer Megatron-LM.
Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective activation recomputation and optimized fused kernels.
MegatronLMDummyScheduler
class accelerate.utils.MegatronLMDummyScheduler
< source >( optimizertotal_num_steps = Nonewarmup_num_steps = 0**kwargs )
Dummy scheduler presents model parameters or param groups, this is primarily used to follow conventional training loop when scheduler config is specified in the deepspeed config file.
MegatronLMDummyDataLoader
class accelerate.utils.MegatronLMDummyDataLoader
< source >( **dataset_kwargs )
Dummy dataloader presents model parameters or param groups, this is primarily used to follow conventional training
AbstractTrainStep
Abstract class for batching, forward pass and loss handler.
GPTTrainStep
class accelerate.utils.GPTTrainStep
< source >( acceleratorargs )
GPT train step class.
BertTrainStep
class accelerate.utils.BertTrainStep
< source >( acceleratorargs )
Bert train step class.
T5TrainStep
class accelerate.utils.T5TrainStep
< source >( acceleratorargs )
T5 train step class.
avg_losses_across_data_parallel_group
accelerate.utils.avg_losses_across_data_parallel_group
< source >( losses )
Average losses across data parallel group.