megatron_gpt_config.yaml

defaults:
  - _self_
  - optional tp_overlap@model.ub_tp_comm_overlap_cfg:

name: megatron_gpt

trainer:
  devices: 1
  num_nodes: 1
  accelerator: gpu
  precision: bf16
  logger: False # logger provided by exp_manager
  enable_checkpointing: False
  use_distributed_sampler: False
  max_epochs: -1 # PTL default. In practice, max_steps will be reached first.
  max_steps: 100000 # consumed_samples = global_step * micro_batch_size * data_parallel_size * accumulate_grad_batches
  log_every_n_steps: 10
  val_check_interval: 100
  limit_val_batches: 50
  limit_test_batches: 500
  accumulate_grad_batches: 1 # do not modify, grad acc is automatic for training megatron models
  gradient_clip_val: 1.0
  benchmark: False
  enable_model_summary: False # default PTL callback for this does not support model parallelism, instead we log manually

# Used for S3 Checkpointing
s3_checkpointing:
  # write_concurrency * tp * pp * 1.15 (buffer) should be within 3500 S3 TPS limit per partition
  max_write_concurrency: 10
  # read_concurrency * tp * pp * 1.15 (buffer) should be within 5500 S3 TPS limit per partition
  max_read_concurrency: 15
  chunk_size_MB: 64
  # enables asynchronous checkpoint writing to S3 dirpath. the feature is experimental and currently does not check if the past save succeeded. Therefore, use in conjunction with save_top_k. 
  enable_async_checkpointing: False

exp_manager:
  explicit_log_dir: null
  exp_dir: null
  name: megatron_gpt
  create_wandb_logger: False
  wandb_logger_kwargs:
    project: null
    name: null
  create_neptune_logger: false
  neptune_logger_kwargs:
    project: null
    name: null
    prefix: train
    log_model_checkpoints: false
    tags: null # can specify as an array of strings in yaml array format
    description: null
  resume_if_exists: True
  resume_ignore_no_checkpoint: True
  resume_from_checkpoint: ${model.resume_from_checkpoint}
  create_checkpoint_callback: True
  checkpoint_callback_params:
    dirpath: null # to use S3 checkpointing, set the dirpath in format s3://bucket/key
    monitor: val_loss
    save_top_k: 10
    mode: min
    always_save_nemo: False # saves nemo file during validation, not implemented for model parallel
    save_nemo_on_train_end: False # not recommended when training large models on clusters with short time limits
    filename: 'megatron_gpt--{val_loss:.2f}-{step}-{consumed_samples}'
    model_parallel_size: ${multiply:${model.tensor_model_parallel_size}, ${model.pipeline_model_parallel_size}}
    async_save: False # Set to True to enable async checkpoint save. Currently works only with distributed checkpoints
    save_last_n_optim_states: -1 # a number of last checkpoints to be saved with optimizer states

