Megatron backend configuration reference

This page lists the flat configuration keys exposed by Primus when framework: megatron. Unless a section says otherwise, values are the defaults from primus/configs/modules/megatron/trainer_base.yaml and related model presets. The effective pretraining preset is pre_trainer.yaml, which extends trainer_base.yaml and overrides several high-impact training defaults.

Where parameters live. Set overrides under modules.pre_trainer.overrides: in your experiment YAML. Model architecture keys usually come from models.<role>.overrides: (or your chosen model preset), but the same names map to Megatron’s argparse namespace either way.

Presets.

• Module presets: primus/configs/modules/megatron/ (the main pretraining bundle is pre_trainer.yaml, which extends trainer_base.yaml and Primus Megatron add-ons).

• Model presets: primus/configs/models/megatron/ (for example language_model.yaml).

Mapping to Megatron-LM. Keys are passed through 1:1 to Megatron’s training arguments (same names as argparse / Namespace). Primus builds that namespace with MegatronArgBuilder.

Upstream reference. Full flag semantics and newer options are defined in Megatron-LM: megatron/training/arguments.py.

Example (experiment YAML)

framework: megatron

modules:
  pre_trainer:
    overrides:
      global_batch_size: 256
      train_iters: 50000
      tensor_model_parallel_size: 2

models:
  pre_train:
    overrides:
      hidden_size: 2048
      num_layers: 32

---

1. Base module parameters

Source: primus/configs/modules/module_base.yaml (merged into Megatron presets; trainer_base.yaml sets trainable: true).

Parameter	Default	Description
trainable	true	When true, this module participates in training workflows. (module_base.yaml alone defaults to false; Megatron trainer_base.yaml overrides to true.)
sink_level	null	Log level for the structured sink (Primus module plumbing); null uses framework default.
file_sink_level	DEBUG	Minimum level for file-backed logging.
stderr_sink_level	INFO	Minimum level for stderr logging.

---

2. Training and batching

Source: primus/configs/modules/megatron/trainer_base.yaml; effective pre_trainer.yaml overrides are noted where they differ.

Parameter	Default	Description
yaml_cfg	null	Reserved; not supported as a Megatron override in this preset.
spec	null	Optional trainer spec hook (unused in defaults).
micro_batch_size	2	Samples per microbatch per data-parallel rank (per forward/backward step before gradient accumulation).
batch_size	null	Deprecated; use micro_batch_size / global_batch_size.
global_batch_size	128 (16 in pre_trainer.yaml)	Total batch size across the data-parallel world (before or after splitting, per Megatron semantics).
rampup_batch_size	null	Optional batch-size ramp schedule string / config.
decrease_batch_size_if_needed	false	Allow shrinking batch if memory is insufficient.
check_for_nan_in_loss_and_grad	true	Abort on NaNs in loss or gradients.
check_for_spiky_loss	false	Detect abnormal loss spikes.
check_for_large_grads	false	Detect abnormally large gradients.
make_vocab_size_divisible_by	128	Pads vocabulary size for efficient kernels / partitioning.
exit_signal_handler	false	Install handlers for graceful shutdown signals.
exit_duration_in_mins	null	Stop training after this many minutes.
exit_interval	null	Exit after this many iterations (if set).
onnx_safe	null	ONNX export compatibility tweaks.
bert_binary_head	true	Use BERT binary classification head when applicable.
use_flash_attn	false (true in pre_trainer.yaml)	Prefer FlashAttention kernels when available.
seed	1234	RNG seed for reproducibility.
data_parallel_random_init	false	Random init that varies across data-parallel ranks.
init_method_xavier_uniform	false	Use Xavier uniform for some weights.
test_mode	false	Lightweight test path (fewer steps / checks).
train_iters	null (1000 in pre_trainer.yaml)	Total training iterations (mutually exclusive with sample-based stopping in typical setups).
train_samples	null	Total training samples (when using sample-based training).
eval_iters	32 (0 in pre_trainer.yaml)	Validation iterations per eval.
eval_interval	2000 (1000 in pre_trainer.yaml)	Run validation every this many iterations.
full_validation	false	Run a full pass over validation data.
multiple_validation_sets	false	Multiple validation datasets / passes.
skip_train	false	Only run eval / test, no training updates.
train_sync_interval	null	Periodic distributed sync barrier for debugging.
adlr_autoresume	false	ADLR autoresume integration.
adlr_autoresume_interval	1000	Autoresume checkpoint interval.
manual_gc	false	Force Python GC on a schedule.
manual_gc_interval	1	GC every N steps when manual_gc is enabled.
manual_gc_eval	false	Run manual GC during evaluation.
mask_type	random	Masking strategy for MLM / similar objectives.
mask_factor	1.0	Masking strength multiplier.
iter_per_epoch	1250	Iterations interpreted as one “epoch” for logging.

---

3. Mixed precision

Source: trainer_base.yaml.

Parameter	Default	Description
fp16	false	Enable FP16 mixed precision training.
bf16	true	Enable BF16 mixed precision training.
grad_reduce_in_bf16	false	All-reduce gradients in BF16 (saves bandwidth).
calculate_per_token_loss	false	Normalize loss per token instead of per sample.
loss_scale	null	Static loss scale for FP16; null uses dynamic scaling.
initial_loss_scale	4294967296	Initial dynamic loss scale.
min_loss_scale	1.0	Floor for dynamic loss scale.
loss_scale_window	1000	Window for dynamic loss scaling updates.
hysteresis	2	Hysteresis steps for loss-scale decreases.
accumulate_allreduce_grads_in_fp32	false	Accumulate and reduce gradients in FP32.
fp16_lm_cross_entropy	false	Compute LM cross-entropy in FP16.
fp8	null	FP8 recipe selection (e4m3, hybrid, etc.); null disables.
fp8_margin	0	FP8 scaling margin.
fp8_recipe	delayed	FP8 recipe variant (e.g. delayed scaling).
fp8_interval	1	Deprecated FP8 interval (kept for compatibility).
fp8_amax_history_len	1024	History length for FP8 amax statistics.
fp8_amax_compute_algo	"max"	How to combine amax history (max, etc.).
fp8_wgrad	true	Run weight gradients in FP8 where supported.
fp8_param_gather	false	FP8 parameter gather for distributed optimizer paths.
te_rng_tracker	false	Transformer Engine RNG tracker for FP8.
inference_rng_tracker	false	Separate RNG tracker for inference FP8.
fp4	null	FP4 mode; null disables.
fp4_recipe	nvfp4	FP4 recipe name.
fp4_param	false	Store parameters in FP4.
first_last_layers_bf16	false	Keep first/last layers in BF16 for stability.
num_layers_at_start_in_bf16	1	Count of early layers forced to BF16 when enabled.
num_layers_at_end_in_bf16	1	Count of final layers forced to BF16 when enabled.
no_fp8_weight_transpose_cache	false	Primus: disable FP8 weight transpose cache (see primus_megatron_module.yaml).

