M+Adam: Low-Precision Pre-Training

Paper | Slides

This repository supports pre-training language models on the C4 dataset with M+Adam, an additive-multiplicative optimizer designed for persistent low-precision training.

M+Adam combines:

• an additive Adam-style branch for local refinement, sign changes, and zero revival; and
• a multiplicative exponent-space branch for scale-aware updates that remain effective when low-precision spacing becomes coarse.

The repository supports:

• BF16 master weights with BF16 compute;
• FP8 master-weight experiments with BF16 compute;
• FP8 master-weight experiments with Transformer Engine FP8 GEMMs;
• distributed training with torchrun;
• FlashAttention 2 and Liger Kernel;
• Weights & Biases logging;
• LLaMA configurations from 60M to larger model sizes.

---

Repository Structure

.
├── main.py
│   └── Main entry point for BF16 master-weight training
│
├── train_fp8_true_master.py
│   └── Entry point for persistent FP8 master-weight experiments
│
├── fp8_block_ops.py
│   └── Blockwise FP8 quantization and dequantization utilities
│
├── optimizer_torch/
│   ├── __init__.py
│   └── m_adam.py
│       └── M+Adam optimizer implementation
│
├── llama_configs/
│   ├── llama_60m.json
│   ├── llama_130m.json
│   ├── llama_350m.json
│   └── ...
│
├── download_c4.py
│   └── C4 download/preparation utility
│
├── scripts/
│   └── m+Adam.sh
│       └── Example BF16 M+Adam launcher
│
├── requirements.txt
└── README.md

Which files should I use?

Training mode	Main file	Supporting files
BF16 weights + BF16 compute	main.py	optimizer_torch/m_adam.py
FP8 master weights + BF16 compute	train_fp8_true_master.py	optimizer_torch/m_adam.py, fp8_block_ops.py
FP8 master weights + FP8 compute	train_fp8_true_master.py	optimizer_torch/m_adam.py, fp8_block_ops.py, Transformer Engine
Example BF16 launcher	scripts/m+Adam.sh	Calls main.py

---

Environment Setup

Install the base dependencies:

pip install -r requirements.txt

The project has been used with a Python 3.10, CUDA 12.x, and PyTorch 2.5.x environment. Exact package versions may need to be adjusted for the CUDA toolkit and GPU available on your system.

Optional acceleration packages:

• flash-attn for FlashAttention 2;
• liger-kernel for Liger kernels;
• transformer-engine for hardware FP8 GEMMs.

The BF16 experiments do not require Transformer Engine. Hardware FP8 compute requires a Hopper-class GPU such as H100 or H200.

---

Preparing C4

The training programs accept local Hugging Face datasets saved with datasets.save_to_disk.

Use the repository utility as the starting point:

python download_c4.py --help

After preparing the train and validation datasets, define their paths:

export TRAIN_DATA=/path/to/c4_train
export VAL_DATA=/path/to/c4_validation

Both training entry points accept:

--local_path_train
--local_path_val

---

Important M+Adam Learning-Rate Mapping

The current implementation uses the following command-line mapping:

M+Adam branch	Paper notation	Command-line argument
Additive Adam branch	eta_a	--lr
Multiplicative/exponent branch	eta_m	--m_adam_lr_e

For the 60M, 1x Chinchilla BF16 configuration:

eta_a = 8e-5
eta_m = 0.004

the correct command-line arguments are:

--lr 8e-5 \
--m_adam_lr_e 0.004

Do not swap these two values.

---

BF16 Master-Weight Training

Use main.py for BF16 weights with BF16 compute.

The following example trains LLaMA-60M for a 1x Chinchilla budget with two GPUs and a global batch size of 512:

