Ring Attention w/ document packing produces different results
Please check that this issue hasn't been reported before.
- I searched previous Bug Reports didn't find any similar reports.
Expected Behavior
When using document packing, ring attention should produce similar or bitwise identical outputs to varlen_llama3 or padded batches.
Current behaviour
When using specifically document packing with ring attention, the loss on a given step is much higher and different than would otherwise be expected. Ring attention on a single non-packed document works. Document packing with varlen_llama3 attention works.
Multi-doc packed data (data/packed.pt, pack 0 = 5 docs with
position_id resets at each boundary):
| config | step-0 loss | ratio to cp=1 |
|---|---|---|
cp=1 (baseline) |
3.171875 |
1.000× |
cp=2, ring_attn_func=batch_ring |
5.843750 |
1.843× (bug) |
cp=2, ring_attn_func=varlen_llama3 |
3.156250 |
0.995× (fix) |
Single-doc control (data/packed_single.pt, one 1024-token doc,
position_ids 0..1023 contiguous, zero resets). Every step is
bitwise-identical across all three configs:
| config | loss (all steps) |
|---|---|
cp=1 (baseline) |
3.031250 |
cp=2, ring_attn_func=batch_ring |
3.031250 |
cp=2, ring_attn_func=varlen_llama3 |
3.031250 |
Steps to reproduce
I don't have the skill set to debug exactly what is wrong with it, but Claude was able to produce a minimal reproduction of the issue which can be found here: https://github.com/Ueaj-Kerman/axolotl-cp-repro
Config yaml
Possible solution
I'd bet the document ids aren't being rotated along the ring, causing the issue
Which Operating Systems are you using?
- Linux
- macOS
- Windows
Python Version
3.11.15
axolotl branch-commit
08fc7de87e79f38c367f6776c5111b40a914062e
Acknowledgements
- My issue title is concise, descriptive, and in title casing.
- I have searched the existing issues to make sure this bug has not been reported yet.
- I am using the latest version of axolotl.
- I have provided enough information for the maintainers to reproduce and diagnose the issue.
Source: axolotl-ai-cloud/axolotl