base.yml

# Copyright 2023–2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This sentinel is a reminder to choose a real run name.
# If there is already a checkpoint under this run, that checkpoint will auto-resume.
#
# NOTE: Some unit/integration tests in MaxText do not always run this file directly.
# When running in decoupled mode (DECOUPLE_GCLOUD=TRUE), tests may use
# `decoupled_base_test.yml` instead of `base.yml` via `tests/utils/test_helpers.py`.
run_name: ""

model_name: "default" # override config settings to match a specific model. other than the override, nothing should use this!
override_model_config: false # When set to true allows overriding model parameters via CLI (or kwargs or env vars) for the purpose of debugging/testing.
override_logical_axis_rules: false # When set overrides logical axis rules instead of merging them.
debug:
  rl: false # RL-specific debugging

normalization_layer_epsilon: 1.e-05 # epsilon value for rmsnorm, layernorm.

################################## CHECKPOINTING ##################################
# Checkpointing makes the following choices in the following order, starting with (1):
#   (1) If there is already a checkpoint for this run_name, we load the latest entire checkpoint.
#     This ensures if we're resuming a run after preemption or hardware failure we lose minimum state.
#   (2) Same priority and mutually exclusive -- you can't set both!
#      * If load_parameters_path is set, we load a parameter only checkpoint from that path.
#      * If load_full_state_path is set, we load a full state checkpoint from that path.
#   (3) We don't load a checkpoint and initialize state instead!

# Loads a just parameters from a specific directory
# e.g. gs://my-base-output-directory/my-previous-run-name/checkpoints/items/NUMBER or NUMBER/items
load_parameters_path: ""

# LoRA adapter support configs
lora_input_adapters_path: ""    # Input GCS path for a parent directory which has all the LoRA adapters (lora_id as subdir)
hf_lora_adapter_path: ""        # Input HF repo ID or local path for HF LoRA adapter

# Loads a full checkpoint including optimizer state and step count from a specific directory
# e.g. gs://my-base-output-directory/my-previous-run-name/checkpoints/items/NUMBER or NUMBER/items
load_full_state_path: ""

# If enable_checkpointing is true, an asynchronous checkpointer will be used if
# async_checkpointing is true, else a synchronous one is used. If you have
# problems with the checkpointer we recommend trying the synchronous one.
enable_checkpointing: true
save_checkpoint_on_completion: true
async_checkpointing: true
checkpoint_period: 10_000
max_num_checkpoints_to_keep: None
enable_continuous_checkpointing: false
# enables one replica to read the ckpt then broadcast to the rest
enable_single_replica_ckpt_restoring: false
# Subdirectory to move checkpoints to before deletion. For example: ".todelete" (Ignored if directory is prefixed with gs://)
checkpoint_todelete_subdir: None
# Full path to move checkpoints to before deletion.
checkpoint_todelete_full_path: None

force_unroll: false # during generate_param_only_checkpoint should we unroll the loop?

# checkpointing using orbax has two important parameters: array driver
# and its underlying storage - the kvstore (preferably ocdbt)
# orbax supports setting a target file size, chunking a single
# large arrays into small physical files (<2GB) can speed up distributed and over
# the network loading enormously
checkpoint_storage_target_data_file_size_bytes: 2147483648
checkpoint_storage_use_ocdbt: true
checkpoint_storage_use_zarr3: true
# larger models requires higher concurrent GB for I/O
# default concurrent gb for PytreeCheckpointHandler is 96GB
checkpoint_storage_concurrent_gb: 96

# Bool flag for enabling Orbax v1.
enable_orbax_v1: false
# function for processing loaded checkpoint dict into a format maxtext can understand. (for other formats, i.e. safetensors)
checkpoint_conversion_fn: none
# optional checkpoint context to use for loading. options: "orbax", "safetensors"
source_checkpoint_layout: "orbax"

# Only applicable to Single Controller/Pathways on Cloud. Experimental feature, under testing
colocated_python_checkpointing: false

# enables autocheckpoint, which saves a checkpoint at the preemption step.
enable_autocheckpoint: false
############################### end checkpointing ##################################


reuse_example_batch: 0 # for testing tpu performance, this options repeated uses the same batch.


metrics_file: "" # for testing, local file that stores scalar metrics. if empty, no metrics are written.
# if true save metrics such as loss and tflops to gcs in {base_output_directory}/{run_name}/metrics/
gcs_metrics: false

# if true save config to gcs in {base_output_directory}/{run_name}/
save_config_to_gcs: false

# gradient dtype
grad_dtype: "float32"

# activation dtypes.
dtype: "bfloat16"
# used to configure quantization in the transformer layers, defaults to null implying bf16.
# possible alternative settings are as follows:
# 'int8' for dynamic range quantization using 8-bits
# 'intmp' for mixed precision quantization for inference as described here: src/maxtext/configs/quantization/readme.md
# 'fp8' for 8-bit floating-point gemms on nvidia gpus.
# 'nanoo_fp8' for 8-bit floating-point gemms on amd mi300/mi325 gpus.
# 'fp8_full' for fp8 quantization with static scaling.
quantization: ""
# used to configure constant_bound_config in aqt lib for static scaling, e.g. constant_bound_config='0.5, 0.5, 0.5, 0.5, 0.5, 0.5'
constant_bound_config: ""
# choose one of default, high, and highest.
# https://kolonist26-jax-kr.readthedocs.io/en/latest/jax.lax.html#jax.lax.precision
matmul_precision: "default"
activations_in_float32: false # sets activations to float32 before nonlinearity it true, else dtype
# used to replicate the quantization scale to avoid the inefficient xla fusion for 2d sharding.
replicate_quant_scale: false
# path to file with quantization config for intmp.
quant_cfg_path: ""
quantize_kvcache: false # set to true to quantize kv cache values, defaults to false
# valid kv_quant_axis values:
#   - "" is valid only when quantize_kvcache is false
#   - "dkv" indicates quantize kv cache over the cache_kv, i.e. kv dimension axis
#   - "heads_and_dkv" indicates quantize kv cache over cache_heads and cache_kv axes
# default to "heads_and_dkv" for faster compution, kv_quant_axis is not used when quantize_kvcache is false
#   - "dkv" is expected with better accuracy but degraded computation
kv_quant_axis: "heads_and_dkv"
kv_quant_dtype: "int8"
checkpoint_is_quantized: false # set to true if reading from a saved aqt quantized checkpoint
# saves params quantized on fly at following path
save_quantized_params_path: ""
#used to configure the mode in which model is called
# when left as is, corresponds to training
# accepted values are "inference"
model_call_mode: ""
use_qwix_quantization: false # [DEPRECATED: AQT will be removed in a future release. It is strongly recommended to set use_qwix_quantization to true] whether to use qwix for quantization. if set to true, the model will be quantized using qwix.
use_manual_quantization: false # a flag if to use manual quantization for batch split. Only used if use_batch_split_schedule is true.
# quantization calibration method used for weights and activations. supported methods can be found in https://github.com/google/qwix/blob/dc2a0770351c740e5ab3cce7c0efe9f7beacce9e/qwix/qconfig.py#l70-l80
weight_quantization_calibration_method: "absmax"
act_quantization_calibration_method: "absmax"
bwd_quantization_calibration_method: "absmax"
# shard the range finding operation for quantization. by default this is set to number of slices.
quantization_local_shard_count: -1
# The 'N' in N:M sparsity, representing the maximum number of non-zero values in each block.
weight_sparsity_n: null
# The 'M' in N:M sparsity, representing the number of values in each block.
weight_sparsity_m: null
# The step size to update the sparsity masks.
weight_sparsity_update_step: 10
# The first number of steps before updating the sparsity masks.
weight_sparsity_start_step: 50

