extends:
  - ../module_base.yaml
  - primus_megatron_module.yaml
  - primus_turbo.yaml
  - torch_compile.yaml
  - zero_bubble.yaml
  - primus_pipeline.yaml

trainable: true

# megatron trainer args:
# training, optimzer, checkpoint, loss, distributed, recompute, data, profile, logging

# training
yaml_cfg: null # not support
spec: null
micro_batch_size: 2
batch_size: null # deprecated
global_batch_size: 128
rampup_batch_size: null
decrease_batch_size_if_needed: false
check_for_nan_in_loss_and_grad: true
check_for_spiky_loss: false
check_for_large_grads: false
make_vocab_size_divisible_by: 128
exit_signal_handler: false
exit_duration_in_mins: null
exit_interval: null
onnx_safe: null
bert_binary_head: true

use_flash_attn: false
seed: 1234
data_parallel_random_init: false
init_method_xavier_uniform: false
test_mode: false

# mixed precision
fp16: false
bf16: true
grad_reduce_in_bf16: false
calculate_per_token_loss: false
loss_scale: null
initial_loss_scale: 4294967296
min_loss_scale: 1.0
loss_scale_window: 1000
hysteresis: 2
accumulate_allreduce_grads_in_fp32: false
fp16_lm_cross_entropy: false

# fp8
fp8: null # e4m3, hybrid
fp8_margin: 0
fp8_recipe: delayed
fp8_interval: 1 # deprecated
fp8_amax_history_len: 1024
fp8_amax_compute_algo: "max"
fp8_wgrad: true
fp8_param_gather: false
te_rng_tracker: false
inference_rng_tracker: false

# fp4
fp4: null
fp4_recipe: nvfp4
fp4_param: false

first_last_layers_bf16: false
num_layers_at_start_in_bf16: 1
num_layers_at_end_in_bf16: 1

# Optimizer
optimizer: adam
lr: 2.5e-4
lr_decay_style: cosine
lr_decay_iters: null
lr_decay_samples: null
lr_warmup_fraction: null
lr_warmup_iters: 0
lr_warmup_samples: 0
lr_warmup_init: 0.0
min_lr: 2.5e-5
lr_wsd_decay_style: exponential
lr_wsd_decay_samples: null
lr_wsd_decay_iters: null
head_lr_mult: 1.0
weight_decay: 0.01
start_weight_decay: null
end_weight_decay: null
weight_decay_incr_style: constant
clip_grad: 1.0
adam_beta1: 0.9
adam_beta2: 0.95
adam_eps: 1.0e-08
sgd_momentum: 0.9
override_opt_param_scheduler: false
use_checkpoint_opt_param_scheduler: false
warmup: null
decoupled_lr: null
decoupled_min_lr: null
# muon (read by Megatron-LM's native get_megatron_optimizer_config() whenever
# args has these attrs -- not consumed by any Primus code, but the only way
# to tune Muon's internals via Primus config since most have no CLI flag)
muon_extra_scale_factor: 1.0
muon_scale_mode: "spectral"
muon_fp32_matmul_prec: "medium"
muon_num_ns_steps: 5
muon_tp_mode: "blockwise"
muon_use_nesterov: false
muon_split_qkv: true
muon_momentum: 0.95

optimizer_cpu_offload: false
optimizer_offload_fraction: 1.0 # float
use_torch_optimizer_for_cpu_offload: false
overlap_cpu_optimizer_d2h_h2d: false
pin_cpu_grads: true
pin_cpu_params: true

# checkpointing arguments
save: null
save_interval: 20000
save_retain_interval: null
no_save_optim: null
no_save_rng: null
load: null
load_main_params_from_ckpt: false
no_load_optim: null
no_load_rng: null
finetune: false
use_checkpoint_args: false
use_mp_args_from_checkpoint_args: false
use_tokenizer_model_from_checkpoint_args: true
exit_on_missing_checkpoint: true
non_persistent_save_interval: null # int
non_persistent_ckpt_type: null # 'global', 'local', 'in_memory', null
non_persistent_global_ckpt_dir: null # str
non_persistent_local_ckpt_dir: null # str
non_persistent_local_ckpt_algo: "fully_parallel" # 'fully_parallel', 'atomic'
dist_ckpt_save_pre_mcore_014: null
dist_ckpt_optim_fully_reshardable: null

