Enable remat checkpoints to host instead of TPU memory #643
Conversation
axlearn/common/attention.py
Outdated
| @@ -3874,6 +3874,7 @@ def build_remat_spec( | |||
| ], | |||
| self_attention: bool = True, | |||
| feed_forward: bool = False, | |||
| offload: bool = False, | |||
There was a problem hiding this comment.
Instead of a bool, should we allow the caller to customize offload_dst directly?
| offload: bool = False, | |
| offload_dst: Optional[Literal["pinned_host"]] = None, |
This will make the API more extensible and closer to the JAX API.
There was a problem hiding this comment.
Agree. I will take another stab at this PR with focus on staying closer to JAX API and extensibility.
There was a problem hiding this comment.
Resolved, could you review again?
axlearn/common/attention.py
Outdated
| @@ -3891,6 +3904,7 @@ def build_remat_spec( | |||
| stack_cfg: A transformer config. | |||
| self_attention: Checkpoint self attention layer activations if true. | |||
| feed_forward: Checkpoint feed-forward layer activations if true. | |||
| offload_dst: Destination of remat checkptoing offloading. | |||
There was a problem hiding this comment.
Add a link to the JAX documentation on offset_dst on the potential values?
There was a problem hiding this comment.
There are no docs yet for this. Do you want me to link to the maxtext code as a comment? https://github.com/google/maxtext/blob/ebd39aa64d670fa13a313b6f776e01ad9e450321/MaxText/layers/models.py#L230
There was a problem hiding this comment.
No ah I misunderstood Ruoming's comment and thought he was fine there being no docs. Let me add the link to maxtext as a comment.
| @@ -188,6 +188,7 @@ def get_trainer_kwargs( | |||
| num_kv_heads=None if version == Version.V1 else 8, | |||
| rope_theta=rope_theta, | |||
| flash_attention=flash_attention, | |||
| remat_offload_dst="pinned_host", | |||
There was a problem hiding this comment.
Add a comment on the observed MFU and step time?
| @@ -188,6 +188,7 @@ def get_trainer_kwargs( | |||
| num_kv_heads=None if version == Version.V1 else 8, | |||
| rope_theta=rope_theta, | |||
| flash_attention=flash_attention, | |||
| remat_offload_dst="pinned_host", | |||
There was a problem hiding this comment.
Is this option limited to 70B? Do we want to apply it to 7B and other models?
There was a problem hiding this comment.
Only use remat checkpoint offload into host when you could benefit from the extra TPU memory.
If there is plenty TPU memory then remat checkpoint into TPU memory.
If TPU memory is low or you want to squeeze as much per device batch as possible, then use offload_dst=pinned_Host.
So yes, it could make a lot of sense for 7B too on V5e and Trilium since then we can possibly increase per device batch size.
There was a problem hiding this comment.
I would prefer to have this PR focus on enabling remat offload and 70B. As a follow up I can do the following as part of V5E perf benchmarking:
- Enable remat offload for 7B and compare performance before and after
- See if I can increase per device batch size for 7B after enabling remat_offload
Would that work for you?
|
@markblee could you trigger CI, review and merge if good? |
Co-authored-by: Mark Lee <[email protected]>
* remat checkpoints to host * update golden configs * Change offload to offload_dst * add step time and MFU for v5e fuji v2 70b * Add maxtext example in code comment * Update axlearn/common/attention.py Co-authored-by: Mark Lee <[email protected]> --------- Co-authored-by: Mark Lee <[email protected]>
This allowed us to get MFU of fuji v2 70B from 58.50% to 61.83%