torchrun --standalone --nproc_per_node=2 main.py \
    --model_config llama_configs/llama_60m.json \
    --local_path_train "${TRAIN_DATA}" \
    --local_path_val "${VAL_DATA}" \
    --batch_size 256 \
    --total_batch_size 512 \
    --max_length 256 \
    --num_training_steps 10000 \
    --scheduler_effective_training_steps 0 \
    --optimizer m_adam \
    --dtype bfloat16 \
    --lr 8e-5 \
    --m_adam_lr_e 0.004 \
    --m_adam_beta1 0.87 \
    --m_adam_beta2 0.98 \
    --m_adam_eps 1e-20 \
    --m_adam_weight_decay_m 0.1 \
    --m_adam_weight_decay_e 0.01 \
    --max_grad_norm 1.0 \
    --warmup_steps 2000 \
    --scheduler linear \
    --min_lr_ratio 0.0 \
    --m_adam_norm_type none \
    --m_adam_sched_e logcosine \
    --m_adam_total_steps_e 10000 \
    --m_adam_warmup_steps_e 2000 \
    --m_adam_min_lr_ratio_e 0.0 \
    --m_adam_logcosine_alpha_e 3.0 \
    --m_adam_tie_e_to_m false \
    --eval_every 10000 \
    --num_eval_tokens 10000000 \
    --save_every 10000 \
    --save_dir checkpoints/m_adam_bf16_60m_1x \
    --workers 1 \
    --seed 0 \
    --use_flash_attention \
    --use_liger_kernel \
    --wandb_project "M+Adam" \
    --wandb_entity "YOUR_WANDB_ENTITY" \
    --wandb_tag "60M-mAdam-1xChinchilla-BF16" \
    --name "m_adam_bf16_60m_1x"

Remove the following flags when the corresponding packages are unavailable:

--use_flash_attention
--use_liger_kernel

The convenience launcher in scripts/m+Adam.sh can also be used as a starting point for BF16 experiments.

---

FP8 Master-Weight Training

Use train_fp8_true_master.py for FP8 master-weight experiments.

The FP8 path uses fp8_block_ops.py to block-quantize selected LLaMA linear weights. Attention and MLP projection weights are included, while embedding, normalization, and language-model-head weights remain outside the FP8 storage path.

After every optimizer update, the selected weights are quantized blockwise and dequantized to the compute dtype for the next forward and backward pass.

FP8 Master Weights with BF16 Compute

This mode keeps model computation in BF16 while persistently quantizing selected weights to FP8 after each update.

torchrun --standalone --nproc_per_node=2 train_fp8_true_master.py \
    --model_config llama_configs/llama_60m.json \
    --tokenizer hf-internal-testing/llama-tokenizer \
    --local_path_train "${TRAIN_DATA}" \
    --local_path_val "${VAL_DATA}" \
    --batch_size 256 \
    --total_batch_size 512 \
    --max_length 256 \
    --num_training_steps 10000 \
    --scheduler_effective_training_steps 0 \
    --optimizer m_adam \
    --dtype bfloat16 \
    --lr 8e-5 \
    --m_adam_lr_e 0.004 \
    --m_adam_beta1 0.87 \
    --m_adam_beta2 0.98 \
    --m_adam_eps 1e-20 \
    --m_adam_weight_decay_m 0.1 \
    --m_adam_weight_decay_e 0.01 \
    --max_grad_norm 1.0 \
    --warmup_steps 2000 \
    --min_lr_ratio 0.0 \
    --m_adam_norm_type none \
    --m_adam_sched_e logcosine \
    --m_adam_total_steps_e 10000 \
    --m_adam_warmup_steps_e 2000 \
    --m_adam_min_lr_ratio_e 0.0 \
    --m_adam_logcosine_alpha_e 3.0 \
    --m_adam_tie_e_to_m false \
    --fp8_true_master \
    --fp8_format e5m2 \
    --fp8_block_size 256 \
    --fp8_pow2_scale \
    --fp8_adam_state_dtype bfloat16 \
    --eval_every 10000 \
    --num_eval_tokens 10000000 \
    --check_nonfinite_every 1 \
    --plot_every 200 \
    --workers 1 \
    --seed 0 \
    --use_flash_attention \
    --use_liger_kernel \
    --out_dir checkpoints/m_adam_fp8master_60m_1x \
    --run_name m_adam_fp8master_60m_1x \
    --wandb_project "M+Adam" \
    --wandb_entity "YOUR_WANDB_ENTITY" \
    --wandb_tag "60M-mAdam-1xChinchilla-FP8Master-BF16Compute" \
    --wandb_name "m_adam_fp8master_60m_1x"

FP8 format

Available formats:

e4m3
e5m2

Example:

--fp8_format e5m2

FP8 block size

The default experiment configuration uses:

--fp8_block_size 256

Optimizer-side dtype

The optimizer-side shadow parameter and M+Adam states can use BF16 or FP32:

--fp8_adam_state_dtype bfloat16

or:

--fp8_adam_state_dtype float32

Use BF16 for lower state memory and FP32 when additional numerical stability is preferred.

---

FP8 Master Weights with Hardware FP8 Compute

To use FP8 master weights together with Transformer Engine FP8 GEMMs, add the following flags to the previous command:

--te_fp8_gemm \
--te_recipe hybrid_delayed \
--te_amax_history 16 \
--te_margin 0