---

4. Optimizer and learning rate

Source: trainer_base.yaml.

Parameter	Default	Description
optimizer	adam	Optimizer family (adam, sgd, etc.).
lr	2.5e-4 (2.0e-05 in pre_trainer.yaml)	Peak learning rate.
lr_decay_style	cosine	LR decay schedule (cosine, linear, constant, WSD, etc.).
lr_decay_iters	null	Decay duration in iterations.
lr_decay_samples	null	Decay duration in samples.
lr_warmup_fraction	null	Warmup as a fraction of total train steps.
lr_warmup_iters	0 (40 in pre_trainer.yaml)	Linear warmup steps.
lr_warmup_samples	0	Warmup in samples.
lr_warmup_init	0.0	LR at the start of warmup.
min_lr	2.5e-5 (0.0 in pre_trainer.yaml)	Minimum LR after decay.
lr_wsd_decay_style	exponential	Weight-decay schedule style for WSD when used.
lr_wsd_decay_samples	null	WSD decay window in samples.
lr_wsd_decay_iters	null	WSD decay window in iterations.
head_lr_mult	1.0	LR multiplier for attention/head modules when supported.
weight_decay	0.01 (0.0 in pre_trainer.yaml)	AdamW / L2-style weight decay.
start_weight_decay	null	Starting weight decay for schedules.
end_weight_decay	null	Ending weight decay for schedules.
weight_decay_incr_style	constant	How weight decay changes between start/end.
clip_grad	1.0	Global gradient norm clip.
adam_beta1	0.9	Adam first moment decay.
adam_beta2	0.95 (0.999 in pre_trainer.yaml)	Adam second moment decay.
adam_eps	1.0e-08	Adam epsilon.
sgd_momentum	0.9	SGD momentum when optimizer is SGD.
override_opt_param_scheduler	false (true in pre_trainer.yaml)	Override optimizer parameter groups’ schedulers.
use_checkpoint_opt_param_scheduler	false	Load optimizer scheduler state strictly from checkpoint.
warmup	null	Alternate warmup specification (legacy / schedule hooks).
decoupled_lr	null	Decoupled LR for certain param groups.
decoupled_min_lr	null	Minimum for decoupled LR.
muon_extra_scale_factor	1.0	Muon optimizer scaling.
muon_scale_mode	"spectral"	Muon scaling mode.
muon_fp32_matmul_prec	"medium"	Muon matmul precision hint.
muon_num_ns_steps	5	Muon Newton–Schulz iterations.
muon_tp_mode	"blockwise"	Muon tensor-parallel mode.
muon_use_nesterov	false	Muon Nesterov momentum.
muon_split_qkv	true	Split QKV for Muon.
muon_momentum	0.95	Muon momentum.
muon_weight_decay	0.01	Muon-specific decay.
muon_weight_decay_method	"decoupled"	How Muon applies decay.
optimizer_cpu_offload	false	Offload optimizer state to CPU.
optimizer_offload_fraction	1.0	Fraction of optimizer state offloaded.
use_torch_optimizer_for_cpu_offload	false	Use PyTorch optimizer for offload path.
overlap_cpu_optimizer_d2h_h2d	false	Overlap CPU optimizer device transfers.
pin_cpu_grads	true	Pin memory for CPU gradients.
pin_cpu_params	true	Pin memory for CPU params in offload.
use_precision_aware_optimizer	false	Use precision-aware optimizer (main grads/params in lower precision).
main_grads_dtype	fp32	Dtype for main gradients (fp32, bf16).
main_params_dtype	fp32	Dtype for master params.
exp_avg_dtype	fp32	Optimizer first moment dtype (fp32, fp16, fp8).
exp_avg_sq_dtype	fp32	Optimizer second moment dtype.

---

5. Parallelism and distribution

Sources: trainer_base.yaml (distributed runtime) and primus/configs/models/megatron/language_model.yaml (model-parallel sizes and TP communication).

5.1 Data / distributed runtime (trainer)