decoder_block: "llama2" # which style of decoderblock to use.
# global parameter scale needs to be a power of 2. if you want finer grained control of the model sizes
# then you should explicitly set base_embed_dim, base_num_query_heads, base_num_kv_heads,
# base_mlp_dim, base_num_decoder_layers and/or head_dim.
weight_dtype: "float32"
global_parameter_scale: 1
base_emb_dim: 2048
base_num_query_heads: 16
base_num_kv_heads: 16
base_mlp_dim: 7168
dense_init_scale: 1.0
base_num_decoder_layers: 16
head_dim: 128
attention_output_dim: -1
# Those parameters are only used with global attention for Gemma4.
global_head_dim: 0
global_num_kv_heads: 0
mlp_activations: ["silu", "linear"]
mlp_activations_limit: -1.0
dropout_rate: 0.0
logits_via_embedding: false
normalize_embedding_logits: true  # whether to normalize pre-softmax logits if logits_via_embedding is true
logits_dot_in_fp32: false  # whether to use fp32 in logits_dense or shared_embedding dot product for stability
cast_logits_to_fp32: true # whether to cast the logits to fp32. the higher precision is generally beneficial, but it can vary slightly.
float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product
float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax
float32_weight_sum: true # whether to use full fp32 precision to sum expert weights for numerical stability
float32_gate_logits: false # whether to cast inputs to fp32 to compute MoE gate logits for numerical stability

# multi-token prediction configs
# the number of auxiliary prediction layers to use for mtp.
# set to 0 to disable the feature.
mtp_num_layers: 0
# the scaling factor (lambda) for the mtp auxiliary loss. the final loss is:
# main_loss + mtp_loss_scaling_factor * avg_mtp_loss
mtp_loss_scaling_factor: 0.1
# specifies which mtp layer (1-indexed) is used to calculate metrics like the
# acceptance rate during evaluation. for example, a value of `1` targets the
# first auxiliary prediction head. set to 0 to disable this specific evaluation
mtp_eval_target_module: 0

# mixture of experts (moe)
num_experts: 1
num_experts_per_tok: 1
megablox: true
sparse_matmul: true
capacity_factor: -1.0 # a factor to decide expert capacity for token dropping, and no dropping by default
ragged_buffer_factor: -1.0 # a factor to determine the size of the ragged buffer for routed MoE activations.
# By default (-1), the routed buffer is worst case size to ensure no dropping.
# When set to 1.0 this buffer if set to the size assuming perfectly balanced. If the routing dictates
# a size larger than this then tokens are dropped.
# In general if ragged_buffer_factor > 0, the ragged_buffer_size is balanced_size * ragged_buffer_factor.
moe_expert_input_dim: -1 # feature dimension of the tokens entering the MoE expert blocks.
base_moe_mlp_dim: -1 # intermediate dimension at MoE layer. For a fully MoE model, base_mlp_dim must be equal to base_moe_mlp_dim.
load_balance_loss_weight: 0.0 # weight for the load balance loss
use_random_routing: false # whether to use random routing for debug/test purpose
use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul
use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism
# If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM
# stays expert-parallel (AllToAll); false keeps 'expert' on the batch dim (activation_batch_moe).
# Only affects the dense (dense_matmul) MoE path; the sparse (shard_map) path is unaffected.
moe_dispatch_no_expert_sharding: false
use_ragged_sort: false # whether to use the Pallas ragged-sort kernels in the MoE permute path; valid both with and
                       # without `use_ring_of_experts` (with EP > 1). When `use_ring_of_experts=True` the kernels run
                       # inside `permute`/`unpermute`; otherwise they run inside `local_permute`/local-unpermute.
use_gather_mosaic_kernel: false # whether to use a custom mosaic kernel for token gather ops
ragged_gather_fallback: false # when true, unconditionally use the JAX reference implementation instead of the
                              # ragged gather SparseCore kernel. When false (default), use the SparseCore kernel.
ragged_gather_reduce_fallback: false # when true, unconditionally use the JAX reference implementation instead of the
                                     # ragged gather reduce SparseCore kernel. When false (default), use the SparseCore kernel.
ragged_gather_cost_estimate_flops: -1 # -1 means auto-compute, any > 0 value overrides the flop cost estimate for the ragged gather kernel
ragged_gather_reduce_cost_estimate_flops: -1 # -1 means auto-compute, any > 0 value overrides the flop cost estimate for the ragged gather reduce kernel
ragged_gather_cost_estimate_bytes_accessed: -1 # -1 means auto-compute, any > 0 value overrides the bytes_accessed cost estimate for the ragged gather kernel
ragged_gather_reduce_cost_estimate_bytes_accessed: -1 # -1 means auto-compute, any > 0 value overrides the bytes_accessed cost estimate for the ragged gather reduce kernel
# tunable tiling dimensions used for mlp gmm
# megablox/jax ragged dot - supports forward pass only (6 configs: `wi_tile_fwd...` and `wo_tile_fwd_...`)
# tokamax ragged dot - supports all 18 configs
wi_tile_fwd_batch_seq: 512
wi_tile_fwd_embed_dim: 1024
wi_tile_fwd_mlp_dim: 1024
wi_tile_dlhs_batch_seq: 512
wi_tile_dlhs_embed_dim: 1024
wi_tile_dlhs_mlp_dim: 1024
wi_tile_drhs_batch_seq: 512
wi_tile_drhs_embed_dim: 1024
wi_tile_drhs_mlp_dim: 1024

wo_tile_fwd_batch_seq: 512
wo_tile_fwd_embed_dim: 1024
wo_tile_fwd_mlp_dim: 1024
wo_tile_dlhs_batch_seq: 512
wo_tile_dlhs_embed_dim: 1024
wo_tile_dlhs_mlp_dim: 1024
wo_tile_drhs_batch_seq: 512
wo_tile_drhs_embed_dim: 1024
wo_tile_drhs_mlp_dim: 1024

merge_gating_gmm: false

norm_topk_prob: false # boolean to enable the top-k probability normalization. qwen3-specific normalization of router weights.

# when moe weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls
moe_fsdp_use_two_stage_all_gather: false
# Shard the expert dimension of the MLP weights on the FSDP axis.
# This configuration is recommended only when num_experts is a multiple of fsdp_parallelism
shard_exp_on_fsdp: false

# deepseek moe
first_num_dense_layers: 0 # number of initial dense layers in the model
shared_experts: 0
routed_scaling_factor: 1.0 # scaling factor for routing scores
routed_score_func: "" # scoring function for routing
routed_bias: false # a flag if a learnable bias is added for routing
routed_bias_update_rate: 0.0 # a flag indicate the update rate applied to the router bias term
mlp_bias: false # a flag if a learnable bias is added for MLP matmul, and originally implemented to support the GPT-OSS model architecture.
n_routing_groups: -1 # number of groups for routing, disabled by default
first_num_hash_layers: 0 # number of hash routing layers, used in DeepSeek V4 (0 means disabled)
topk_routing_group: -1 # number of top groups to route inputs. For EP,
# Splits the batch to allow for better scheduling when using expert parallelism by overlapping the
# all-to-all communication with compute. Currently only implemented with DeepSeek sparse layers.
use_batch_split_schedule: false # a flag if splitting batch into micro-batches to hide communications that yields performance benefits.
batch_split_factor: 1 # the factor by which to split the batch. Only used if use_batch_split_schedule is true.

# For complex architectures like llama4 there are repeated sets of
# inhomogeneous layers. E.g. maverick uses [dense+rope, moe+rope, dense+rope, moe+nope]
# which can only be scanned together in one large block of inhomogeneous_layer_cycle_interval=4 layers.
inhomogeneous_layer_cycle_interval: 1

# pipeline parallelism
# The number of decoder layers is equal to the product of num_stages, num_layers_per_pipeline_stage and num_pipeline_repeats.
# There is a tradeoff between the num_layers_per_pipeline_stage and num_pipeline_repeats: The more layers per stage the easier
# it is to hide the pipeline communication behind the compute since there is more compute per stage, however there will be a larger bubble
# since there are fewer repeats. Similarly, there is tradeoff for num_pipeline_microbatches - more microbatches leads to a smaller bubble,
# but a smaller size per microbatch which may hurt per-stage performance. Additionally, note when microbatches > num_stages we have the opportunity to
# perform the circular transfer (last stage to first) asynchronously.
# The bubble fraction is (num_stages - 1) / (num_pipeline_repeats * num_pipeline_microbatches + num_stages - 1)
num_layers_per_pipeline_stage: 1
# The number of repeats will be set to num_decoder_layers / (num_pipeline_stages * num_layers_per_pipeline_stage)
num_pipeline_repeats: -1
pipeline_parallel_layers: -1 # Pipeline only this number of layers - for the remaining layers the "stage" mesh axes will act like data parallelism.
# This option helps when the number of layers does not have friendly divisors since SPMD pipelining requires that the
# PP degree divides the number of layers.
# By default (when set to -1) we pipeline all of the decoder layers.

