Enable ZeRO-3 linear wrapper for existing models (#8189)
Passing an already-constructed model to `deepspeed.initialize()` with
ZeRO-3 and `memory_efficient_linear=true` does not install the ZeRO-3
Linear wrapper. The wrapper is currently installed only when the model
is constructed inside a `deepspeed.zero.Init()` context.
Without the wrapper, the standard Linear implementation can retain the
gathered weight storage until backward completes, significantly
increasing memory usage.
This PR activates the existing ZeRO-3 Linear wrapper for the
`deepspeed.initialize(model=...)` path when
`memory_efficient_linear=true`, without requiring a
`deepspeed.zero.Init()` context.
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