Parameter	Default	Description
overlap_p2p_comm	true	Overlap pipeline P2P with compute.
distributed_backend	nccl	Process-group backend (nccl, gloo, …).
distributed_timeout_minutes	10 (60 in pre_trainer.yaml)	Collective timeout.
defer_embedding_wgrad_compute	false	Defer embedding weight gradients.
wgrad_deferral_limit	0	Max deferred embedding wgrad steps.
align_grad_reduce	true	Align gradient reductions for efficiency.
ddp_num_buckets	null	Number of DDP buckets.
ddp_bucket_size	null	DDP bucket size in elements.
ddp_pad_buckets_for_high_nccl_busbw	false	Pad buckets for NCCL bus bandwidth.
ddp_average_in_collective	false	Average inside collective vs outside.
overlap_grad_reduce	false	Overlap gradient all-reduce with backward.
overlap_param_gather	false	Overlap param all-gather (distributed optimizer).
overlap_param_gather_with_optimizer_step	false	Overlap param gather with optimizer step.
align_param_gather	true	Align param gather for distributed optimizer.
scatter_gather_tensors_in_pipeline	true	Scatter/gather tensors across PP ranks.
use_ring_exchange_p2p	false	Ring-exchange P2P for PP.
local_rank	null	Local rank override (normally from launcher).
lazy_mpu_init	null	Defer Megatron parallel state init.
account_for_embedding_in_pipeline_split	false	Account for embedding in PP partition.
account_for_loss_in_pipeline_split	false	Account for loss partition in PP.
empty_unused_memory_level	0	Aggressiveness of torch.cuda.empty_cache.
standalone_embedding_stage	false	Dedicated PP stage for embeddings.
use_distributed_optimizer	false (true in pre_trainer.yaml)	Shard optimizer state across data parallel.
use_sharp	false	Use SHARP for collectives when available.
sharp_enabled_group	null	Which group SHARP applies to (dp, dp_replica).
use_custom_fsdp	false	Custom FSDP integration path.
use_megatron_fsdp	false	Megatron FSDP path.
init_model_with_meta_device	false	Build model on meta device first.
data_parallel_sharding_strategy	no_shard	FSDP / ZeRO style sharding (no_shard, optim, …).
gradient_reduce_div_fusion	true	Fuse division into reduce-scatter.
suggested_communication_unit_size	400000000	Suggested communication chunk size.
keep_fp8_transpose_cache_when_using_custom_fsdp	false	Keep FP8 transpose cache with custom FSDP.
num_distributed_optimizer_instances	1	Sharded optimizer instances per rank group.
use_torch_fsdp2	false	Use PyTorch FSDP2 integration.
nccl_communicator_config_path	null	JSON config for NCCL communicators.
use_tp_pp_dp_mapping	false	Custom TP/PP/DP process mapping.
replication	false	Data replication mode for certain schedules.
replication_jump	null	Stride between replicated ranks.
replication_factor	null	Replication factor.
deterministic_mode	false	Prefer deterministic algorithms (slower).
check_weight_hash_across_dp_replicas_interval	null	Periodically hash weights across DP replicas for debugging.
overlap_moe_expert_parallel_comm	false	Overlap MoE expert-parallel communication.
decoder_pipeline_manual_split_list	null	Primus: manual PP split points for decoder (list of ints).
patch_moe_overlap	false	Primus: patch MoE compute/comm overlap.

5.2 Model parallelism (model preset)

Parameter	Default	Description
model_parallel_size	null	Legacy combined MP size override.
tensor_model_parallel_size	1	Tensor parallelism degree (intra-layer split).
encoder_tensor_model_parallel_size	0	Encoder TP size when encoder/decoder differ.
pipeline_model_parallel_size	1	Pipeline parallelism stages.
pipeline_model_parallel_layout	null	Optional explicit PP layout string.
pipeline_model_parallel_comm_backend	null	nccl or ucc for PP collectives.
encoder_pipeline_model_parallel_size	0	Encoder PP stages (encoder–decoder models).
pipeline_model_parallel_split_rank	null	Rank where encoder/decoder split.
decoder_first_pipeline_num_layers	null	Layers on first decoder PP stage.
decoder_last_pipeline_num_layers	null	Layers on last decoder PP stage.
virtual_pipeline_model_parallel_size	null	Virtual PP (interleaved) depth.
num_layers_per_virtual_pipeline_stage	null	Layers per virtual stage.
num_virtual_stages_per_pipeline_rank	null	Virtual stages per physical PP rank.
microbatch_group_size_per_vp_stage	null	Microbatch grouping for interleaved PP.
sequence_parallel	true	Sequence parallelism when TP > 1.
context_parallel_size	1	Context (sequence) parallelism degree.
cp_comm_type	p2p	Context-parallel comm pattern (p2p, a2a, allgather, a2a+p2p).
hierarchical_context_parallel_sizes	null	Hierarchical CP group sizes.
expert_model_parallel_size	1	Expert parallelism for MoE.
expert_tensor_parallel_size	null	Expert tensor-parallel degree.
high_priority_stream_groups	[]	Named groups that get high-priority CUDA streams.

5.3 Tensor-parallel communication overlap (model)

Parameter	Default	Description
async_tensor_model_parallel_allreduce	true	Async TP all-reduces for column-parallel layers.
tp_comm_overlap	false	Enable TP communication overlap planner.
tp_comm_overlap_cfg	null	Extra JSON / path for overlap configuration.
tp_comm_overlap_ag	true	Overlap all-gather in TP backward.
tp_comm_overlap_rs	true	Overlap reduce-scatter in TP backward.
tp_comm_overlap_rs_dgrad	false	Overlap RS for data-grad path.
tp_comm_split_ag	true	Split all-gather for overlap.
tp_comm_split_rs	true	Split reduce-scatter for overlap.
tp_comm_bulk_wgrad	true	Bulk weight-gradient path for TP comm.
tp_comm_bulk_dgrad	true	Bulk data-gradient path for TP comm.
barrier_with_L1_time	true	Barrier using L1 timing hooks for TP comm profiling.
tp_comm_bootstrap_backend	nccl	Backend used to bootstrap TP communicators.

---

6. Checkpointing

Source: trainer_base.yaml.

