deepspeedai / deepspeedai/DeepSpeed
How to calculate Transformer/Bert layer's activation memory?
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
In the deepspeed docs, the activation memory is calculated by:
XXX: For Transformers is probably around (2* seq * attn_heads + 16 * hidden_size) * sequence * batch/gpu
In the ZeRO-Infinity paper, Section 3, the activation memory is calculated by:
2 * bsz * seq * hd * nl/ci
where bsz is batch size, seq is sequence length, hd is hidden dimension, nl is number of Transformer layers and ci is the number of Transformer blocks between two .activation checkpoints.
If we don't consider activation checkoutpoints, we get:
2 * bsz * seq * hd * nl
Which one is correct? Would anybody please tell me how to calculate the activation memory of Transformer/Bert layer with multi-head self-attetion?
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Review the DeepSpeed memory documentation and Section 3 of the ZeRO-Infinity paper, focusing on the definitions of sequence length, hidden dimension, layer count, and activation checkpointing. Reconcile the two formulas for a Transformer/BERT layer with multi-head self-attention, then update the documentation with a consistent explanation and an example calculation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Documentation
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100