📸 `flash-attn` on Sunspot

Sam Foreman 2024-06-17

• Update: 2024-06-16
• 🐛 Impact on Loss [Bug?]
• 📦 LLM Framework Release
  • 📸 flash 🤝 📷 no-flash
  • 🚧 Broken MPI
• 🕵🏻‍ Framework Comparison
  • 🥸 Fix in Disguise
  • ✅ 2024.0 Fix
  • 📊 lr-decay-iters Comparison
• 📈 lr-decay-iters dependence
• 🏎️ Performance Improvement in 2024.1

Update: 2024-06-16

After an interactive debug session with Intel, the root behavior of the apparent discrepancy was identified.

In particular, we found that the ALCF/Megatron-DeepSpeed repo was NOT explicitly setting the dropout values to 0.0 (and so, was using the default values of 0.1) for both --attention-dropout and --hidden-dropout.

After making this change, the losses were observed to agree, as can be seen below in

Figure 1: After correctly setting the dropout values, the loss curves were observed to agree.

🐛 Impact on Loss [Bug?]

In the q4-drop, it was observed that toggling flash-attn on / off seemed to produce different loss curves (with otherwise identical configs)

TP: 1
PP: 1
GAS: 1
OPT: adamw
dtype: bf16
NLAYERS: 10
MICRO_BATCH: 2
WORLD_SIZE: 24

This can be seen clearly in the figure below:

This was identified, and to be addressed in upcoming release.

📦 LLM Framework Release

On 05/14/2024, Intel dropped their new LLM frameworks release:

📸 flash 🤝 📷 no-flash

With this new release, Intel observed that the loss curves agreed exactly for flash / no-flash, using the learning rate settings below:

lr: 0.00015
lr_warmup_frac: 0.01
lr_decay_iters: 320000

Testing with Jerome’s new release:

module use -a /home/jmitche1/anl_release/2024/q2
module load frameworks_2024_5_v2

I was able to independently confirm these results, shown in 📸 flash 🤝 📷 no-flash below.

📸 flash vs. 📷 no-flash

🚧 Broken MPI¹

For whatever reason, things seemed to have spontaneously broken on the night of 2024-04-14 ??

When trying to run experiments the following day (05/15/2024) I was met with this²:

Abort(15): Fatal error in internal_Init_thread: Other MPI error

which was discussed further in this thread on slack.

It seems Subrata also encountered a similar issue [see: slack thread]

Abort(15): Fatal error in internal_Init_thread: Other MPI error

module use -a /home/jmitche1/anl_release/2024/q2
module load frameworks_2024_5_v2
module use /home/ftartagl/graphics-compute-runtime/modulefiles
module load graphics-compute-runtime/agama-ci-devel-803.29
module load spack-pe-gcc/0.6.1-23.275.2 gcc/12.2.0
module use /soft/preview-modulefiles/24.086.0
module load oneapi/release/2024.04.15.001

🕵🏻‍ Framework Comparison

As I was re-building MPI, and after talking to Jerome, I realized that most of the dependencies are already present in the provided frameworks/ modules on Sunspot.

As a simple test, I tried building a new environment built on the base conda environment³ provided by theframeworks/2023.12.15.001 module, which worked without modification and had ) most of what I needed already installed:

>>> import torch
>>> torch.__version__
'2.1.0a0+cxx11.abi'
>>> import intel_extension_for_pytorch as ipex
>>> ipex.__version__
'2.1.10+xpu'
>>> from mpi4py import MPI

The remaining dependencies were installed according to the instructions from the new release frameworks_2024_5_v2.

Details included below.

$ module use frameworks/2023.12.15.001
$ export PBS_O_WORKDIR=$(pwd) ; source ALCF/helpers.sh && setup_venv_from_conda

Looking around the available modules a bit, I noticed a newer frameworks release (frameworks/2024.04.15.002) that had a newer version of both torch and ipex:

module use /soft/preview-modulefiles/24.086.0
module load frameworks/2024.04.15.002.lua
python3 -c 'from mpi4py import MPI; print(MPI.__file__)'
# /soft/datascience/aurora_nre_models_frameworks-2024.1_preview_u1/lib/python3.9/site-packages/mpi4py/MPI.cpython-39-x86_64-linux-gnu.so

>>> import torch
>>> torch.__version__
'2.1.0.post2+cxx11.abi'
>>> import intel_extension_for_pytorch as ipex
>>> ipex.__version__
'2.1.30+xpu'
>>> from mpi4py import MPI; print(MPI.__file__)
/soft/datascience/aurora_nre_models_frameworks-2024.1_preview_u1/lib/python3.9/site-packages/mpi4py/MPI.cpython-39-x86_64-linux-gnu.so

The remaining dependencies were installed identically to what was just done previously for the frameworks/2023.12.15.001 module.

NOTE: In the figures below, we denote these two environments as:

• 2024.0:
  • module load frameworks/2023.12.15.001
• 2024.1:
  • module use /soft/preview-modulefiles/24.086.0
  • module load frameworks/2024.04.15.002.lua
• anl_24_q2_release:
  • eval "$(~/miniconda3/bin/conda shell.zsh hook)"
  • conda activate anl_24_q2_release

🥸 Fix in Disguise

Armed now with functional environment(s) for argonne-lcf/Megatron-DeepSpeed, I was able to resume my previous experiments.

From the discussion with Intel, it was hard to understand / reason about why the flash-attn fix would have any dependence on the learning rate schedule (warmup + decay).

If the flash-attn fix works for a particular learning rate schedule, you would reasonably expect that it should work for any learning rate schedule.

An additional source of confusion for me was that the discrepancy in the loss curves (seemingly) disappeared when using the learning rate settings provided by Intel⁴, but not when using the ALCF defaults⁵.

After thinking about it for a bit and trying to reason about possible causes, I wondered if it might not be a mix of multiple different factors:

1. Small learning rate
2. Very long decay
3. [maybe ?] somehow dependent on the learning rate warmup fraction
   1. preliminary experiments seemed to suggest this was not the case

So, I was curious what would happen if I used the (larger) learning rate value from the ALCF defaults (lr=0.003) with the very long lr-decay-iters: 320000 from Intel.

These results are shown below.

In particular, for all three experiments the following learning rate settings were used:

lr: 0.0003
lr-warmup-frac: 0.05
lr-decay-iters: 320000

Looking at this figure ^, it appears that up until the very very end, all three loss curves agree identically.

However, if we look closely at the very end, it looks like there might be a slight difference beginning to appear between the 2024.0 (brown line) and {anl_24_q2_release, 2024.1} ({dark, light} blue lines, respectively).

Thinking that I might be onto something, I then tried again with a smaller lr-decay-iters: 5000.

This result is shown below:

In particular, we can now more clearly see the difference beginning to appear between the 2024.0 and 2024.1 loss curves.

Continuing on, we see this effect become increasingly dramatic with even smaller values of lr-decay-iters:

In each of these experiments, it appears that:

• 2024.0:
  • Not impacted by this lr-decay-iters dependence
  • Continue to decrease for the duration of training
• 2024.1:
  • Impacted by the lr-decay-iters dependence
  • Plateaus towards the end of training

Older Figs

✅ 2024.0 Fix

Everything seems to work with

module load frameworks/2023.12.15.001

📊 lr-decay-iters Comparison

• 2024.0:
  • 
• 2024.1:
  • 

📈 lr-decay-iters dependence

🏎️ Performance Improvement in 2024.1

lr: 0.0003
lr-warmup-frac: 0.05
lr-decay-iters: null