Parameter	Default	Description
save	null	Path prefix / pattern for checkpoints to write.
save_interval	20000 (1000 in pre_trainer.yaml)	Save every N iterations.
save_retain_interval	null	Retain checkpoints at this interval.
no_save_optim	null	Skip optimizer state in checkpoints when truthy.
no_save_rng	null	Skip RNG state in checkpoints when truthy.
load	null	Checkpoint path to load.
load_main_params_from_ckpt	false	Load only main parameters.
no_load_optim	null	Skip loading optimizer state.
no_load_rng	null	Skip loading RNG state.
finetune	false (true in pre_trainer.yaml)	Finetune mode (do not require full optimizer match).
use_checkpoint_args	false	When true, restore training args from checkpoint metadata.
use_mp_args_from_checkpoint_args	false	Restore model-parallel args from checkpoint.
use_tokenizer_model_from_checkpoint_args	true	Restore tokenizer path from checkpoint args.
exit_on_missing_checkpoint	true	Fail if load is set but checkpoint is missing.
non_persistent_save_interval	null	Ephemeral checkpoint interval.
non_persistent_ckpt_type	null	global, local, in_memory, or null.
non_persistent_global_ckpt_dir	null	Directory for non-persistent global checkpoints.
non_persistent_local_ckpt_dir	null	Directory for non-persistent local checkpoints.
non_persistent_local_ckpt_algo	"fully_parallel"	fully_parallel or atomic.
pretrained_checkpoint	null	Load weights from a pretrained checkpoint path.
ckpt_step	null	Specific step to load within a distributed checkpoint.
use_dist_ckpt_deprecated	false	Use deprecated distributed checkpoint format.
use_persistent_ckpt_worker	false	Background worker for checkpoint IO.
auto_detect_ckpt_format	false	Infer checkpoint format automatically.
dist_ckpt_format_deprecated	null	Legacy format hint.
ckpt_format	torch_dist	torch, torch_dist, or zarr.
ckpt_convert_format	null	Target format for one-shot conversion.
ckpt_convert_save	null	Output path for conversion.
ckpt_convert_update_legacy_dist_opt_format	false	Update legacy distributed-optimizer layout when converting.
ckpt_fully_parallel_save_deprecated	false	Deprecated fully-parallel save toggle.
ckpt_fully_parallel_save	true	Save shards in parallel across ranks.
async_save	null	Async checkpoint save (null = framework default).
ckpt_fully_parallel_load	false	Load shards in parallel.
ckpt_assume_constant_structure	false	Assume identical layer structure across ranks.
dist_ckpt_strictness	assume_ok_unexpected	How to handle unexpected keys in distributed ckpt.
dist_ckpt_save_pre_mcore_014	null	Compatibility flag for older Megatron-Core checkpoints.
dist_ckpt_optim_fully_reshardable	null	Optimizer state fully reshardable layout.
auto_continue_train	false	Primus: resume from latest checkpoint in the save directory when enabled.
disable_last_saving	false	Primus: skip writing the final checkpoint at shutdown.

---

7. Data

Source: trainer_base.yaml.

Parameter	Default	Description
data_path	null	Single blended dataset path / list.
data_sharding	true	Shard data across ranks.
split	"99,1,0" (null in pre_trainer.yaml)	Train/valid/test split ratios as comma string.
train_data_path	null	Training data blend.
valid_data_path	null	Validation data blend.
test_data_path	null	Test data blend.
data_args_path	null	External JSON/YAML of dataset arguments.
per_split_data_args_path	null	Per-split dataset args file.
data_cache_path	null	On-disk cache for indexed datasets.
mock_data	false	Use synthetic data (no real files).
merge_file	null	Merge file for blended datasets.
seq_length	4096 (1024 in pre_trainer.yaml)	Training sequence length.
encoder_seq_length	null	Encoder sequence length (encoder–decoder).
decoder_seq_length	null	Decoder sequence length.
retriever_seq_length	256	Sequence length for retriever models.
sample_rate	1.0	Sampling rate for dataset blending.
mask_prob	0.15	MLM mask probability.
short_seq_prob	0.1	Probability of shorter sequences in BERT-style data.
num_workers	8	DataLoader worker processes per rank.
reset_position_ids	false	Reset position IDs at document boundaries.
reset_attention_mask	false	Reset attention mask at boundaries.
eod_mask_loss	false	Mask loss at end-of-document tokens.
dataloader_type	null (cyclic in pre_trainer.yaml)	Dataloader implementation (single, cyclic, external, …).
mmap_bin_files	true	Memory-map .bin index files when supported.
create_attention_mask_in_dataloader	true	Build attention masks in the dataloader.
num_dataset_builder_threads	1	Threads to build dataset indices.

---

8. Recomputation (activation checkpointing)

Sources: trainer_base.yaml and primus_megatron_module.yaml.

Parameter	Default	Description
recompute_activations	false	Enable activation recomputation globally.
recompute_granularity	null	full or selective checkpointing.
recompute_method	null	uniform or block selective recomputation.
recompute_num_layers	null	Layers to recompute per block / schedule.
recompute_layer_ids	null	Primus: explicit global layer indices to recompute (0 … num_layers-1).
distribute_saved_activations	false	Distribute saved activations across TP/PP for memory balance.
checkpoint_activations	false	Deprecated alias for activation checkpointing.
moe_layer_recompute	false	Recompute MoE layer activations (model preset).

---

9. Logging and profiling

Sources: trainer_base.yaml and primus_megatron_module.yaml.

9.1 Logging

