Summary of Bart memory improvement workstream

memory_changes.md

Summary of Impact

All experiments were run using BartForConditionalGeneration on a batch size of 6 long CNN articles, of uneven length, so some were padded to 1024.

• transformers/master
  • FWD pass: 6.8 GB
  • generate (9 steps): 7.982GB
• fairseq/master:
  • forward: 5.0 GB
  • generate (9 steps): 5.3 GB
• transformers/after_changes
  • FWD pass: 4.8 GB
  • generate: 5.3 GB

Pull Requests

I propose 5 independent changes, with the details of each described in the associated PR description:

• Bart: not returning attn_weights if they are not needed (600 MB)
  • PR Link
  • Scope of code change: bart_modeling.py
  • Scope of speed up: Bart only.
• BartDecoder: separate causal_mask and padding_mask (These were previously combined into one broadcastable mask).
  • (800 MB for forward when not generating)
  • PR Link
• PretrainedModel.generate: call encoder before expanding input_ids (1.5 GB)
  • PR Link
• Bart: not materializing self.lm_head (200 MB)
  • PR Link
• examples/summarization/evaluate_cnn.py: drop columns that are exclusively pad_token_id from input_ids before calling model()
  • PR Link
  • this reduces the runtime to compute EVAL on the CNN examples from 2h to 1:37 before any other changes.
  • Does not effect peak memory usage.

Next Steps:

• There is still a small speed gap that I have not deeply investigated, but it is largely resolved.

memory_logging_tool.md

Memory Logging

I wrote a very lightweight Mixin to support logging: https://github.com/sshleifer/durbango/blob/b87dcf76e3e6f2ae3c07a2183dc2599c002f3d99/durbango/logging_utils.py#L106

Like many testing tools, it requires some setup and doesn't always help you find your bug. In this case, it was helpful for staying organized and being able to very quickly measure the memory usage associated with a specific code change. It is also very immature:

• haven't tested it/cleaned it up.
• only supports torch
• requires pandas (For making nice tables, saving csvs)

Instructions

1. To use it, replace all nn.Module inheritance in your stack (all the modules you want to log from and their parents) with LoggingModule. (or inherit from (nn.Module, LoggingMixin))
2. Add self.log_mem(statement) instead of print(statement) where you want to know the memory. Helpful statements are similar to normal debugging: shapes of things and where you are in the code.
3. Aggregate the logs at the end using self.summary (if you don't want the aggregation, you can call LoggingMixin.collect_log_data())

self.summary Example:

cpu_mem            2.941MB
time                  11.1
gpu_mem_chg        1.908GB
gpu_mem_peak       5.382GB

Warning: there is roughly 100ms overhead associated with every log_mem call, so you can't really compare speed while you are comparing memory.