pretrained_checkpoint: null
ckpt_step: null
use_dist_ckpt_deprecated: false
use_persistent_ckpt_worker: false
auto_detect_ckpt_format: false
dist_ckpt_format_deprecated: null
ckpt_format: torch_dist # 'torch', 'torch_dist', 'zarr'
ckpt_convert_format: null # 'torch', 'torch_dist', 'zarr'
ckpt_convert_save: null
ckpt_convert_update_legacy_dist_opt_format: false
ckpt_fully_parallel_save_deprecated: false
ckpt_fully_parallel_save: true
async_save: null
ckpt_fully_parallel_load: false
ckpt_assume_constant_structure: false
dist_ckpt_strictness: assume_ok_unexpected

# distributed arguments
overlap_p2p_comm: true
distributed_backend: nccl
distributed_timeout_minutes: 10
defer_embedding_wgrad_compute: false
wgrad_deferral_limit: 0 # int
align_grad_reduce: true
ddp_num_buckets: null # int
ddp_bucket_size: null # int
ddp_pad_buckets_for_high_nccl_busbw: false
ddp_average_in_collective: false
overlap_grad_reduce: false
overlap_param_gather: false
overlap_param_gather_with_optimizer_step: false
align_param_gather: true
scatter_gather_tensors_in_pipeline: true
use_ring_exchange_p2p: false
local_rank: null
lazy_mpu_init: null
account_for_embedding_in_pipeline_split: false
account_for_loss_in_pipeline_split: false
empty_unused_memory_level: 0
standalone_embedding_stage: false
use_distributed_optimizer: false
use_sharp: false
sharp_enabled_group: null  # options: [dp, dp_replica]
use_custom_fsdp: false
use_megatron_fsdp: false
init_model_with_meta_device: false
data_parallel_sharding_strategy: no_shard # 'no_shard', 'optim', 'optim_grads', 'optim_grads_params'
gradient_reduce_div_fusion: true
suggested_communication_unit_size: 400000000 # int
keep_fp8_transpose_cache_when_using_custom_fsdp: false
num_distributed_optimizer_instances: 1 # int
data_parallel_replicate_degree: 1
use_torch_fsdp2: false

# === ODC (On-Demand Communication) ===
# Master switch for ODC's rocSHMEM/XGMI point-to-point gradient reduction on the
# torch-FSDP2 path (ROCm/MI300X). DEFAULT false => ODC is a complete no-op and
# every ODC integration patch is skipped (byte-for-byte stock behaviour). When
# true (requires use_torch_fsdp2=true), Primus:
#   * routes FSDP2 gradient reduce-scatter through ODC (scatter_accumulate),
#   * skips #856's eager-RCCL device_id injection (distributed_init_patches.py),
#   * skips #808's FSDP2 forward all-gather prefetch (it cannot overlap ODC's
#     serial transport); backward prefetch is kept.
# Formerly gated by the ODC_ENABLE env var; now config-driven.
enable_odc: false
# ODC wiring depth (integer): 1 = install the __init__/lazy_init hooks only
# (construction + first fwd/bwd sanity, gradients NOT routed); 2 = also hook
# train_step / optimizer.step for full gradient routing. Production = 2.
odc_phase: 2
# Guard against ODC async-reduce's occasional non-deterministic grad spike: at
# the optimizer step, if the global grad norm exceeds this threshold, the step is
# skipped (grads zeroed) exactly like Megatron skips nan/inf iters. <= 0 disables.
odc_grad_spike_threshold: 1000.0
# ODC symmetric-memory / comm-kernel tuning (consumed by the ODC library via its
# runtime config; these replace the former ODC_* env vars). Defaults reproduce the
# previous env defaults exactly.
odc_p2p_backend: mori          # symmetric-memory backend: mori | rocshmem
odc_mori_init: pg              # MORI init method: pg | uid
odc_max_buffer_size: 67108864  # BufferSplitter max global buffer (bytes, 64 MiB)
odc_rocshmem_gda: false        # rocSHMEM GPU-direct (GDA) device path (multi-node)
odc_rocshmem_lib: null         # optional librs_host_gda.so ctypes override path
odc_gda_rs_blocks: 64          # GDA reduce-scatter kernel grid blocks
odc_gda_pipe: 1                # GDA peer-pipeline depth (also bridged to turbo env)
odc_gda_defer_reduce: auto     # defer cross-node reduce to per-minibatch: auto|1|0
odc_gda_warmup_mode: strided   # write-visibility settle: strided|full|hdp|fence|hdpfence
odc_gda_stride_bytes: 65536    # strided warm-up page stride (bytes)
nccl_communicator_config_path: null
use_tp_pp_dp_mapping: false
replication: false
replication_jump: null # int
replication_factor: null # int
deterministic_mode: false
check_weight_hash_across_dp_replicas_interval: null
overlap_moe_expert_parallel_comm: false