Parameter	Default	Description
log_avg_skip_iterations	2	Skip first N iterations for throughput averaging.
log_avg_reset_interval	10	Reset moving averages periodically.
log_params_norm	false	Log L2 norms of parameters.
log_num_zeros_in_grad	false	Log fraction of zero gradients.
log_throughput	false (true in pre_trainer.yaml)	Log tokens/sec and timing.
log_progress	false	Verbose progress logging.
timing_log_level	0	Verbosity for timing logs.
timing_log_option	minmax	Aggregate style for timing (minmax, all, …).
tensorboard_log_interval	1	Steps between TensorBoard scalars.
tensorboard_queue_size	1000	TensorBoard event queue size.
log_timers_to_tensorboard	false (true in pre_trainer.yaml)	Write timer stats to TensorBoard.
log_batch_size_to_tensorboard	false (true in pre_trainer.yaml)	Log batch size.
log_learning_rate_to_tensorboard	true	Log LR.
log_validation_ppl_to_tensorboard	false	Log validation perplexity.
log_memory_to_tensorboard	false	Log memory usage.
log_world_size_to_tensorboard	false	Log distributed world size.
log_loss_scale_to_tensorboard	true	Log FP16/FP8 loss scale.
wandb_project	null	Weights & Biases project name.
wandb_exp_name	null	W&B run name.
wandb_save_dir	null	W&B local directory.
wandb_entity	null	W&B entity / team.
enable_one_logger	true	Enable NVIDIA OneLogger integration.
one_logger_project	megatron-lm	OneLogger project string.
one_logger_run_name	null	OneLogger run name.
log_interval	100 (1 in pre_trainer.yaml)	Console log interval in iterations.
tensorboard_dir	null	TensorBoard output directory.
logging_level	null	Python logging level override.
config_logger_dir	""	Directory for dumped config logs.
one_logger_async	false	Async OneLogger flushing.
app_tag_run_name	null	Application tag for telemetry.
app_tag_run_version	0.0.0	Application tag version.
disable_tensorboard	true	Primus: disable TensorBoard integration in Primus-wrapped runs.
disable_wandb	true	Primus: disable W&B.
disable_mlflow	true	Primus: disable MLflow.
mlflow_run_name	null	Primus: MLflow run name.
mlflow_experiment_name	null	Primus: MLflow experiment name.
use_rocm_mem_info	false	Primus: collect ROCm memory info via rocm-smi every step when true.
use_rocm_mem_info_iters	[1, 2]	Primus: iterations at which to log memory if use_rocm_mem_info is false.

9.2 Profiling

Parameter	Default	Description
profile	false	Enable lightweight Nsight / CUDA profiling hooks.
use_pytorch_profiler	false	Enable torch.profiler regions.
profile_ranks	[0]	Ranks to profile.
profile_step_start	10	First step to profile.
profile_step_end	12	Last step to profile.
iterations_to_skip	null	Skip listed iterations in profiling.
result_rejected_tracker_filename	null	Log rejected samples to this file.
enable_gloo_process_groups	true	Create auxiliary Gloo groups for CPU-side ops.
record_memory_history	false	Record CUDA memory history (debug).
memory_snapshot_path	snapshot.pickle	Path for memory snapshot dumps.
disable_profiler_activity_cpu	false	Primus: omit CPU activities from profiler traces.
torch_profiler_record_shapes	true	Primus: record tensor shapes in PyTorch profiler.
torch_profiler_with_stack	true	Primus: capture Python stacks in profiler.
torch_profiler_use_gzip	false	Primus: gzip profiler outputs.

---

10. Model architecture

Sources: primus/configs/models/megatron/language_model.yaml and primus/configs/models/megatron/primus_megatron_model.yaml.

10.1 Core architecture

Parameter	Default	Description
use_legacy_models	false	Use legacy Megatron model code paths.
deprecated_use_mcore_models	false	Deprecated flag for Megatron-Core models; prefer current transformer_impl + stack.
model_type	gpt	gpt or mamba family.
num_layers	24	Transformer layers (decoder or unified stack).
encoder_num_layers	null	Encoder depth (encoder–decoder).
decoder_num_layers	null	Decoder depth.
hidden_size	1024	Hidden / model width.
num_attention_heads	16	Attention heads.
attention_backend	auto	Attention kernel backend selection.
group_query_attention	false	Enable grouped-query attention (GQA).
qk_layernorm	false	LayerNorm on Q/K projections.
qk_l2_norm	false	L2-normalize Q/K vectors.
num_query_groups	null	Number of query groups for GQA; null means MHA.
add_position_embedding	false	Add absolute position embeddings (non-RoPE stacks).
position_embedding_type	learned_absolute	Position embedding style.
max_position_embeddings	null	Maximum sequence positions (context length cap).
original_max_position_embeddings	null	Original pretrained length for interpolation / scaling.
untie_embeddings_and_output_weights	true	Separate input embedding and LM head weights.
ffn_hidden_size	null	FFN hidden size; null often defaults via hidden_size heuristics.
kv_channels	null	Per-head KV channels override.
hidden_dropout	0.1	Dropout on residual / hidden states.
attention_dropout	0.1	Attention dropout.
fp32_residual_connection	false	Accumulate residuals in FP32.
apply_residual_connection_post_layernorm	false	Apply residual after (vs before) norm where supported.
add_bias_linear	false	Biases in linear / column-parallel layers.
add_qkv_bias	false	Biases in QKV projections.
swiglu	true	SwiGLU activation in FFN.
quick_geglu	false	Faster GeGLU path.
openai_gelu	false	OpenAI GELU variant.
squared_relu	false	Squared ReLU activation.
rotary_base	10000	RoPE base frequency.
rotary_percent	1.0	Fraction of head dim spanned by RoPE.
rotary_interleaved	false	Interleaved RoPE layout.
rotary_seq_len_interpolation_factor	null	Positional interpolation factor for long contexts.
use_rotary_position_embeddings	null	Force RoPE on/off; null follows model type.
use_rope_scaling	false	Enable LLaMA-style rope scaling.
rope_scaling_factor	8.0	Scaling factor for extended contexts (LLaMA-3 style).
transformer_impl	transformer_engine	Backend library (transformer_engine, local, …).
rope_type	null	rope or yarn style extensions.
norm_epsilon	1.0e-05	LayerNorm / RMSNorm epsilon.
normalization	"LayerNorm"	Norm type (LayerNorm, RMSNorm with TE, …).
apply_layernorm_1p	false	LayerNorm with +1 offset trick.
clone_scatter_output_in_embedding	true	Clone embedding scatter for autograd safety.
perform_initialization	true	Run weight initialization.
use_cpu_initialization	null	Initialize on CPU then move to GPU.
use_te_activation_func	false	Use Transformer Engine activation kernels.
gradient_accumulation_fusion	true	Fuse gradient accumulation kernels.
delay_wgrad_compute	false	Delay weight-gradient computation for scheduling.

10.2 Tokenizer and vocabulary

