Skip to content

maxtext-parameters.md

Latest commit

 

History

History
148 lines (91 loc) · 10.4 KB

maxtext-parameters.md

File metadata and controls

148 lines (91 loc) · 10.4 KB
Raw

MaxText backend configuration reference

Primus routes experiment YAML into the MaxText stack (JAX / XLA). Configuration is a flat map of keys (no nested training.* trees like TorchTitan): Primus merges module and model presets, writes a temporary YAML, and MaxText’s pyconfig.initialize loads it on top of upstream defaults.

The Primus overlay keeps base_config: "base.yml" so MaxText loads its own configs/base.yml at runtime. This page lists Primus-defined defaults and commonly overridden Primus fields. For the full upstream parameter set (hundreds of keys), see the MaxText documentation and upstream base.yml.

How parameters flow

  1. YAML presets under primus/configs/modules/maxtext/ and primus/configs/models/maxtext/ are merged with CLI overrides.
  2. MaxTextAdapter.convert_config passes the merged namespace through MaxTextConfigBuilder (currently a thin pass-through).
  3. export_params_to_yaml writes a flat YAML file; MaxText ignores unknown Primus-private keys via pydantic filtering.
  4. Unknown keys from upstream still resolve through environment overrides inside MaxText (pyconfig), not shown here.

1. Base module parameters

Shared with all Primus modules via module_base.yaml and trainer extensions.

Parameter Default (Primus) Description
trainable true in trainer_base.yaml (overrides module_base’s false) When true, the module participates in training orchestration.
sink_level null Structured logging sink level for the module (if the logging stack is configured to use it).
file_sink_level DEBUG File sink verbosity.
stderr_sink_level INFO Stderr sink verbosity.

2. Training

From pre_trainer.yaml (extends trainer_base.yaml).

Parameter Default Description
base_config "base.yml" Upstream MaxText base file loaded by pyconfig._load_config.
hardware "gpu" Hardware target string consumed by MaxText.
steps 1000 Global optimizer steps for the run.
log_period 100 Steps between log emissions.

3. Data

Parameter Default Description
dataset_type "hf" Dataset backend selector (Hugging Face in the default path).
hf_path "allenai/c4" Hugging Face dataset repo or identifier.
hf_data_dir "en" Subdirectory / config slice within the HF dataset.
hf_train_files "" Optional explicit train file list (format per MaxText HF loader).
packing true Sequence packing for efficiency when supported by the data pipeline.

4. Checkpointing

These are Primus overlay defaults. MaxText also loads upstream base.yml at runtime through base_config: "base.yml", where upstream checkpoint defaults might differ. When debugging effective behavior, distinguish the Primus YAML written by the adapter from the upstream MaxText defaults loaded afterward.

Parameter Default Description
enable_checkpointing false See Training section.
async_checkpointing false When enable_checkpointing is true, use async checkpoint workers.

5. Profiling

Parameter Default Description
profiler "xplane" Profiler backend (e.g. XPlane for JAX).
skip_first_n_steps_for_profiler 3 Warmup steps excluded from capture.
profiler_steps 1 Number of steps to profile once active.

6. Memory and recomputation

Parameter Default Description
remat_policy 'full' Activation rematerialization policy (none, minimal, full, etc.—see MaxText).
optimizer_memory_host_offload false Offload optimizer state to host memory when supported.
scan_layers true Use scanned layer implementation where applicable.
param_scan_axis 1 Axis for parameter scanning / partitioning layout.

7. Precision and quantization

Parameter Default Description
dtype "bfloat16" Default compute dtype for many ops.
quantization "" Quantization mode string (empty = none; set per MaxText AQT recipes).
quantize_kvcache false Quantize KV cache tensors.
kv_quant_axis "heads_and_dkv" KV quantization axis naming for kernels.
kv_quant_dtype "int8" Storage dtype for KV cache when quantization is on.
weight_dtype bfloat16 Weight storage / compute dtype for non-quantized paths.
checkpoint_is_quantized false Set true when loading an AQT-quantized checkpoint.
logits_dot_in_fp32 false Compute logits matmul in float32 for numerical stability.

fp8 + MoE (v26.6): MoE fp8 configs (e.g. mixtral_8x7B-fp8, qwen3_30B_A3B-fp8) must set pure_nnx_decoder: false. v26.6 defaults to true, which invokes the legacy Linen Fp8Einsum without a binding scope. See Advanced. Dense fp8 configs are unaffected.

8. Model

From model_base.yaml and per-model files such as llama3_8B.yaml.

Parameter Default Description
model_name "default" in model_base; e.g. "llama3-8b" in llama3_8B.yaml Selects MaxText’s bundled model YAML when present.
override_model_config true When true, CLI / kwargs override values from the loaded model config.
attention "autoselected" Attention implementation. gemma4 requires cudnn_flash_te: its local layers use sliding-window attention, and only cudnn_flash_te passes the window to the kernel (as window_size). autoselected resolves to a pallas kernel that builds a plain causal mask and silently drops the window.
use_iota_embed true Use iota-based embedding for performance on accelerator backends.
tokenizer_path e.g. "meta-llama/Meta-Llama-3-8B" Hugging Face tokenizer id or local path.

9. Advanced

Parameter Default Description
shardy false Enable Shardy-related integration in MaxText when building shardings.
pure_nnx_decoder true Run the decoder as pure Flax NNX. As of v26.6 the base config defaults this to true. fp8 MoE configs must override this to false — the legacy per-module Linen Fp8Einsum (used by the MoE sparse-matmul quant path) needs an active Linen binding scope, and the pure-NNX decoder invokes it unbound, crashing at the first step. See the note in Precision and quantization.

Related reading

  • MaxText documentation—full parameter reference and recipes.
  • Primus implementation: primus/backends/maxtext/argument_builder.py, maxtext_pretrain_trainer.py, maxtext_adapter.py.