Optimizer offloading through weight-only offload #867
Conversation
axlearn/common/optimizers.py
Outdated
| Only wrap the optimizer that you actually want to offload with this function to avoid | ||
| unneseccary overhead. This is usually the optimizer that occupies the most HBM. For example, | ||
| when you have chained optimizers: |
There was a problem hiding this comment.
Where does the overhead come from? Is it from the states of clip_by_global_norm being offloaded? If so, could we use regular expressions to specify which states to offload?
Before the optimizer can be invoked, the offloaded optimizer states need to be transferred to device memory space. If we remove these device_put calls, we will get errors like |
hanzhi713
force-pushed
the
weight-only-offload-cleanup
branch
from
3adee99 to
5ea7bb4
Compare
|
@markblee Can you take a look at the pytype errors in CI? Should I change Nested type to include tuple/namedtuple or should I just ignore the errors? |
Maybe we can relax the partition fn return type to include named tuple? Changing Nested may have undesirable impact elsewhere. |
Thanks. I included the return type in Nested[...]. Just NamedTuple wouldn't work. |
This PR requires jax >= 0.4.34, != 0.4.35, >=0.4.36: it works on jax 0.4.34, but is broken on jax 0.4.35 due to libtpu bug. It worked on nightly jax 0.4.36 as of 10/30.
This PR represents effort to enable optimizer offloading. The approach we use in this PR is weight-only offloading, which is based on similar building blocks as activation offloading (aka remat offload). When offloading is enabled, optimizer states are stored on CPU pinned memory. Before apply optimizer to calculate updates, optimizer states are moved from CPU memory to HBM via
jax.device_put. The new optimizer states are moved back from HBM to CPU.An alternative approach to this PR is host computation. Host computation means that optimizer transformations are computed on CPU. Before the start of the computation, gradients and weights are transferred to CPU, and after the computation, their new values are transferred back to HBM. This method has lower HBM footprint, but it's much 2x ~ 3x slower due to slow CPU computation. Also, it's very buggy.
TLDR: to be merged after upgrading jax to 0.4.36.