Parameter	Default	Description
tokenizer_type	null	Tokenizer class name (GPT2BPETokenizer, HuggingFaceTokenizer, …).
tokenizer_model	null	Path to tokenizer model / vocabulary file.
vocab_size	null	Vocabulary size (often inferred from tokenizer).
vocab_file	null	Vocabulary file path for BPE/WP tokenizers.
vocab_extra_ids	0	Extra reserved token slots.
tiktoken_pattern	null	Regex pattern for tiktoken.
tiktoken_num_special_tokens	1000	Special token count for tiktoken setup.
tiktoken_special_tokens	null	Serialized special tokens for tiktoken.
legacy_tokenizer	false	Legacy tokenizer behavior.
trust_remote_code	false	trust_remote_code for Hugging Face tokenizers.

10.3 Initialization and attention numerics

Parameter	Default	Description
init_method_std	0.02	Standard deviation for weight init.
apply_query_key_layer_scaling	false	Scale Q/K by layer index (deprecated GPT-3 trick).
attention_softmax_in_fp32	false	Force softmax in FP32.

10.4 Kernel fusion flags

Parameter	Default	Description
bias_gelu_fusion	true	Fuse bias + GELU.
cross_entropy_loss_fusion	false	Fused cross-entropy + softmax.
cross_entropy_fusion_impl	"native"	native or te fused CE.
bias_swiglu_fusion	true	Fuse bias + SwiGLU.
masked_softmax_fusion	true	Fused masked softmax.
no_persist_layer_norm	false	Non-persistent LayerNorm mode in TE.
bias_dropout_fusion	true	Fuse bias + dropout.
apply_rope_fusion	true	Fused RoPE kernels.

10.5 Multi-latent attention (MLA)

Parameter	Default	Description
multi_latent_attention	false	Enable MLA blocks instead of standard MHA.
q_lora_rank	null	Low-rank query projection rank.
kv_lora_rank	32	Low-rank KV compression rank.
qk_head_dim	128	Q/K head dimension for MLA.
qk_pos_emb_head_dim	64	Positional head dimension for MLA.
v_head_dim	128	Value head dimension for MLA.
rotary_scaling_factor	1.0	RoPE scaling inside MLA (distinct from rope_scaling_factor above).
mscale	1.0	Yarn / scaling m-factor.
mscale_all_dim	1.0	Yarn scaling on all dims.

10.6 Mixture-of-experts (MoE)

Parameter	Default	Description
num_experts	null	Experts per MoE layer; null means dense model.
moe_layer_freq	1	Every Nth layer is MoE (1 = every layer).
moe_ffn_hidden_size	null	Expert FFN hidden size.
moe_shared_expert_overlap	false	Shared expert overlaps routing.
moe_shared_expert_intermediate_size	null	Shared expert FFN size.
moe_grouped_gemm	false	Grouped GEMM for experts.
moe_router_load_balancing_type	"aux_loss"	Router balancing (aux_loss, seq_aux_loss, sinkhorn, none).
moe_router_dtype	null	Router activation dtype (fp32, fp64).
moe_router_score_function	softmax	softmax or sigmoid routing scores.
moe_router_topk	2	Experts to select per token.
moe_router_pre_softmax	false	Apply softmax before top-k.
moe_router_num_groups	null	Group-limited routing: number of expert groups.
moe_router_group_topk	null	Groups to pick before top-k inside groups.
moe_router_topk_scaling_factor	null	Scaling for routing logits.
moe_router_enable_expert_bias	false	Learnable per-expert bias.
moe_router_bias_update_rate	1.0e-03	Update rate for expert bias.
moe_use_legacy_grouped_gemm	false	Legacy grouped GEMM path.
moe_aux_loss_coeff	0.0	Auxiliary load-balancing loss weight.
moe_z_loss_coeff	null	Router z-loss coefficient.
moe_input_jitter_eps	null	Input jitter for router stability.
moe_token_dispatcher_type	allgather	Token dispatch algorithm (allgather, alltoall, flex, alltoall_seq).
moe_enable_deepep	false	DeepEP-style expert parallelism.
moe_per_layer_logging	false	Per-layer MoE statistics logging.
moe_expert_capacity_factor	null	Capacity factor for token dropping / padding.
moe_pad_expert_input_to_capacity	false	Pad expert batches to capacity.
moe_token_drop_policy	probs	Token dropping policy when over capacity.
moe_extended_tp	false	Extended tensor-parallel for experts.
moe_use_upcycling	false	Expert upcycling initialization.
moe_permute_fusion	false	Fuse token permutation for MoE.
disable_primus_topk_router	false	Primus: disable Primus top-k router patch.
moe_router_force_load_balancing	false	Primus: force load-balanced routing.
use_deprecated_20241209_moe_layer	false	Primus: legacy MoE layer implementation.
moe_router_force_load_balancing_type	even	Primus: Control the force load balancing type for the MoE router. Choices: even, uniform.

10.7 Logit softcapping (Primus / Grok-style)

Parameter	Default	Description
final_logit_softcapping	null	Softcap value for final logits; null disables.
attn_logit_softcapping	null	Softcap for attention logits.
router_logit_softcapping	null	Softcap for MoE router logits.

---

11. Primus extensions

11.1 Build and compile

Parameter	Default	Description
disable_compile_dependencies	true	Primus: avoid compiling dependency stacks in the trainer wrapper.

11.2 Primus-Turbo (primus_turbo.yaml)