A complete command is:

torchrun --standalone --nproc_per_node=2 train_fp8_true_master.py \
    --model_config llama_configs/llama_60m.json \
    --tokenizer hf-internal-testing/llama-tokenizer \
    --local_path_train "${TRAIN_DATA}" \
    --local_path_val "${VAL_DATA}" \
    --batch_size 256 \
    --total_batch_size 512 \
    --max_length 256 \
    --num_training_steps 10000 \
    --scheduler_effective_training_steps 0 \
    --optimizer m_adam \
    --dtype bfloat16 \
    --lr 8e-5 \
    --m_adam_lr_e 0.004 \
    --m_adam_beta1 0.87 \
    --m_adam_beta2 0.98 \
    --m_adam_eps 1e-20 \
    --m_adam_weight_decay_m 0.1 \
    --m_adam_weight_decay_e 0.01 \
    --max_grad_norm 1.0 \
    --warmup_steps 2000 \
    --min_lr_ratio 0.0 \
    --m_adam_norm_type none \
    --m_adam_sched_e logcosine \
    --m_adam_total_steps_e 10000 \
    --m_adam_warmup_steps_e 2000 \
    --m_adam_min_lr_ratio_e 0.0 \
    --m_adam_logcosine_alpha_e 3.0 \
    --m_adam_tie_e_to_m false \
    --fp8_true_master \
    --fp8_format e5m2 \
    --fp8_block_size 256 \
    --fp8_pow2_scale \
    --fp8_adam_state_dtype bfloat16 \
    --te_fp8_gemm \
    --te_recipe hybrid_delayed \
    --te_amax_history 16 \
    --te_margin 0 \
    --eval_every 10000 \
    --num_eval_tokens 10000000 \
    --check_nonfinite_every 1 \
    --plot_every 200 \
    --workers 1 \
    --seed 0 \
    --use_flash_attention \
    --use_liger_kernel \
    --out_dir checkpoints/m_adam_fp8master_fp8compute_60m_1x \
    --run_name m_adam_fp8master_fp8compute_60m_1x \
    --wandb_project "M+Adam" \
    --wandb_entity "YOUR_WANDB_ENTITY" \
    --wandb_tag "60M-mAdam-1xChinchilla-FP8Master-FP8Compute" \
    --wandb_name "m_adam_fp8master_fp8compute_60m_1x"

Hardware FP8 compute requires Transformer Engine and a Hopper-class GPU. The training script checks the GPU compute capability before enabling this mode.

Evaluation remains in BF16 by default. To evaluate under Transformer Engine FP8 autocast as well, add:

--te_fp8_eval

---

Scaling to Other Model Sizes

Select another model configuration through --model_config:

--model_config llama_configs/llama_130m.json

or:

--model_config llama_configs/llama_350m.json

When changing model size, token budget, precision regime, or global batch size, the learning rates, warmup, epsilon, and weight decay may need to be re-tuned.

The 60M configuration above is intended as a reproducible starting point, not a universal setting for every scale.

---

Distributed Training

The effective global batch size is:

batch_size x number_of_GPUs x gradient_accumulation

For example:

batch_size = 256
number_of_GPUs = 2
total_batch_size = 512

gives:

gradient_accumulation = 1

Run on four GPUs by changing:

--nproc_per_node=4

and choose a compatible --total_batch_size.

---

Outputs and Logging

main.py

The BF16 training path supports:

• model checkpoints through --save_dir;
• checkpoint frequency through --save_every;
• validation through --eval_every;
• W&B logging through --wandb_project, --wandb_entity, and --name.

train_fp8_true_master.py

The FP8 master-weight path stores:

• argument configuration;
• training and evaluation JSONL metrics;
• loss plots;
• W&B metrics;
• run-specific outputs under --out_dir.

---

Notes

1. --lr is the additive-branch learning rate.
2. --m_adam_lr_e is the multiplicative/exponent-branch learning rate.
3. FP8 master-weight storage and FP8 matrix multiplication are separate options.
4. --fp8_true_master enables persistent FP8 weight quantization.
5. --te_fp8_gemm enables Transformer Engine FP8 compute.
6. --te_fp8_eval is optional; evaluation is BF16 by default.
7. Remove FlashAttention or Liger flags when those packages are not installed or unsupported by the target GPU.

---

Acknowledgements

This repository builds on components and ideas from:

• Hugging Face Transformers
• Hugging Face Datasets
• FlashAttention
• Liger Kernel
• NVIDIA Transformer Engine
• GaLore