Heuristic:Allenai Open instruct BFloat16 Training
| Knowledge Sources | |
|---|---|
| Domains | Optimization, Deep_Learning |
| Last Updated | 2026-02-07 00:00 GMT |
Overview
Always use bfloat16 precision for model weights and training to maximize stability and throughput.
Description
Open Instruct hardcodes bfloat16 (bf16) as the training precision across all training pipelines. Unlike float16, bfloat16 has the same exponent range as float32, preventing overflow/underflow issues that are common in large-scale training. The training throughput is nearly identical to float16 while being significantly more numerically stable.
Usage
Apply this heuristic for all training and inference in Open Instruct. Both the policy model and reference model use bf16. The DeepSpeed config also specifies bf16 mode.
The Insight (Rule of Thumb)
- Action: Set `dtype=torch.bfloat16` when loading models and `bf16=True` in DeepSpeed config.
- Value: bfloat16 (not float16 or float32).
- Trade-off: Slight precision loss compared to float32, but negligible impact on model quality. Requires Ampere (A100) or newer GPU architecture.
Reasoning
Float16 has a narrow exponent range (5 bits) that causes gradient underflow/overflow, especially in large models with mixed loss scales (policy loss + KL penalty). Bfloat16 has 8 exponent bits (same as float32), making it immune to these issues. The 7-bit mantissa (vs float16's 10-bit) has no measurable impact on final model quality for language models.
Code Evidence
Model loading from `open_instruct/grpo_fast.py:252`:
self.policy: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
model_config.model_name_or_path,
revision=model_config.model_revision,
dtype=torch.bfloat16,
...
)
DeepSpeed config from `open_instruct/grpo_fast.py:224`:
ds_config = get_train_ds_config(
...
bf16=True,
...
)
TF32 matmul enabled in `open_instruct/dpo_utils.py:303`:
torch.backends.cuda.matmul.allow_tf32 = True