# num_pipeline_microbatches must be a multiple of the number of pipeline stages. By default it is set to the number of stages.
# Note the microbatch_size is given by global_batch_size / num_pipeline_microbatches, where global_batch_size = per_device_batch_size * num_devices
num_pipeline_microbatches: -1
pipeline_delay_activation_forwarding: false # This delays the activation forwarding one loop iteration simplifying XLA's task of overlapping since
# the communication and compute in each iteration are now independent. However this comes at the cost of doubling the pipeline bubble,
# and you must set the number of microbatches to at least 2 * num_stages (the minimum 2 * num_stages is set by default with this delay).

pipeline_fsdp_ag_once: false # If set to true then all gather all of the weights over FSDP before the first pipeline iteration.
# This is a memory/time tradeoff - we now have to store the FSDP gathered weights and gradients (typically in bf16), as opposed
# to only one stage's worth, however we only execute one all-gather and reduce across per repeat, as opposed
# to every microbatch. This is similar to zero-1 sharding, since we also don't need to all gather the FSDP weights in the backward pass.
# An alternative to setting this to true may be to replace any FSDP with DP and use optimizer offloading if necessary.
pipeline_fsdp_ag_per_repeat: false
# Pipeline weight prefetching per repeat is an advanced SPMD pipeline parallelism improvement technique
# When enabled, it prefetches necessary weight gathering ahead of microbatched computation, therefore reducing collectives

# There are two loops for PP:
#  1)  Outer loop over microbatches (pipeline iterations)
#  2)  Inner loop over layers (layers per stage)
# We have observed extra remat when a remat policy and scanning is performed on both, and recommend the default
# settings below of scanning and setting a remat policy only over the pipeline iterations.
# It may be useful to do the reverse when the layers_per_stage is very large.
# The below settings only have effect when using pipeline parallelism.
scan_pipeline_iterations: true
scan_pipeline_repeats: false
scan_layers_per_stage: false
set_remat_policy_on_pipeline_iterations: true
set_remat_policy_on_layers_per_stage: false


# Choose 'remat_policy' between 'minimal_with_context', 'minimal', 'save_dot_with_context_except_mlp', 'save_dot_except_mlpwi', 'save_dot_except_mlp',
# 'save_qkv_proj', 'qkv_proj_offloaded', 'custom', 'minimal_offloaded', 'save_out_proj' and 'full'.
# These options offer a trade-off between speed (fastest to slowest) and HBM usage (highest to lowest)
remat_policy: 'full'
# If "custom" remat_policy is chosen, you can select tensors from the following list to offload on host memory, rematerialize or save on device memory.
# Pick one of these options for following tensors: ['remat','device','offload']
decoder_layer_input: 'device' # this tensor cannot be rematerialized - it serves as periodic checkpoints that act as the remat start points
context: 'remat' # From https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/attention.py#L581-L583
mlpwi: 'remat'
mlpwi_0: 'remat'
mlpwi_1: 'remat'
mlpwo: 'remat'
moe_mlpwi_0: 'remat'
moe_mlpwi_1: 'remat'
moe_mlpwo: 'remat'
query_proj: 'remat'
key_proj: 'remat'
value_proj: 'remat'
qkv_proj: 'remat'
out_proj: 'remat'
query_wa_proj: 'remat'
kv_wa_proj: 'remat'
mla_q: 'remat'
mla_kv: 'remat'
attention_out: 'remat'
engram: 'remat'

optimizer_memory_host_offload: false
parameter_memory_host_offload: false
scan_layers: true # We recommend setting this to false when using pipeline parallelism, instead scanning the PP iterations.
param_scan_axis: 1

# The attention parameter dictates the specific algorithm/methodology used to compute the attention scores
# The attention_type parameter determines the variants of attention, e.g. global or local_sliding
attention: 'autoselected' # Supported attention: autoselected, dot_product, flash, cudnn_flash_te
attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla
share_kv_projections: false # Note: Not compatible with attention_type='mla'
attention_bias: false # If true, adds a learnable bias to the query, key, and value projections
attention_sink: false
sliding_window_size: 0
chunk_attn_window_size: 0
attn_logits_soft_cap: 0.0
final_logits_soft_cap: 0.0
z_loss_multiplier: 0.0
use_post_attn_norm: false
use_post_ffw_norm: false
v_norm_with_scale: true
qk_norm_with_scale: true
mla_naive_kvcache: true

# Gemma 4 small (E2B / E4B) specific config fields.
# When `hidden_size_per_layer_input` is positive, each decoder layer adds an
# extra "Per-Layer-Embedding" sub-block at the end (see models/gemma4_small.py).
hidden_size_per_layer_input: 0
vocab_size_per_layer_input: 0
# Number of *trailing* decoder layers that reuse K/V from the last non-shared
# layer of the same attention type (sliding↔sliding, full↔full). Those layers
# do not have their own k/v projections.
num_kv_shared_layers: 0
# When True, KV-shared layers double their MLP intermediate size to compensate
# for the lost K/V parameters.
use_double_wide_mlp: False


# Adding Mixture of Block Attention Support (MoBA): https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf
moba: false
moba_chunk_size: 1024
moba_topk: 8

# DeepSeek Sparse Attention (DSA)
# deepseek3.2 introduces indexer in MLA
use_indexer: false
indexer_head_dim: 128
indexer_n_heads: 64
indexer_topk: 2048
# Determines the training strategy for the indexer:
# - false (Dense Warm-up): Computes indexer loss over all tokens. Used with `trainable_parameters_mask` to freeze other model parameters.
# - true (Sparse Training): Computes indexer loss over top-k tokens only and detaches the indexer input for independent optimization.
# Note: This is only active when `indexer_loss_scaling_factor` > 0.
indexer_sparse_training: false
# Multiplier for the indexer KL divergence loss
indexer_loss_scaling_factor: 0.0

# MLA parameters
q_lora_rank: 0
kv_lora_rank: 512
qk_nope_head_dim: 128
qk_rope_head_dim: 64
v_head_dim: 128

# Compressed Attention parameters
o_lora_rank: 0 # Output LoRA rank for Compressed Attention.
o_groups: 0 # Output groups for Compressed Attention.
compress_ratios: [] # Per-layer compression ratios (0, 4, 128, etc).
compressed_rope_max_timescale: 160_000 # If positive, used for Compressed Sparse/Heavy Attention.

# QK-Clip (Muon Clip) Configuration
use_qk_clip: false  # Enable QK-Clip (supported in MLA with DotProduct or Tokamax Splash)
qk_clip_threshold: 100.0  # Threshold for clipping (tau in the paper)

# Combine matmuls for QKV and MLP
fused_qkv: false
fused_mlp: false

record_internal_nn_metrics: 0

# Output directory
# Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/"
base_output_directory: ""

# Multi-tier checkpointing is an experimental Orbax feature that: periodically saves to persistent storage(GCS bucket) dictated by `multi_tier_checkpointing_backup_interval_minutes` and,
# saves to a local directory for smaller checkpoint intervals(local_checkpoint_period).
# The local checkpoint directory must be specified when enabling multi-tier checkpointing.
# During restore, if a local copy is available in any slice, it will be broadcast to other slices without having to fetch from persistent storage.
# See more details on https://github.com/google/orbax/tree/main/checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing.
# Example for enabling multi-tier checkpointing
# enable_multi_tier_checkpointing=true local_checkpoint_directory="/local" local_checkpoint_period=20 multi_tier_checkpointing_backup_interval_minutes=20
enable_multi_tier_checkpointing: false

# The interval to backup local checkpoints to the persistent storage(GCS bucket) in minutes.
# It should be a positive number when enabling multi-tier checkpointing.
multi_tier_checkpointing_backup_interval_minutes: 0

# Number of identical pipelines in job, should be equal to ICI data parallelism * DCN data parallelism.
# It should be a positive number when enabling multi-tier checkpointing. If set to 0, it will be set to num of slices.
mtc_data_parallelism: 0


# Whether to enable emergency checkpoint. If true, `local_checkpoint_directory` and a non-zero `local_checkpoint_period` must also be specified.
# Emergency checkpoint is an experimental Orbax feature that: periodically saves to persistent storage and, with a larger invertal, saves to a local directory.
# During restore, if a local copy is available in any slice, it will be broadcast to other slices without having to fetch from persistent storage.
# See more details on https://github.com/google/orbax/tree/main/checkpoint/orbax/checkpoint/experimental/emergency.
enable_emergency_checkpoint: false