Parameter	Default	Description
enable_primus_turbo	false	Master switch for Primus-Turbo integrations. Many sub-features require this plus specific kernels.
use_turbo_attention	false	Turbo attention implementation.
use_sink_attention	false	GPT-OSS-style learned sink attention.
sink_sliding_window	0	Sliding-window size for sink attention (GPT-OSS uses 128).
sink_window_even_layers_only	true	Apply the sliding window only to even layers (GPT-OSS pattern).
use_turbo_gemm	false	Active Turbo GEMM flag for Dense paths.
use_turbo_parallel_linear	(removed)	Removed—use use_turbo_gemm. Passing this key now raises an assertion error (use_turbo_parallel_linear has been removed; please use use_turbo_gemm instead).
use_turbo_grouped_gemm	false	Active Turbo grouped GEMM flag for MoE paths.
use_turbo_grouped_mlp	(removed)	Removed—use use_turbo_grouped_gemm. Passing this key now raises an assertion error (use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead).
moe_use_fused_router_with_aux_score	false	Fused MoE router with auxiliary scores.
enable_turbo_attention_float8	false	FP8 path inside Turbo attention (spacing in YAML is normalized to this key).
use_turbo_deepep	false	Turbo DeepEP expert communication.
turbo_deepep_num_cu	32	DeepEP compute units / channels.
turbo_deepep_use_comm_stream	false	Use a dedicated communication stream for DeepEP.
turbo_sync_free_moe_stage	0	Stage selector for sync-free MoE.
use_turbo_fused_act_with_probs	false	Fuse activation + probability tensors to remove redundant work.
use_turbo_rms_norm	false	Turbo RMSNorm kernels.

11.3 Zero-bubble pipeline (zero_bubble.yaml)

Parameter	Default	Description
patch_zero_bubble	false	Install Primus zero-bubble PP patches when true.
debug_scheduler_table	false	Print PP scheduler tables (also in primus_pipeline.yaml; last merge wins—defaults match).
enable_zb_runtime	true	Unified runtime for zero-bubble and related schedules.
pre_communication_optimization	false	Issue a tiny comm before real comm to tune overlap.
zero_bubble_pipeline_timers_start_iter	100	Start iter for auto-scheduler timers.
zero_bubble_pipeline_timers_end_iter	110	End iter for auto-scheduler timers.
zero_bubble_max_pending_backward	auto	Max pending backward ops (ZB1p vs ZB2p style); auto adapts.
zero_bubble_adaptive_memory_limit_percentile	85	GPU memory percentile cap for adaptive ZB.
enable_optimizer_post_validation	false	Post-optimizer validation step (needs FSDP path).
enable_exactly_numeric_match	true	Require bitwise match in post validation when enabled.
enable_zero_bubble	true	Enable zero-bubble schedule features in the ZB runtime.
zero_bubble_v_schedule	false	Zero-bubble “V” schedule without extra memory vs some baselines.
zero_bubble_v_schedule_mem_setup	half	Memory setup variant: half, min, or zb.
enable_1f1b_v	false	1F1B-V schedule variant.
allow_padding_num_layers	true	Allow PP layer padding for divisibility.
profile_memory_iter	-1	Iteration to profile memory (-1 disables).
interleave_group_size	0	Interleaved PP group size.
offload_chunk_num	0	Activation offload chunk count.
offload_time	1.0	Time budget for offload (scheduler hint).
auto_offload_time	true	Auto-tune offload timing.
offload_overlap_sr	true	Overlap save/resume in offload path.
num_seq_splits	1	Splits along sequence dimension for ZB.
cpu_offload	false	CPU offload of activations in ZB path.

11.4 Primus pipeline (primus_pipeline.yaml)

Parameter	Default	Description
patch_primus_pipeline	false	Enable Primus pipeline scheduling patches.
pp_algorithm	"1f1b-interleaved"	Schedule name (1f1b, 1f1b-interleaved, zero-bubble, zero-bubble-heuristic, zbv-formatted, v-half, v-min).
communication_method	"async_p2p"	async_p2p or batch_p2p PP transfers.
offload	false	Generic PP activation offload toggle in Primus pipeline.
offload_ops	""	Comma-separated offload targets (attn today; other ops listed in-file are not supported yet).
pp_max_mem	null	zero-bubble-heuristic only: max activation memory per stage (null = unlimited).
pp_cost_f	null	zero-bubble-heuristic only: forward cost per stage (scalar or list; null = default 1000).
pp_cost_b	null	zero-bubble-heuristic only: backward cost per stage (scalar or list; null = default 1000).
pp_cost_w	null	zero-bubble-heuristic only: weight-grad cost per stage (scalar or list; null = default 1000).

pp_warmup and dump_pp_data are Primus helpers defined in primus_megatron_module.yaml (not primus_pipeline.yaml):

Parameter	Default	Description
pp_warmup	false	Primus: warm-up PP stages to reduce first-iteration latency.
dump_pp_data	false	Primus: dump PP tensors for debugging.

---

12. Reinforcement learning and GRPO-related settings

Source: trainer_base.yaml. Names follow Megatron’s grpo_* / rl_* prefixes (there is no rl_grpo single flag in these presets).

Parameter	Default	Description
perform_rl_step	false	Run RL / preference optimization steps (GRPO / LangRL integration).
rl_prompts_per_eval	32	Prompts per RL evaluation pass.
grpo_prompts_per_step	32	GRPO prompts sampled per training step.
grpo_group_size	2	Samples per prompt group for GRPO.
grpo_iterations	2	Inner GRPO iterations.
grpo_clamp_eps_lower	0.01	PPO-style lower clip epsilon.
grpo_clamp_eps_upper	0.01	Upper clip epsilon.
grpo_kl_beta	0.001	KL penalty weight toward reference policy.
grpo_entropy_term_weight	0.0	Entropy bonus weight.
grpo_filter_groups_with_same_reward	false	Drop groups with identical rewards.
grpo_default_temperature	1.0	Default softmax temperature for rollouts.
grpo_default_top_p	0	Top-p sampling (0 often means disabled / greedy—see Megatron RL docs).
langrl_inference_server_type	inplace_megatron	LangRL inference backend.
langrl_inference_server_conversation_template	null	Conversation template path / name.
langrl_env_config	null	Environment / task YAML for LangRL.
rl_offload_optimizer_during_inference	false	Offload optimizer to CPU during rollout inference.
rl_offload_kv_cache_during_training	false	Offload KV cache while training forward runs.
rl_remove_kv_cache_during_training	false	Drop KV cache between RL phases to save memory.
rl_reset_cuda_graphs	false	Reset CUDA graphs when switching RL phases.
rl_partial_rollouts	false	Partial sequence rollouts.
rl_inference_logprobs_is_correction	false	Interpret inference logprobs as IS correction term.
rl_importance_sampling_truncation_coef	null	Truncate importance ratios at this value.
rl_calculate_intra_group_similarity	false	Log similarity within GRPO groups.