model:
  # The following two settings are used for continual training:
  restore_from_path: null # Set this to a .nemo file path to restore only the model weights
  restore_from_ckpt: null # Set this to a training ckpt path to restore both model weights and optimizer states

  # use GPTModel from megatron.core
  mcore_gpt: True
  
  # Skip training and do only validation step
  skip_train: False

  # specify micro_batch_size, global_batch_size, and model parallelism
  # gradient accumulation will be done automatically based on data_parallel_size
  micro_batch_size: 4 # limited by GPU memory
  global_batch_size: 8 # will use more micro batches to reach global batch size
  rampup_batch_size: null # Should be a list of 3 values: [<start_batch_size>, <batch_size_increment>, <rampup_samples>]
  tensor_model_parallel_size: 1 # intra-layer model parallelism
  pipeline_model_parallel_size: 1 # inter-layer model parallelism
  virtual_pipeline_model_parallel_size: null # interleaved pipeline
  expert_model_parallel_size: 1 # expert model parallelism

  # model architecture
  encoder_seq_length: 512
  max_position_embeddings: ${.encoder_seq_length}
  num_layers: 12
  hidden_size: 768
  ffn_hidden_size: 3072 # Transformer FFN hidden size. Usually 4 * hidden_size.
  num_attention_heads: 12
  init_method_std: 0.02 # Standard deviation of the zero mean normal distribution used for weight initialization.')
  use_scaled_init_method: True # use scaled residuals initialization
  hidden_dropout: 0.1 # Dropout probability for hidden state transformer.
  attention_dropout: 0.1 # Dropout probability for attention
  ffn_dropout: 0.0 # Dropout probability in the feed-forward layer.
  kv_channels: null # Projection weights dimension in multi-head attention. Set to hidden_size // num_attention_heads if null
  apply_query_key_layer_scaling: False # scale Q * K^T by 1 / layer-number.
  normalization: 'layernorm' # Normalization layer to use. Options are 'layernorm', 'rmsnorm'
  layernorm_epsilon: 1e-5
  do_layer_norm_weight_decay: False # True means weight decay on all params
  make_vocab_size_divisible_by: 128 # Pad the vocab size to be divisible by this value for computation efficiency.
  pre_process: True # add embedding
  post_process: True # add pooler
  persist_layer_norm: True # Use of persistent fused layer norm kernel.
  bias: True # Whether to use bias terms in all weight matrices.
  activation: 'gelu' # Options ['gelu', 'geglu', 'swiglu', 'reglu', 'squared-relu', 'fast-geglu', 'fast-swiglu', 'fast-reglu']
  headscale: False # Whether to learn extra parameters that scale the output of the each self-attention head.
  transformer_block_type: 'pre_ln' # Options ['pre_ln', 'post_ln', 'normformer']
  openai_gelu: False # Use OpenAI's GELU instead of the default GeLU
  normalize_attention_scores: True # Whether to scale the output Q * K^T by 1 / sqrt(hidden_size_per_head). This arg is provided as a configuration option mostly for compatibility with models that have been weight-converted from HF. You almost always want to se this to True.
  position_embedding_type: 'learned_absolute' # Position embedding type. Options ['learned_absolute', 'rope', 'alibi', 'kerple' , 'xpos', 'sandwich'] xpos and sandwich are experimental.
  rotary_percentage: 1.0 # If using position_embedding_type=rope, then the per head dim is multiplied by this.
  attention_type: 'multihead' # Attention type. Options ['multihead']
  share_embeddings_and_output_weights: True # Share embedding and output layer weights.
  overlap_p2p_comm: False # Overlap p2p communication with computes. This argument is valid only when `virtual_pipeline_model_parallel_size` is larger than 1
  batch_p2p_comm: True # Batch consecutive inter-peer send/recv operations. This argument is valid only when `virtual_pipeline_model_parallel_size` is larger than 1
  seq_len_interpolation_factor: null # RoPE Interpolation factor for sequence length. This is used to build long-context models with RoPE ex: https://arxiv.org/abs/2306.15595.
  num_query_groups: null # Number of query groups for group query attention. If None, normal attention is used.
  scale_positional_embedding: False # Apply scaling for RoPE frequencies 
  drop_layers: [] # Select which layers should be ignored. Should be a list of layer numbers such as [0, 1, 4].

  ## Reset learning rate schedule.
  # 1. reset_lr=True, reset_lr_steps=False. When pre-training an existing checkpoint "from scratch" on a different dataset.
  # 2. reset_lr=True, reset_lr_steps=True. When continuing training from an existing checkpoint with the same configuration.
  #    Learning rate's max_steps and decay_steps will be recalculated as follows: max_steps -= completed_steps, decay_steps -= completed_steps where completed_steps is the number of steps already completed at the checkpoint.
  #    This will help to reach the min_lr value by the end of training without changing trainer.max_steps.
  reset_lr: False # Set to True to reset learning rate to initial learning rate. Only supported with distributed optmizer and megatron_amp_O2.
  reset_lr_steps: False # Set to True to adjust learning rate's max_steps and decay_steps by subtracting number of steps already completed at the checkpoint.

  tokenizer:
    library: 'megatron'
    type: 'GPT2BPETokenizer'
    model: null
    vocab_file: null
    merge_file: null
    delimiter: null # only used for tabular tokenizer
    sentencepiece_legacy: False # Legacy=True allows you to add special tokens to sentencepiece tokenizers.

  # Mixed precision
  native_amp_init_scale: 4294967296 # 2 ** 32
  native_amp_growth_interval: 1000
  hysteresis: 2 # Gradient scale hysteresis
  fp32_residual_connection: False # Move residual connections to fp32
  fp16_lm_cross_entropy: False # Move the cross entropy unreduced loss calculation for lm head to fp16

  # Megatron O2-style half-precision
  megatron_amp_O2: True # Enable O2-level automatic mixed precision using main parameters
  grad_allreduce_chunk_size_mb: 125

  # Fusion
  grad_div_ar_fusion: True # Fuse grad division into torch.distributed.all_reduce. Only used with O2 and no pipeline parallelism..
  gradient_accumulation_fusion: False # Fuse weight gradient accumulation to GEMMs. Only used with pipeline parallelism and O2.
  bias_activation_fusion: True # Use a kernel that fuses the bias addition from weight matrices with the subsequent activation function.
  bias_dropout_add_fusion: True # Use a kernel that fuses the bias addition, dropout and residual connection addition.
  masked_softmax_fusion: True # Use a kernel that fuses the attention softmax with it's mask.
  get_attention_mask_from_fusion: True # When using fused softmax it will create the attention mask so we won't copy it to the pipeline stages.
  apply_rope_fusion: False # Use a kernel to add rotary positional embeddings. Only used if position_embedding_type=rope


  # Miscellaneous
  seed: 1234
  resume_from_checkpoint: null # manually set the checkpoint file to load from
  use_cpu_initialization: False # Init weights on the CPU (slow for large models)
  onnx_safe: False # Use work-arounds for known problems with Torch ONNX exporter.
  apex_transformer_log_level: 30 # Python logging level displays logs with severity greater than or equal to this
  gradient_as_bucket_view: True # PyTorch DDP argument. Allocate gradients in a contiguous bucket to save memory (less fragmentation and buffer memory)
  sync_batch_comm: False # Enable stream synchronization after each p2p communication between pipeline stages
  nccl_communicator_config_path: null # Path to the yaml file with NCCL communicator options (min_ctas, max_ctas, and cga_cluster_size)
  validation_param_sync_overlap: False # Overlap parameter AllGather with validation step.

  # FSDP
  fsdp: False # Enable training with torch FSDP.
  fsdp_sharding_strategy: 'full' # Method to shard model states. Available options are 'full', 'hybrid', and 'grad'.
  fsdp_grad_reduce_dtype: 32 # Gradient reduction data type.
  fsdp_sharded_checkpoint: False # Store and load FSDP shared checkpoint.

  # Distributed checkpoint setup
  dist_ckpt_format: 'torch_dist' # Set to 'torch_dist' to use PyTorch distributed checkpoint format.
  dist_ckpt_load_on_device: True # whether to load checkpoint weights directly on GPU or to CPU
  dist_ckpt_parallel_save: True # if true, each worker will write its own part of the dist checkpoint
  dist_ckpt_parallel_save_within_dp: False # if true, save will be parallelized only within a DP group (whole world otherwise), which might slightly reduce the save overhead
  dist_ckpt_parallel_load: False # if true, each worker will load part of the dist checkpoint and exchange with NCCL. Might use some extra GPU memory
  dist_ckpt_torch_dist_multiproc: 2 # number of extra processes per rank used during ckpt save with PyTorch distributed format
  dist_ckpt_assume_constant_structure: False # set to True only if the state dict structure doesn't change within a single job. Allows caching some computation across checkpoint saves.
  dist_ckpt_parallel_dist_opt: True # parallel save/load of a DistributedOptimizer. 'True' allows performant save and reshardable checkpoints. Set to 'False' only in order to minimize the number of checkpoint files.
  dist_ckpt_load_strictness: null # defines checkpoint keys mismatch behavior (only during dist-ckpt load). Choices: assume_ok_unexpected (default - try loading without any check), log_all (log mismatches), raise_all (raise mismatches)

  ## Activation Checkpointing
  # NeMo Megatron supports 'selective' activation checkpointing where only the memory intensive part of attention is checkpointed.
  # These memory intensive activations are also less compute intensive which makes activation checkpointing more efficient for LLMs (20B+).
  # See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
  # 'full' will checkpoint the entire transformer layer.
  activations_checkpoint_granularity: null # 'selective' or 'full'
  activations_checkpoint_method: null # 'uniform', 'block'
  # 'uniform' divides the total number of transformer layers and checkpoints the input activation
  # of each chunk at the specified granularity. When used with 'selective', 'uniform' checkpoints all attention blocks in the model.
  # 'block' checkpoints the specified number of layers per pipeline stage at the specified granularity
  activations_checkpoint_num_layers: null
  # when using 'uniform' this creates groups of transformer layers to checkpoint. Usually set to 1. Increase to save more memory.
  # when using 'block' this this will checkpoint the first activations_checkpoint_num_layers per pipeline stage.
  num_micro_batches_with_partial_activation_checkpoints: null
  # This feature is valid only when used with pipeline-model-parallelism.
  # When an integer value is provided, it sets the number of micro-batches where only a partial number of Transformer layers get checkpointed
  # and recomputed within a window of micro-batches. The rest of micro-batches in the window checkpoint all Transformer layers. The size of window is
  # set by the maximum outstanding micro-batch backpropagations, which varies at different pipeline stages. The number of partial layers to checkpoint
  # per micro-batch is set by 'activations_checkpoint_num_layers' with 'activations_checkpoint_method' of 'block'.
  # This feature enables using activation checkpoint at a fraction of micro-batches up to the point of full GPU memory usage.
  activations_checkpoint_layers_per_pipeline: null
  # This feature is valid only when used with pipeline-model-parallelism.
  # When an integer value (rounded down when float is given) is provided, it sets the number of Transformer layers to skip checkpointing at later
  # pipeline stages. For example, 'activations_checkpoint_layers_per_pipeline' of 3 makes pipeline stage 1 to checkpoint 3 layers less than
  # stage 0 and stage 2 to checkpoint 6 layers less stage 0, and so on. This is possible because later pipeline stage
  # uses less GPU memory with fewer outstanding micro-batch backpropagations. Used with 'num_micro_batches_with_partial_activation_checkpoints',
  # this feature removes most of activation checkpoints at the last pipeline stage, which is the critical execution path.

  ## Sequence Parallelism
  # Makes tensor parallelism more memory efficient for LLMs (20B+) by parallelizing layer norms and dropout sequentially
  # See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
  sequence_parallel: False

  ## Transformer Engine
  transformer_engine: False
  fp8: False # enables fp8 in TransformerLayer forward
  fp8_e4m3: False # sets fp8_format = recipe.Format.E4M3
  fp8_hybrid: True # sets fp8_format = recipe.Format.HYBRID
  fp8_margin: 0 # scaling margin
  fp8_interval: 1 # scaling update interval
  fp8_amax_history_len: 1024 # Number of steps for which amax history is recorded per tensor
  fp8_amax_compute_algo: max # 'most_recent' or 'max'. Algorithm for computing amax from history
  reduce_amax: True # Perform reduction to sync amax tensors across GPUs after every iteration
  use_emha: False # Use fused multi-head attention for large sequence-length. Note this is not yet supported. Please set to False.
  ub_tp_comm_overlap: False
  # Use userbuffer backend to overlap tensor-parallel communications with computes.
  # This feature is only available with Transformer Engine and squence parallelism enabled and, currently, supports only GPT models.
  ub_tp_comm_overlap_cfg: null
  # A yaml file with userbuffer communicator configurations. This file should provide `method`, `dtype`, `num_sm`, `num_splits`,
  # `cga_size`, `num_splits`, `set_sm_margin`, and `aggregate` for the communicators to use custom settings.
  # If the configuration file is not provided a default setting is used for all communicators.

  ## Flash Attention
  use_flash_attention: False # Use flash attention in self-attention module, this config does nothing when transformer_engine=True

  ##Offloading Activations/Weights to CPU
  cpu_offloading: False
  cpu_offloading_num_layers: ${sum:${.num_layers},-1} #This value should be between [1,num_layers-1] as we don't want to offload the final layer's activations and expose any offloading duration for the final layer
  cpu_offloading_activations: True
  cpu_offloading_weights: True

  ## Network
  sharp: False # Enable the use of SHARP for NCCL data-parallel communications. This is going to be ignored if the network doesn't support SHARP.
  
  ## Megatron timers
  enable_megatron_timers: False
  megatron_timer_kwargs:
    log_every_n_steps: 10
    log_mode: minmax
    barrier: False

  data:
   # Path to data must be specified by the user.
    # Supports List, String and Dictionary
    # List : can override from the CLI: "model.data.data_prefix=[.5,/raid/data/pile/my-gpt3_00_text_document,.5,/raid/data/pile/my-gpt3_01_text_document]",
    # Or see example below:
    # data_prefix:
    #   - .5
    #   - /raid/data/pile/my-gpt3_00_text_document
    #   - .5
    #   - /raid/data/pile/my-gpt3_01_text_document
    # Dictionary: can override from CLI "model.data.data_prefix"={"train":[1.0, /path/to/data], "validation":/path/to/data, "test":/path/to/test}
    # Or see example below:
    # "model.data.data_prefix: {train:[1.0,/path/to/data], validation:[/path/to/data], test:[/path/to/test]}"
    data_prefix: ???
    index_mapping_dir: null # path to save index mapping .npy files, by default will save in the same location as data_prefix
    data_impl: mmap
    mmap_bin_files: True
    splits_string: 900,50,50
    seq_length: ${model.encoder_seq_length}
    skip_warmup: True
    num_workers: 2
    num_dataset_builder_threads: 1
    dataloader_type: single # cyclic
    reset_position_ids: False # Reset position ids after end-of-document token
    reset_attention_mask: False # Reset attention mask after end-of-document token
    eod_mask_loss: False # Mask loss for the end of document tokens
    validation_drop_last: True # Set to false if the last partial validation samples is to be consumed
    no_seqlen_plus_one_input_tokens: False # Set to True to disable fetching (sequence length + 1) input tokens, instead get (sequence length) input tokens and mask the last token
    pad_samples_to_global_batch_size: False # Set to True if you want to pad the last partial batch with -1's to equal global batch size
    shuffle_documents: True # Set to False to disable documents shuffling. Sample index will still be shuffled
    exchange_indices_distributed: False # Set to True to exchange indices via torch.distributed instead of filesystem
    data_cache_generation_only: False # Set to True to generate only the data cache and stop the training script
    renormalize_blend_weights: False # Renormalize the blend weights to account for mid-level dataset oversampling done to ensure fulfillmenet of the of the requested number of samples.

  # Nsys profiling options
  nsys_profile:
    enabled: False
    start_step: 10  # Global batch to start profiling
    end_step: 10 # Global batch to end profiling
    ranks: [0] # Global rank IDs to profile
    gen_shape: False # Generate model and kernel details including input shapes

  optim:
    name: fused_adam
    lr: 2e-4
    weight_decay: 0.01
    betas:
    - 0.9
    - 0.98
    sched:
      name: CosineAnnealing
      warmup_steps: 500
      constant_steps: 0
      min_lr: 2e-5

  gc_interval: 0
  # Interval of the host memory garbage collection. When it is zero, collectiion relies on the automatic garbage collector.
  # If an interger value larger than zero is set, collection is done manually by the batch step interval of `gc_interval`.