# It should be specified when and only when `enable_emergency_checkpoint` is true. Or when `enable_multi_tier_checkpointing` is true.
local_checkpoint_directory: ""

# It should be a positive number when and only when `enable_emergency_checkpoint` or `enable_multi_tier_checkpointing` is true.
local_checkpoint_period: 0

# Jax cache directory
jax_cache_dir: "~/jax_cache"

# Hardware
hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu', 'gpu_multiprocess' and 'cpu'

# internal_compile allows bypassing open-source topology name mappings when using internal topologies directly via get_topology_desc.
internal_compile: false
internal_compile_num_devices: -1 # You must specify the number of devices when using internal_compile.
compile_xla_flags: "" # Compiler options e.g. compile_xla_flags="--xla_tpu_num_sparse_cores_for_gather_offloading=1 --xla_tpu_scoped_vmem_limit_kib=65536"

# Parallelism
shard_mode: "auto" # can be either auto or explicit
custom_mesh_and_rule: "" # replace default mesh and logical rule by specifying yml name under config/mesh_and_rule/.
mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']
logical_axis_rules: [
                      # ==========================================
                      # Vocabulary Embedding
                      # ==========================================
                      # Vocab Activations
                      ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']],
                      ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'expert']],
                      ['activation_vocab', ['tensor', 'tensor_transpose', 'tensor_sequence']],
                      ['activation_vocab', ['tensor', 'tensor_transpose']],
                      ['activation_vocab', 'tensor_sequence'],
                      # Vocab Weights
                      ['vocab', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
                      ['embed_vocab', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
                      # ==========================================
                      # Attention
                      # ==========================================
                      # Attention Activations
                      ['activation_batch_attn', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                      ['activation_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
                      ['activation_kv_heads', ['tensor', 'tensor_transpose', 'tensor_sequence']],
                      ['activation_length_attn', ['context']],
                      ['activation_q_length', ['context']],
                      ['activation_kv_length', []],
                      ['activation_embed_attn', ['tensor', 'tensor_transpose']],
                      ['activation_kv', ['tensor', 'tensor_transpose', 'tensor_sequence']],
                      ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                      ['activation_kv_head_dim', ['tensor', 'tensor_transpose', 'tensor_sequence']],
                      # Attention Weights
                      ['heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
                      ['q_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
                      ['kv_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']],
                      ['qkv', []],
                      ['kv', []],
                      ['kv_head_dim', []],
                      ['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'tensor_transpose', 'expert']],
                      ['q_lora', ['fsdp', 'context', 'tensor_transpose', 'expert']],
                      ['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
                      ['q_lora', ['fsdp', 'context', 'expert']],
                      ["q_lora_up_proj", []],
                      ['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'tensor_transpose', 'expert']],
                      ['kv_lora', ['fsdp', 'context', 'tensor_transpose', 'expert']],
                      ['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
                      ['kv_lora', ['fsdp', 'context', 'expert']],
                      ["kv_lora_up_proj", []],
                      # ==========================================
                      # Mixture of Experts (MoE)
                      # ==========================================
                      # MoE Activations
                      ['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                      ['activation_length_moe', ['context']],
                      ['activation_norm_length_moe', ['tensor_sequence', 'context']],
                      ['activation_embed_moe', ['tensor', 'tensor_transpose']],
                      ['activation_mlp_moe', ['tensor', 'tensor_transpose', 'tensor_sequence']],
                      ['activation_exp', ['expert']],
                      # MoE Weights
                      ['exp', 'expert'],
                      ['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
                      ['embed_moe', ['fsdp', 'fsdp_transpose', 'tensor_transpose', 'context']],
                      ['embed_moe', ['fsdp', 'tensor_transpose', 'context']],
                      ['embed_moe', ['fsdp', 'fsdp_transpose', 'context']],
                      ['embed_moe', ['fsdp', 'context']],
                      # ==========================================
                      # Standard MLP / Dense Layers / Model Structure
                      # ==========================================
                      # Dense Activations
                      ['activation_mlp', ['tensor', 'tensor_transpose', 'tensor_sequence']],
                      # Note activation batch and length also get used in vocab
                      ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                      ['activation_length', ['context']],
                      ['activation_norm_length', ['tensor_sequence', 'context']],
                      ['activation_embed', ['tensor', 'tensor_transpose']],
                      ['activation_stage', 'stage'],
                      # General Weights
                      ['mlp', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
                      # GDN (linear-attention) projections shard like 'mlp' during training; the
                      # vLLM serving config overrides this to match tpu-inference's ATTN_HEAD order.
                      ['gdn_head', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']],
                      ['embed', ['fsdp', 'fsdp_transpose', 'tensor_transpose', 'context', 'expert']],
                      ['embed', ['fsdp', 'tensor_transpose', 'context', 'expert']],
                      ['embed', ['fsdp', 'fsdp_transpose', 'context', 'expert']],
                      ['embed', ['fsdp', 'context', 'expert']],
                      ['norm', ['tensor', 'tensor_transpose']],
                      ['layers', 'stage'],
                      ['diloco', 'diloco'],
                      ['engram_dim', ['tensor']],
                      ['dense_layers', []],
                      ['moe_layers', []],
                      ['mhc', []],
                      # ==========================================
                      # Inference(Prefill, Decode, Cache)
                      # ==========================================
                      ['prefill_activation_length', ['context']],
                      ['prefill_activation_norm_length', ['tensor_sequence', 'context']],
                      ['activation_prefill_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                      ['decode_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                      ['decode_length', []],
                      ['cache_heads', ['autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence']],
                      ['cache_heads', ['autoregressive', 'tensor', 'tensor_sequence']],
                      ['paged_kv_heads', ['tensor']],
                      ['cache_batch_prefill', []],
                      ['cache_batch', []],
                      ['cache_heads_none', []],
                      ['cache_kv', []],
                      ['cache_sequence', []],
                      ['num_pages', []],
                      ['tokens_per_page', []],
                      ['paged_kv_head_dim_size', []],
                      # ==========================================
                      # Deprecated / Scheduled for Removal
                      # ==========================================
                      ['mlp_no_fsdp', ['tensor', 'tensor_sequence', 'autoregressive']],
                      ['embed_tensor_transpose', ['tensor_transpose']],
                      ['exp_with_fsdp', 'fsdp'],
                  ]
# Axes used for DCN must be earlier in this list than ICI, see (b/339009148) for details
data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']]
input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activation_norm_length']
# Determines which physical axis plays the role of context parallelism for input data processing and load balancing
# only supports "context" or "expert" (when custom_mesh_and_rule=ep-as-cp)
context_sharding: "context"

# sharding tolerance: float between 0.0 and 1.0 representing the allowed percentage of non-sharded parameters.
sharding_tolerance: 0.02

# One axis for each parallelism type may hold a placeholder (-1)
# value to auto-shard based on available slices and devices.
# By default, product of the DCN axes should equal number of slices
# and product of the ICI axes should equal number of devices per slice.
dcn_diloco_parallelism: 1
dcn_data_parallelism: -1  # recommended DCN axis to be auto-sharded
dcn_fsdp_parallelism: 1
dcn_fsdp_transpose_parallelism: 1
dcn_sequence_parallelism: 1  # never recommended
dcn_context_parallelism: 1
dcn_context_autoregressive_parallelism: 1
dcn_tensor_parallelism: 1 # never recommended
dcn_tensor_transpose_parallelism: 1
dcn_tensor_sequence_parallelism: 1 # never recommended
dcn_pipeline_parallelism: 1
dcn_expert_parallelism: 1
dcn_autoregressive_parallelism: 1 # never recommended
ici_diloco_parallelism: 1
ici_data_parallelism: 1
ici_fsdp_parallelism: -1 # recommended ICI axis to be auto-sharded
ici_fsdp_transpose_parallelism: 1
ici_sequence_parallelism: 1
ici_context_parallelism: 1
ici_context_autoregressive_parallelism: 1
ici_tensor_parallelism: 1
ici_tensor_transpose_parallelism: 1
ici_tensor_sequence_parallelism: 1
ici_autoregressive_parallelism: 1
ici_pipeline_parallelism: 1
ici_expert_parallelism: 1

# Enabling check_vma is recommended for improved performance. Only supported for EP / FSDP ICI parallelisms, shard_mode: "auto", use_ragged_sort: False, use_ring_of_experts: False, and use_tokamax_gmm=False.
check_vma: False

# Enable ZeRO-1 optimizer sharding over data axis
shard_optimizer_over_data: false

# Unless explicitly specified, the number of TPU slices is automatically determined. It should only be set for
# disaggregated reinforcement learning workloads using multiple slices. For ahead of time compilation,
# you should set compile_toplogy_num_slices, which will in turn set this value. For non-TPU environments this is set to 1.
num_slices: -1

# Vocab Tiling Configs
# Enables a memory-saving optimization by computing the cross-entropy loss in chunks.
# The logits are tiled into `num_vocab_tiling` parts along the batch-sequence axis,
# reducing peak memory usage. This is highly recommended for models with large
# vocabularies (e.g., Gemma). Set to a value greater than 1 to enable.
num_vocab_tiling: 1

# Tokenizer
vocab_size: 32_000 # powers of 2 for sharding
tokenizer_path: ""
# grain and tfds pipeline supports tokenizer_type: sentencepiece, huggingface, tiktoken
# hf pipeline only supports huggingface type, and will ignore tokenizer_type flag
tokenizer_type: "sentencepiece" # Currently supporting: "tiktoken", "sentencepiece", "huggingface"
use_chat_template: false
chat_template_path: "" # path to chat template json file
chat_template: "" # Chat template to use with HF tokenizers. It should be a valid Jinja2-formatted template.
tokenize_train_data: true  # false if the dataset is pre-tokenized
tokenize_eval_data: true  # false if the dataset is pre-tokenized
add_bos: true
add_eos: true
# If false, use chunking for long sequences instead of truncation.
# Note: use_truncation=false is only available in grain's pretrain preprocessing pipeline.
# See the TokenizeAndTrim and TokenizeAndChunk classes in
# `src/maxtext/input_pipeline/_grain_tokenizer.py` for implementation details.
use_truncation: true

# Dataset
per_device_batch_size: 12.0
# When expansion_factor_real_data is set to > 1, total_hosts//expansion_factor_real_data will load data.
# Each data-loading host will load per_device_batch_size * expansion_factor_real_data.
# When set to between 0 and 1, it's for grain pipeline to use a smaller chip count to read checkpoint from a larger chip count job.
# Details in https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline/data_input_grain.md#using-grain
expansion_factor_real_data: -1.0
eval_per_device_batch_size: 0.0
max_corpus_chars: 10_000_000
train_data_columns: ['text'] # for DPO dataset containing "chosen" and "rejected"
train_image_column: 'image'
eval_data_columns: ['text'] # for DPO dataset containing "chosen" and "rejected"
eval_image_column: 'image'
packing: true
num_epoch: 1
generate_padding_batch_train: false
generate_padding_batch_eval: false
# Maximum number of segments that can be packed into a single sequence
# This needs to be passed to TransformerEngine's DotProductAttention layer for packing
# This also affects packing for grain, since TransformerEngine may crash or cause
# data corruption if there are more segments packed than specified
# Set this to something like 32 for GPUs when using TransformerEngine
max_segments_per_seq: -1
# Rampup batch size, similar to Megatron-LM, see
# https://github.com/NVIDIA/Megatron-LM/blob/2a01637aa54ccdaf7ea9afc1f1b80f58c53d7f3c/megatron/core/num_microbatches_calculator.py#L233-L237
# The ramp-up proceeds in stages from `per_device_batch_size_start` up to
# the final `per_device_batch_size`. For a clean ramp-up, the total range
# (`per_device_batch_size` - `per_device_batch_size_start`)
# should be evenly divisible by batch size increment.
enable_rampup_batch_size: false
per_device_batch_size_start: 4.0
per_device_batch_size_increment: 2.0
# The target number of training samples to process during the ramp-up phase.
# There is no strict rule for this value, it only needs to be positive.
global_rampup_samples: 500

# direct preference optimization (DPO)
use_dpo: false

# Supervised Fine-Tuning (SFT)
use_sft: false
# sft_train_on_completion_only=false trains on both prompt and completion tokens; trains only on completion tokens otherwise
sft_train_on_completion_only: false

# dataset_type must be synthetic, hf, grain, tfds
# details in: https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline.md
dataset_type: tfds
# for TFDS input pipeline (dataset_type=tfds)
dataset_path: "" # your path given as argument in download_dataset.sh, e.g. "gs://my-maxtext-dataset/"
dataset_name: 'c4/en:3.1.0'
eval_dataset_name: 'c4/en:3.1.0'
train_split: 'train'
eval_split: 'validation'
# for HuggingFace input pipeline (dataset_type=hf)
# Check definition at https://github.com/huggingface/datasets/blob/0feb65dd8733191dd2d1e74215b422fc5939a56a/src/datasets/load.py#L1338-L1408
hf_path: ''
hf_name: ''
hf_data_dir: ''
hf_train_files: ''
hf_eval_split: ''
hf_eval_files: ''
hf_access_token: ''
# for Grain input pipeline (dataset_type=grain)
# Path to grain data files. Can be a single pattern or multiple patterns with weights.
# For multiple patterns, use semicolon (;) to separate and comma (,) to specify weights.
# Example: "path/to/data1.array_record*,0.3;path/to/data2.array_record*,0.7"
# Note: When using multiple files (separated by ';'), only ArrayRecord format is supported.
# For more details, see https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/data_input_pipeline/data_input_grain.md
grain_train_files: ''
grain_eval_files: ''
grain_train_mixture_config_path: '' # Path to a JSON file specifying the mixture weights for Grain training data.
grain_file_type: 'arrayrecord' # arrayrecord or parquet
grain_packing_type: 'first_fit' # 'first_fit', 'best_fit' or 'concat_then_split'. See details of the corresponding module in https://google-grain.readthedocs.io/en/latest/grain.experimental.html
grain_worker_count: 1 # Set to -1 to enable auto-tuning: automatically determines optimal worker count. See https://google-grain.readthedocs.io/en/latest/_autosummary/grain.experimental.pick_performance_config.html
grain_per_worker_buffer_size: 1
# num_threads and prefetch_buffer_size are per-worker per-dataset.
# When using array_records, they are used in ReadOptions (https://google-grain.readthedocs.io/en/latest/tutorials/data_loader_tutorial.html#per-worker-readoptions)
# The default value matches that in the Grain package. If mixing multiple data sources, consider lowering these values to reduce memory usage.
# When using parquet, grain_num_threads is the number of files to read and interleave in parallel
grain_num_threads: 16
grain_prefetch_buffer_size: 500
grain_worker_count_eval: 1
grain_per_worker_buffer_size_eval: 1
grain_ram_budget_mb: 1024 # RAM budget (MB) for auto-tuning worker count. Only used when grain_worker_count is -1.
grain_num_threads_eval: 16
grain_prefetch_buffer_size_eval: 500
grain_data_source_max_workers: 16  # Max workers for ThreadPoolExecutor when mixing multiple Grain data sources.
grain_shuffle_buffer_size: 100 # shuffle buffer when using sequential access formats such as Parquet, TFRecord.
grain_use_elastic_iterator: false # For elastic training, set to this true and packing=false
# for using pathways
colocated_python_data_input: false  # experimental feature, under testing

# OLMo numpy pipeline (dataset_type=olmo_grain). Worker count, buffer size,
# and shuffle seed reuse grain_worker_count / grain_per_worker_buffer_size /
# data_shuffle_seed.
olmo_index_path: '' # JSON from tools/data_generation/build_olmo_npy_index.py
olmo_path_remap_from: '' # rewrite index paths starting with this prefix...
olmo_path_remap_to: ''   # ...to this one (e.g. gs://bucket/ -> /mnt/.../ for gcsfuse).
olmo_apply_ngram_filter: true # mask instances with repetitive n-grams (OLMo-core filter)

# Training loop
steps: 150_001 # If set to -1 then will inherit value from learning_rate_schedule_steps
log_period: 100 # The frequency of Tensorboard flush, gcs metrics writing, and managed profiler metrics updating.

jax_distributed_initialization_timeout: 300 # This is the default timeout in https://github.com/jax-ml/jax/blob/main/jax/_src/distributed.py
# Note there are two separate initializations - the jax coordination service (aka jax.distributed.initialize) and the backend (e.g. PjRT), the timeout above refers
# only to the jax coordination service.
jax_debug_log_modules: "" # Set this to "jax" to enable jax verbose logging such as for the jax coordination service initialization.
skip_jax_distributed_system: false # If true we will not initialize the jax distributed system.
# Currently the jax distributed is needed on cloud TPUs for async checkpointing.
# However when run on google internal TPUs the coordination service is started automatically
# and we should set this to true so we won't try to initialize a second time manually.

# Learning rate schedule structure depends on lr_schedule_type:
#
# Cosine schedule (lr_schedule_type='cosine'):
# Inspired by Llama2's learning rate schedule, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
# 1) Linear warmup from 0 to [learning_rate] over steps 0 to [learning_rate_schedule_steps * warmup_steps_fraction]
# 2) Cosine decay from [learning_rate] to [learning_rate * learning_rate_final_fraction] until learning_rate_schedule_steps
# 3) Constant learning rate of 0 from learning_rate_schedule_steps to steps (if steps > learning_rate_schedule_steps)
#
# WSD schedule (lr_schedule_type='wsd', Warmup-Stable-Decay):
# 1) Linear warmup from 0 to [learning_rate] over steps 0 to [learning_rate_schedule_steps * warmup_steps_fraction]
# 2) Stable phase at [learning_rate] for the majority of training
# 3) Decay from [learning_rate] to [learning_rate * learning_rate_final_fraction] over [learning_rate_schedule_steps * wsd_decay_steps_fraction] steps
#    The decay can be either linear or cosine based on wsd_decay_style
# 4) Constant learning rate of 0 from learning_rate_schedule_steps to steps (if steps > learning_rate_schedule_steps)
#
# The zero learning rate section can be used to more accurately measure the fully trained model's performance.
learning_rate: 3.e-5
lr_schedule_type: 'cosine'  # Options: 'cosine' or 'wsd'
learning_rate_final_fraction: 0.1  # Final LR as fraction of peak LR (applies to both cosine and WSD schedules)
wsd_decay_steps_fraction: 0.1  # Fraction of learning_rate_schedule_steps used for decay phase in WSD (e.g., 0.1 = 10%)
wsd_decay_style: 'linear'  # Decay style for WSD schedule: 'linear' or 'cosine'
warmup_steps_fraction: 0.1  # Fraction of learning_rate_schedule_steps used for warmup phase (applies to both schedules)
learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.
# However you may choose a longer schedule (learning_rate_schedule_steps > steps), in which case the training will end before
# dropping fully down. Or you may choose a shorter schedule, where the unspecified steps will have a learning rate of 0.

max_target_length: 2048 # Maximum sequence length
max_prefill_predict_length: 64 # Maximum length for the prefill when doing autoregression
prompt: "I love to" # Prompt for language model sampling.
load_from_prefill_dir: false # If true, decode.py doesn't "prefill" but just reads from directory
prefill_cache_dir: "" # If set and load_from_prefill_dir, decode.py reads from directory. If set, decode.py writes to directory
autoregressive_decode_assert: ""

# For nsys profiler, pass the training command to nsys command
# e.g. nsys profile -s none --force-overwrite true --capture-range=cudaProfilerApi --capture-range-end=stop {training command}
profiler: "" # Supported profiler: '', xplane, nsys
# If set to true, upload all profiler results from all hosts. Otherwise, only upload the profiler result from the first host.
upload_all_profiler_results: false
# Skip first n steps for profiling, to omit things like compilation and to give
# the iteration time a chance to stabilize.
skip_first_n_steps_for_profiler: 1
# Profile for a small number of steps to avoid a large profile file size.
profiler_steps: 5
hide_profiler_step_metric: false
profile_cleanly: true # If set to true, adds a block_until_ready on train state which aligns the profile for each step.
profile_periodically_period: -1 # If set to a positive integer, profile every profile_periodically_period steps.
# This is useful to debug scenarios where performance is changing.
enable_tpu_profiling_options: false
tpu_num_chips_to_profile_per_task: 1
tpu_num_sparse_core_tiles_to_trace: 1
tpu_num_sparse_cores_to_trace: 2

# Managed ML diagnostics settings. If the feature is enabled, it will
# - create a managed ML diagnostics run with all the MaxText configs
# - upload xplane profiling, if it is enabled.
# - upload training metrics, at the defined log_period interval.
managed_mldiagnostics: false # Whether to enable the managed diagnostics
managed_mldiagnostics_on_demand_profiling: true # Enable on-demand profiling server by default
managed_mldiagnostics_run_group: ""  # Optional. Used to group multiple runs.
managed_mldiagnostics_region: "" # Optional. GCP region for managed mldiagnostics. If empty, it will be auto-detected by the SDK.

# Dump HLO and jaxpr options
dump_hlo: false
dump_step: -1 # Dump modules at the given step if set to a positive integer.
dump_hlo_local_dir: "/tmp/xla_dump/"
dump_hlo_delete_local_after: true # Cleans local directory after its uploaded
dump_hlo_gcs_dir: "" # Defaults to {base_output_directory}/{run_name}/xla_dump
dump_hlo_local_module_name: "jit_train_step" # Filter saving modules locally by this string. Set to empty string to remove any filter.
dump_hlo_module_name: "jit_train_step" # Filter uploading modules by this string. Set to empty string to remove any filter.
dump_hlo_xla_flags: "" # Defaults to "--xla_dump_to={dump_hlo_local_dir} --xla_dump_hlo_module_re={dump_hlo_local_module_name} --xla_dump_large_constants"
dump_hlo_upload_all: false # If true all hosts dump HLO, false only jax.process_index()==0
# All hosts should have identical HLO for SPMD programs, however we have encountered some bugs
# where this is not the case and it is helpful to compare HLO across hosts.
dump_jaxpr: false
dump_jaxpr_local_dir: "/tmp/jaxpr_dump/"
dump_jaxpr_delete_local_after: true
dump_jaxpr_gcs_dir: "" # Defaults to {base_output_directory}/{run_name}/jaxpr_dump

# When dropout is false the model is a deterministic function of the
# data_shuffle_seed and init_weights_seed (i.e. reproducible losses)
enable_dropout: true
enable_data_shuffling: true
data_shuffle_seed: 0
init_weights_seed: 0

# DiLoCo params.
enable_diloco: false
diloco_sync_period: 36
diloco_outer_lr: 0.3
diloco_outer_momentum: 0.9

# You may disable clipping by setting gradient_clipping_threshold to zero.
gradient_clipping_threshold: 1.0

# Instead of updating the weights every step, you may effectively use a larger
# batch by accumulating the gradient over a set of steps.
gradient_accumulation_steps: 1

opt_type: "adamw"  # one of "adamw", "adam_pax", "sgd", or "muon"

# If true, skip the training step when loss or gradient spike is detected
# No updates for both weights and momentums (if applies)
skip_step_on_spikes: false
# The rolling interval to calculate the mean and standard deviation
skip_step_interval: 128
# The scaling factor to determine if a spike occurred
skip_step_scaling_factor: 6.0

# List of parameter names/patterns to train.
# If non-empty, all other parameters will be frozen. Example: ['.*indexer.*'].
# If empty (default), all parameters are trained.
trainable_parameters_mask: []

# AdamW optimizer parameters
# We use AdamW following Llama2's training details, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients.
adam_b2: 0.95 # Exponential decay rate to track the second moment of past gradients.
adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root.
adam_eps_root: 0. # A small constant applied to denominator inside the square root.
adam_weight_decay: 0.1 # AdamW Weight decay
adamw_mask: [] # List of parameter names/patterns to exclude from weight decay in AdamW, like ['bias', '.*norm', '.*ln.*'].
mu_dtype: "" # data type to store "mu" of AdamW tracking the first moment. Inherits from  weight_dtype if unset.
# Setting nu_dtype is not yet supported by optax, instead nu_dtype is always inherited from weights.
# See b/399961932 for more.

# Muon optimizer parameters
# https://github.com/google-deepmind/optax/blob/main/optax/contrib/_muon.py
# "mu_dtype", "adam_eps" are shared by AdamW
# "nesterov", "ns_coeffs", "ns_steps", "weight_decay_mask", "adaptive" use default
muon_beta: 0.95 # Decay rate for the exponentially weighted average of grads.
muon_weight_decay: 0 # Strength of the weight decay regularization. This is multiplied with the learning rate.
muon_consistent_rms: None # If None, apply width scaling to updates. If float, apply consistent rms scaling (recommend 0.2).


# Use iota operator in Embed
use_iota_embed: false
# use positional embedding
use_untrainable_positional_embedding: false
trainable_position_size: -1  # enable gpt3 position embedding with a positive trainable_position_size
# RoPE parameters
rope_type: "default" # one of "default", "llama3.1" or "yarn"
rope_linear_scaling_factor: 1.0 # linear scaling factor for "default" RoPE (see class `RotaryEmbedding` for more)
rope_use_scale: true # apply rope scaling for llama3.1 (see class `LLaMARotaryEmbedding` for more)
rope_min_timescale: 1
rope_max_timescale: 10_000 # Timescale For global Attention
local_rope_max_timescale: -1 # If positive used for local window Attention, otherwise `rope_max_timescale` is used for both local and global
global_rope_max_timescale: -1 # Timescale For global Attention (Gemma 4 specific)
global_rope_proportion: 0.25
local_rope_proportion: 1.0

# yarn RoPE parameters
max_position_embeddings: 163840
original_max_position_embeddings: 4096
rope_factor: 40
beta_fast: 32
beta_slow: 1
mscale: 1.0
rope_interleave: true # RoPE with sin/cos interleaved vs concatenated
rope_truncate: true # Floor lower bound and ceil upper bound for correction range
rope_attention_scaling: false # Scale the rotary embedding output

# Ahead of time Compilation (aka AOT)
# Only set these arguments if you are running train_compile or loading a compiled train step.
compiled_trainstep_file: "" # Name of saved serialized compiled train_step, e.g. compiled_train_v5e-256.pickle
compile_topology: '' # Target hardware version, e.g. 'v5e-256'
compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.

# MaxText Estimator configs
write_estimator_result: False

decode_sampling_strategy: "greedy" # decode_sampling_strategy should be one of greedy, weighted, nucleus, topk, or composite(top_k -> top_p -> weighted temperature)
decode_sampling_nucleus_p: -1 # set if you're doing nucleus / top-p
decode_sampling_top_k: 0 # set if you're doing top-k
decode_sampling_temperature: 1.

eval_interval: -1  # the specific number of train step between eval_step
eval_steps: -1  # run this number of steps for eval, recommend setting this to prevent error due to running out of evel data
target_eval_loss: 0.  # early stop once reaching target eval_loss
abort_on_nan_loss: true # Check for NaN and abort if found in training loss
abort_on_inf_loss: true # Check for Inf and abort if found in training loss

# Goodput parameters
enable_goodput_recording: false
monitor_goodput: false
goodput_upload_interval_seconds: 30
enable_pathways_goodput: false
monitor_step_time_deviation: true
step_deviation_interval_seconds: 30
enable_gcp_goodput_metrics: true
enable_gcp_step_deviation_metrics: true

# GCP workload monitoring
report_heartbeat_metric_for_gcp_monitoring: false
heartbeat_reporting_interval_in_seconds: 5
report_performance_metric_for_gcp_monitoring: false

enable_tensorboard: true

# Vertex AI Tensorboard Configurations - https://github.com/AI-Hypercomputer/maxtext/blob/main/docs/guides/use_vertex_ai_tensorboard.md
# Set to true for GCE, false if running via XPK
use_vertex_tensorboard: false
# Project to create Vertex AI Tensorboard in for GCE, blank if project is set using 'gcloud config set project'
# Set this to blank if running via XPK
vertex_tensorboard_project: ""
# Region to create Vertex AI Tensorboard in for GCE, blank if running via XPK
# Vertex AI supported regions: https://cloud.google.com/vertex-ai/docs/general/locations#available-regions
vertex_tensorboard_region: ""

# If set to true, MaxText will perform extra checks using jax.checkify. Note that this will effect performance.
max_checkify: false

# Inference
inference_microbenchmark_prefill_lengths: "64,128,256,512,1024"
inference_microbenchmark_stages: "prefill,generate"
inference_microbenchmark_loop_iters: 10
inference_microbenchmark_log_file_path: ""
inference_microbenchmark_num_samples: [1, 2, 3, 4, 5]
inference_metadata_file: "" # path to a json file
inference_server: "MaxtextInterleavedServer"  # inference server to start
prefill_slice: "v5e-16" # slice to use for prefill in disaggregation mode
generate_slice: "v5e-16" # slice to use for generatation in disaggregation mode
inference_benchmark_test: false
enable_model_warmup: false
enable_llm_inference_pool: false          # Bool to launch inference server for llm_inference_gateway with their specified APIs
multi_sampling: false
return_log_prob: false

# Stack prefill cache across the layer to reduce the
# Python layer latency.
stack_prefill_result_cache: false

# KV Cache layout control
# Logical layout: 0,1,2,3 ; CACHE_BATCH, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV
# Default layout: 1,2,0,3 ; CACHE_SEQUENCE, CACHE_HEADS, CACHE_BATCH, CACHE_KV
prefill_cache_axis_order: "1,2,0,3"
ar_cache_axis_order: "1,2,0,3"

# Compute layout control
# Default layout: 0,1,2,3 ; BATCH, LENGTH, HEAD, D_KV
# Currently only support compute layout: 0,1,2,3 and 0,2,1,3
compute_axis_order: "0,1,2,3"

reshape_q: false

# Maxengine Metrics
prometheus_port: 0

# Maxengine server
enable_jax_profiler: false
jax_profiler_port: 9999

# TPU power trace level for xprof. 0:POWER_TRACE_NONE, 1:POWER_TRACE_NORMAL, or 2:POWER_TRACE_SPI
xprof_tpu_power_trace_level: 0
xprof_e2e_enable_fw_throttle_event: false
xprof_e2e_enable_fw_power_level_event: false
xprof_e2e_enable_fw_thermal_event: false
profile_power_events: false # Set to true to enable TPU-specific power/thermal profiling events. Defaults to false to avoid breaking GPU xplane tracing.

log_config: true # Prints the config (after defaults have been set by pyconfig logic)
debug_sharding: false # Prints model weights sharding info

# Checkpoint Structured logging
enable_checkpoint_cloud_logger: false

# Single-controller
enable_single_controller: false

custom_mesh: "" # Available options: ['hybrid_ring_64x4', 'hybrid_ring_32x8']
# Split physical axes for https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.mesh_utils.create_device_mesh.html
allow_split_physical_axes: false
# Apply transformations to the mesh to optimize for TPU v6e
optimize_mesh_for_tpu_v6e: false

shardy: true # Whether to use shardy XLA backend (default in Jax starting 0.7.0), or GSPMD (to be fully deprecated ~2026)

remove_size_one_mesh_axis_from_type: true # Whether to remove size one mesh axis from type through jax.config.

use_ragged_attention: false
ragged_block_size: 256

### Splash attention block sizes
# These can be tuned for specific hardware generations, and can be set up to
# the model's sequence length.
sa_block_q: 512
sa_block_kv: 512
sa_block_kv_compute: 512
sa_block_q_dkv: 512
sa_block_kv_dkv: 512
sa_block_kv_dkv_compute: 512
sa_block_q_dq: 512
sa_block_kv_dq: 512
sa_use_fused_bwd_kernel: false
sa_q_layout: "HEAD_DIM_MINOR"
sa_k_layout: "HEAD_DIM_MINOR"
sa_v_layout: "HEAD_DIM_MINOR"
use_max_logit_estimate: -1 # -1 means no estimate, any > 0 value will be used as max logit estimate
cost_estimate_flops_fwd: -1 # -1 means using splash default cost estmiation, any >= 0 value will be used as cost estmiation for splash to overlap for communication (forward)
cost_estimate_flops_bwd: -1 # -1 means using splash default cost estmiation, any >= 0 value will be used as cost estmiation for splash to overlap for communication (backward)
dq_reduction_steps: 0 #the number of reduction steps. For now, only 3 or all the kv steps are supported.
use_splash_scheduler: false # to use tokamax splash attention scheduler.
### Determine if we want to use load balance for context parallelism
context_parallel_load_balance: true
context_parallel_strategy: "all_gather" # "all_gather" or "ring"
context_parallel_reorder_strategy: "auto" # "auto", "dual_chunk_swap", or "striped"


# Chunked Prefill Parameters
prefill_chunk_size: 256
use_chunked_prefill: false

# Prefix Caching parameters in jetstream
enable_prefix_caching: false
prefix_caching_hbm_byte: 10_000_000_000 # 10 GB
prefix_caching_dram_byte: 100_000_000_000 # 100 GB

# This is a temporary flag that will be removed soon after the fix lands in TE
enable_padding_causal_mask: true

# Llama4-specific
# Whether to apply Query/Key normalization.
# NOTE: non-Llama4 models use RMSNorm before RoPE
# whereas Llama4 models use L2Norm after RoPE
use_qk_norm: false
# Every `X` layers will NOT use RoPE
nope_layer_interval: -1
# Every `X` layers is MoE layer
interleave_moe_layer_step: 1
# dynamically scale the attention temperature for each query token based on sequence length
# Recommended for long sequences (e.g., >32k tokens) to maintain stable output results
# See (https://arxiv.org/abs/2501.19399) for more details
temperature_tuning: false

# Multimodal flags
use_multimodal: false
use_audio: false
freeze_vision_encoder_params: true
freeze_audio_encoder_params: true
dtype_mm: "float32"  # Data type for multimodal model's vision encoder
remat_policy_for_vit: "minimal"  # Remat policy for multimodal model's vision encoder. Check `remat_policy` for options.
image_size_for_vit: 896 # Default for Gemma3, and should be overwritten by model's config
image_path: "" # Local image path used for decoding, can be multiple paths separated by comma, exp "/path/image1.jpg,/path/image2.jpg"
video_path: "" # Local video path used for decoding, can be multiple paths separated by comma, exp "/path/video1.mp4,/path/video2.mp4"
audio_path: "" # Local audio path used for decoding, can be multiple paths separated by comma, exp "/path/audio1.wav,/path/audio2.wav"
image_placeholder: "<|image|>"
video_placeholder: "<|video|>"
audio_placeholder: "<|audio|>"
use_audio_in_video: false
posemb_type_for_vit: "learn"
# max_num_images_per_example only applies for training when your image column is a list of images.
# -1 means no limit, and will pad to the max possible number of images determined by sequence length.
# Set it to avoid unnecessary padding if you know the maximum number of images per example.
max_num_images_per_example: -1
vision_output_length: -1 # The output length (number of soft tokens) from vision encoder, used in Gemma4.

### llama4 multi modal configs
hidden_size_for_vit: 1408
intermediate_size_for_vit: 5632
num_attention_heads_for_vit: 16
num_channels_for_vit: 3
tile_size_for_vit: 336
patch_size_for_vit: 14
conv_stride_for_vit: 14
num_hidden_layers_for_vit: 34
projector_input_dim_for_vit: 4096
projector_output_dim_for_vit: 4096
rope_theta_for_vit: 10000
vision_output_dim_for_vit: 4096
pixel_shuffle_ratio_for_vit: 0.5
projector_dropout_for_vit: 0.0

# Qwen3-OmniMoe vision encoder
spatial_merge_size_for_vit: 2
out_hidden_size_for_vit: 512
temporal_patch_size_for_vit: 2
num_position_embeddings_for_vit: 1024
deepstack_visual_indexes_for_vit: []

### Audio encoder configs (Qwen3-OmniMoe)
d_model_for_audio: 256
encoder_attention_heads_for_audio: 4
encoder_ffn_dim_for_audio: 512
encoder_layers_for_audio: 2
attention_dropout_for_audio: 0.0
activation_dropout_for_audio: 0.0
activation_function_for_audio: "gelu"
num_mel_bins_for_audio: 128
max_source_positions_for_audio: 1500
scale_embedding_for_audio: true
n_window_for_audio: 50
n_window_infer_for_audio: 800
conv_chunksize_for_audio: 500
downsample_hidden_size_for_audio: 256
output_dim_for_audio: 512
num_conv_layers_for_audio: 3
max_timescale_for_audio: 10000.0
max_sample_len_for_audio: 10000

use_mrope: false
mrope_section: [24, 20, 20]
position_id_per_seconds: 25

# Subslice shape in the form of "x,y,z" when using pathways (single controller).
# Example: "8,8" to use a 8x8 subgrid (64 chips) of a full pod (16x16) of trillium.
subslice_shape: ""

# NNX
enable_nnx: true
pure_nnx_decoder: true
pure_nnx: true

################################## Qwen3-Next Specific Configs ##################################
# Kernel size for the 1D convolution in the Gated Delta Net
gdn_conv_kernel_dim: 4
# Head dimension for the key/query in the Gated Delta Net
gdn_key_head_dim: 128
# Head dimension for the value in the Gated Delta Net
gdn_value_head_dim: 128
# Number of key/query heads in the Gated Delta Net
gdn_num_key_heads: 16
# Number of value heads in the Gated Delta Net
gdn_num_value_heads: 32
# Chunk size for the parallel scan algorithm in the Gated Delta Net.
gdn_chunk_size: 64
# Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.
use_qk_norm_in_gdn: true
# The ratio of dimension to apply ROPE on
partial_rotary_factor: 1.0

# Use tokamax library for gmm kernel implementation
use_tokamax_gmm: false
# Whether to use GMM v2 for MoE.
# Requires use_tokamax_gmm: true. Currently incompatible with quantization.
use_gmm_v2: false
use_tokamax_splash: false
# Setting this flag will use a non-pallas implementation.
use_jax_splash: false

# vLLM Adapter Configurations
# Path to the HuggingFace-style config directory for the adapter (e.g. src/maxtext/integration/vllm/maxtext_vllm_adapter)
vllm_hf_config_path: ""
# A JSON string of overrides to apply to the HuggingFace-style config for the vLLM adapter.
# This can be used to override specific settings without modifying the original config file.
vllm_hf_overrides: {}
# JSON string containing additional configuration for the vLLM model (e.g. '{"maxtext_config": {...}}')
vllm_additional_config: {}
# When use_jax_splash=true, force the layout of the query tensor to be [..., NUM_HEADS, HEAD_DIM, SEQ_LENGTH]
force_q_layout: false

################################## DeepSeek Manifold-Constrained Hyper Connections (mHC) ##################################
# The number of parallel streams in Hyper Connection.
mhc_expansion_rate: 1
# The number of iterations for the Sinkhorn-Knopp algorithm.
sinkhorn_iterations: 20
# mHC-lite: https://openreview.net/pdf?id=5IJX6kvOif
# - Controls whether to generate the MHC doubly stochastic matrix via
#   permutation-based convex combination rather than Sinkhorn-Knopp.
# - Practical only for a small mhc_expansion_rate (e.g., k=4), as the number
#   of permutation matrices scales factorially (k!).
# - Ideally, this should be used whenever expansion rate is small as it removes
#   the expensive sinkhorn iterations, the downside to this approach is that
#   it is factorial in k.
enable_mhc_lite: False

################################## DeepSeek Engram ##################################
# Indices of transformer layers where Engram are integrated; leave empty [] to disable.
# Example: [1, 4] attaches to the 2nd and 5th layer.
engram_layers: []
# The max 'n' in N-gram. Example: n=3 means it covers both 2-grams and 3-grams.
engram_max_ngram_size: 3
# Number of heads dedicated to the Engram.
engram_num_heads: 8
# Head dimension for heads.
engram_head_dim: 1280
# List of minimum head vocab sizes for each n-gram order.
engram_vocab_bases: []
# Temporal window size for Engram convolution.
engram_kernel_size: 4
# The seed for Engram hash mapping.
engram_seed: 0

##### Distillation parameters
distill_alpha: 0.5
distill_temperature: 1.0
# distill_beta is used for cosine similarity loss between intermediate activataitions of out_proj in teacher/student models.
# 0.0 value disables this feature.
distill_beta: 0.0
# distill_feature_loss_type is the type of loss to use for feature distillation ("cosine" or "l2").
distill_feature_loss_type: "cosine"
distill_layer_indices: None
# Dynamic loss weight scheduling: set *_end to a target value and *_schedule to "linear" or "cosine".
# When *_end is None (default), the corresponding weight stays fixed throughout training.
distill_alpha_end: None
distill_alpha_schedule: "constant"
distill_temperature_end: None
distill_temperature_schedule: "constant"
distill_beta_end: None
distill_beta_schedule: "constant"

##### Elastic training parameters
# Elastic training is Pathways-specific and does not work on McJAX.
elastic_enabled: false
elastic_timeout_seconds: 300
elastic_max_retries: 10
elastic_min_slice_count: -1