---

13. Additional specialized parameters

Source: trainer_base.yaml (remaining domains).

13.1 Vision pretraining

Parameter	Default	Description
vision_pretraining	false	Enable vision backbone pretraining.
vision_pretraining_type	classify	Objective (classify, etc.).
vision_backbone_type	vit	Vision backbone family.
swin_backbone_type	tiny	Swin variant size.
num_classes	1000	Classification classes.
img_h	224	Image height.
img_w	224	Image width.
num_channels	3	Input channels.
patch_dim	16	ViT patch size.
classes_fraction	1.0	Fraction of classes used.
data_per_class_fraction	1.0	Fraction of data per class.

13.2 RETRO

Parameter	Default	Description
retro_project_dir	null	RETRO project directory with indices.
retro_add_retriever	false	Add frozen retriever tower.
retro_cyclic_train_iters	null	Cyclic iterator length.
retro_encoder_layers	2	Retriever encoder layers.
retro_encoder_hidden_dropout	0.1	Retriever dropout.
retro_encoder_attention_dropout	0.1	Retriever attention dropout.
retro_num_neighbors	2	Neighbors per query chunk.
retro_num_retrieved_chunks	2	Chunks concatenated per neighbor set.
retro_attention_gate	1	Gating between retrieval and LM.
retro_verify_neighbor_count	true	Assert neighbor counts for debugging.

13.3 DINO self-supervised

Parameter	Default	Description
dino_local_img_size	96	Local crop size.
dino_local_crops_number	10	Number of local crops.
dino_head_hidden_size	2048	Projection head width.
dino_bottleneck_size	256	Bottleneck dimension.
dino_freeze_last_layer	1	Freeze last layer epochs.
dino_norm_last_layer	false	Normalize last layer weights.
dino_warmup_teacher_temp	0.04	Teacher temperature warmup start.
dino_teacher_temp	0.07	Teacher temperature.
dino_warmup_teacher_temp_epochs	30	Epochs to warm teacher temperature.

13.4 Biencoder / ICT / retriever utilities

Parameter	Default	Description
ict_head_size	null	ICT projection head width.
biencoder_projection_dim	0	Biencoder shared projection dimension.
biencoder_shared_query_context_model	false	Share query/context encoders.
ict_load	null	ICT checkpoint path.
bert_load	null	BERT encoder checkpoint for biencoder.
titles_data_path	null	Titles file for ICT datasets.
query_in_block_prob	0.1	Probability of in-block queries.
use_one_sent_docs	false	Single-sentence pseudo documents.
evidence_data_path	null	Evidence passages for open-domain QA.
retriever_report_topk_accuracies	[]	k values for top-k accuracy logging.
retriever_score_scaling	false	Scale retriever scores.
block_data_path	null	Block JSON data for retrieval.
embedding_path	null	Precomputed embeddings path.
indexer_batch_size	128	Batch size when building ANN index.
indexer_log_interval	1000	Indexer progress log interval.

13.5 Straggler detection

Parameter	Default	Description
log_straggler	false	Log straggler diagnostics.
disable_straggler_on_startup	false	Skip straggler detection at startup.
straggler_ctrlr_port	65535	Controller port for straggler service.
straggler_minmax_count	1	Min/max samples for straggler stats.

13.6 Inference-oriented options

Parameter	Default	Description
inference_batch_times_seqlen_threshold	-1	Heuristic threshold tying batch and sequence length.
inference_dynamic_batching	false	Dynamic batching for inference server.
inference_dynamic_batching_buffer_size_gb	40.0	GPU buffer budget (GB).
inference_dynamic_batching_buffer_guaranteed_fraction	0.2	Minimum reserved fraction of buffer.
inference_dynamic_batching_buffer_overflow_factor	null	Overflow growth factor.
inference_dynamic_batching_max_requests_override	null	Hard cap on concurrent requests.
inference_dynamic_batching_max_tokens_override	null	Hard cap on tokens in flight.
max_tokens_to_oom	12000	Token limit guard before OOM abort.
output_bert_embeddings	false	Return BERT pooled embeddings.
bert_embedder_type	megatron	megatron or huggingface embedder.
flash_decode	false	Flash decode kernels for incremental generation.
enable_cuda_graph	false	Capture CUDA graphs for inference.
cuda_graph_warmup_steps	3	Warm-up steps before capturing graphs.
external_cuda_graph	false	External graph provider hooks.
cuda_graph_scope	full	Graph scope (full or attn).
inference_max_requests	8	Max concurrent requests.
inference_max_seq_length	2560	Max prefill + decode tokens per request.

13.7 Fault tolerance package and tooling

Parameter	Default	Description
enable_ft_package	false	NVIDIA fault-tolerance package hooks.
calc_ft_timeouts	false	Auto-calculate FT timeouts.
run_workload_inspector_server	false	Run workload inspector sidecar.

13.8 Heterogeneous layers and process resilience

Parameter	Default	Description
heterogeneous_layers_config_path	null	JSON describing variable layer widths/types per layer.
heterogeneous_layers_config_encoded_json	null	Inline base64/JSON blob for heterogeneous layers.
inprocess_restart	false	In-process restart for fault recovery experiments.

13.9 Experimental and rerun controls

Parameter	Default	Description
enable_experimental	false	Gate experimental Megatron features.
error_injection_rate	0	Fraction of iterations with injected errors (testing).
error_injection_type	transient_error	correct_result, transient_error, or persistent_error.
rerun_mode	disabled	disabled, validate_results, or report_stats for rerun harness.

---

Related documentation

• Megatron-LM argument definitions: megatron/training/arguments.py

• Primus Megatron presets: primus/configs/modules/megatron/

• Primus Megatron model presets: primus/configs/models/megatron/