🎰 Deterministic `flash-attn`

Sam Foreman

🎰 Deterministic `flash-attn`

Author

Affiliation

Sam Foreman

ALCF

Published

June 17, 2024

[NOTE]: For additional details, refer to the W&B Report.

Simple tests to confirm the loss is exactly reproducible across independent runs (when launched with the same seed).

In particular, we set:
```
output = flash_attn_func(q, k, v, None, self.causal, deterministic=True)
```
in all the flash_attn_func(...) calls from megatron/model/transformer.py

All experiments ran on Polaris @ ALCF, using:

machine: Polaris
args.zero_stage: 1
args.num_layers: 32
args.micro_batch_size: 1
args.optimizer: "adamw"
args.use_flash_attn: true
env.DFL_STEM: "books"
env.GRAD_ACC_STEPS: 8
env.WORLD_SIZE: 8

Figure 1: Plot of the loss curve for 3 independent runs with `deterministic=True`

Citation

BibTeX citation:

@online{foreman2024,
  author = {Foreman, Sam},
  title = {🎰 {Deterministic} `Flash-Attn`},
  date = {2024-06-17},
  url = {https://samforeman.me/posts/AuroraGPT/determinstic-flash-attn/},
  langid = {en}
}

For attribution, please cite this work as:

Foreman, Sam. 2024. “🎰 Deterministic `Flash-Attn`.” June 17, 2024. https://samforeman.me/posts/AuroraGPT/determinstic-flash-attn/.