train_iters: null
eval_iters: 32
full_validation: false
multiple_validation_sets: false
eval_interval: 2000
skip_train: false
train_sync_interval: null # int

adlr_autoresume: false
adlr_autoresume_interval: 1000

# activation recomputation
recompute_activations: false
recompute_granularity: null # full, selective
recompute_method: null # uniform, block
recompute_num_layers: null # int
distribute_saved_activations: false
checkpoint_activations: false # deprecated

# garbage collection
manual_gc: false # default false
manual_gc_interval: 1 # int, default 0
manual_gc_eval: false

#data
data_path: null
data_sharding: true
split: "99,1,0"
train_data_path: null
valid_data_path: null
test_data_path: null
data_args_path: null # str
per_split_data_args_path: null # str
data_cache_path: null
mock_data: false
# mock_dataset: Configuration for synthetic dataset (used when mock_data=true)
#   class: Fully qualified class name (optional, defaults based on model_type)
#   params: Dict of parameters passed to dataset constructor
#   Example:
#     mock_dataset:
#       class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset"
#       params:
#         num_samples: 1000
#         image_size: 512
mock_dataset:
  class: null
  params:
    num_samples: 1000
    image_size: 512
merge_file: null
seq_length: 4096
encoder_seq_length: null
decoder_seq_length: null
retriever_seq_length: 256
sample_rate: 1.0
mask_prob: 0.15
short_seq_prob: 0.1
num_workers: 8
reset_position_ids: false
reset_attention_mask: false
eod_mask_loss: false
train_samples: null
dataloader_type: null
mmap_bin_files: true

#profile:
profile: false
use_pytorch_profiler: false
profile_ranks: [0]
profile_step_end: 12
profile_step_start: 10
iterations_to_skip: null
result_rejected_tracker_filename: null
enable_gloo_process_groups: true
record_memory_history: false
memory_snapshot_path: snapshot.pickle # str

#logging:
log_avg_skip_iterations: 2
log_avg_reset_interval: 10
log_params_norm: false
log_num_zeros_in_grad: false
log_throughput: false
wall_clock_step_timer: false
log_progress: false
timing_log_level: 0
timing_log_option: minmax
tensorboard_log_interval: 1
tensorboard_queue_size: 1000
log_timers_to_tensorboard: false
log_batch_size_to_tensorboard: false
log_learning_rate_to_tensorboard: true
log_validation_ppl_to_tensorboard: false
log_memory_to_tensorboard: false
log_world_size_to_tensorboard: false
log_loss_scale_to_tensorboard: true
wandb_project: null
wandb_exp_name: null
wandb_save_dir: null
wandb_entity: null
enable_one_logger: true
one_logger_project: megatron-lm
one_logger_run_name: null
log_interval: 100
tensorboard_dir: null
logging_level: null # int
config_logger_dir: ""

one_logger_async: false
app_tag_run_name: null
app_tag_run_version: 0.0.0

# rerun
error_injection_rate: 0 # int
error_injection_type: transient_error # str: 'correct_result', 'transient_error', 'persistent_error'
rerun_mode: disabled # str: 'disabled', 'validate_results', 'report_stats'

# Experimental features
enable_experimental: false

# # Hybrid model configuration
# hybrid_attention_ratio: 0.0 # float range [0,0, 1.0]
# hybrid_mlp_ratio: 0.0 # float range [0,0, 1.0]
# hybrid_override_pattern: null # str

