MaxText backend configuration reference

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#

YAML presets under primus/configs/modules/maxtext/ and primus/configs/models/maxtext/ are merged with CLI overrides.
MaxTextAdapter.convert_config passes the merged namespace through MaxTextConfigBuilder (currently a thin pass-through).
export_params_to_yaml writes a flat YAML file; MaxText ignores unknown Primus-private keys via pydantic filtering.
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.

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`	`"cudnn_flash_te"`	Attention implementation (Primus default favors TE flash on AMD GPUs).
`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.