# # Mamba layer configuration
# mamba_state_dim: 128
# mamba_head_dim: 64
# mamba_num_groups: 8
# mamba_num_heads: null
# disable_mamba_mem_eff_path: false

# Args of precision-aware optimizer
use_precision_aware_optimizer: false
main_grads_dtype: fp32 # str: fp32, bf16
main_params_dtype: fp32 # str: fp32, bf16
exp_avg_dtype: fp32 # 'fp32', 'fp16', 'fp8'
exp_avg_sq_dtype: fp32 # 'fp32', 'fp16', 'fp8'

# vision
vision_pretraining: false
vision_pretraining_type: classify
vision_backbone_type: vit
swin_backbone_type: tiny
num_classes: 1000
img_h: 224
img_w: 224
num_channels: 3
patch_dim: 16
classes_fraction: 1.0
data_per_class_fraction: 1.0

# others
retro_project_dir: null
retro_add_retriever: false
retro_cyclic_train_iters: null
retro_encoder_layers: 2
retro_encoder_hidden_dropout: 0.1
retro_encoder_attention_dropout: 0.1
retro_num_neighbors: 2
retro_num_retrieved_chunks: 2
retro_attention_gate: 1
retro_verify_neighbor_count: true
dino_local_img_size: 96
dino_local_crops_number: 10
dino_head_hidden_size: 2048
dino_bottleneck_size: 256
dino_freeze_last_layer: 1
dino_norm_last_layer: false
dino_warmup_teacher_temp: 0.04
dino_teacher_temp: 0.07
dino_warmup_teacher_temp_epochs: 30

mask_type: random
mask_factor: 1.0
iter_per_epoch: 1250
log_straggler: false
disable_straggler_on_startup: false
straggler_ctrlr_port: 65535
straggler_minmax_count: 1
# inference
inference_batch_times_seqlen_threshold: -1
inference_dynamic_batching: false
inference_dynamic_batching_buffer_size_gb: 40.0 # float
inference_dynamic_batching_buffer_guaranteed_fraction: 0.2 # float
inference_dynamic_batching_buffer_overflow_factor: null # float
inference_dynamic_batching_max_requests_override: null # int
inference_dynamic_batching_max_tokens_override: null # int
max_tokens_to_oom: 12000
output_bert_embeddings: false
bert_embedder_type: megatron # "megatron", "huggingface"
flash_decode: false
enable_cuda_graph: false
cuda_graph_warmup_steps: 3 # int
external_cuda_graph: false
cuda_graph_scope: full # full, attn
inference_max_requests: 8 # int
inference_max_seq_length: 2560 # int, (prefill + decode)

create_attention_mask_in_dataloader: true
num_dataset_builder_threads: 1
ict_head_size: null
biencoder_projection_dim: 0
biencoder_shared_query_context_model: false
ict_load: null
bert_load: null
titles_data_path: null
query_in_block_prob: 0.1
use_one_sent_docs: false
evidence_data_path: null
retriever_report_topk_accuracies: []
retriever_score_scaling: false
block_data_path: null
embedding_path: null
indexer_batch_size: 128
indexer_log_interval: 1000

enable_ft_package: false
calc_ft_timeouts: false
run_workload_inspector_server: false

heterogeneous_layers_config_path: null
heterogeneous_layers_config_encoded_json: null
inprocess_restart: false

# rl_args
perform_rl_step: false
rl_prompts_per_eval: 32
grpo_prompts_per_step: 32
grpo_group_size: 2
grpo_iterations: 2
grpo_clamp_eps_lower: 0.01
grpo_clamp_eps_upper: 0.01
grpo_kl_beta: 0.001
grpo_entropy_term_weight: 0.0
grpo_filter_groups_with_same_reward: false
grpo_default_temperature: 1.0
grpo_default_top_p: 0
langrl_inference_server_type: inplace_megatron
langrl_inference_server_conversation_template: null
langrl_env_config: null
rl_offload_optimizer_during_inference: false
rl_offload_kv_cache_during_training: false
rl_remove_kv_cache_during_training: false
rl_reset_cuda_graphs: false
rl_partial_rollouts: false
rl_inference_logprobs_is_correction: false
rl_importance_sampling_truncation_coef: null
rl_calculate_intra_group_similarity: false
