[zero] splitting up context managers into enter/exit trackers

Having converted multiple HF transformers models to zero3, I feel there is a need to make the context managers (zero.Init + GatheredParameters) more code readability/maintenance friendly, since currently each of these 2 forces:

x = y

to be rewritten as:

        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_lm_head.weight, modifier_rank=None):
                x = y
        else:
            x = y

and very often the sections that need wrapping can be multiple lines of code, thus resulting in very difficult to read and maintain code, e.g.:

        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_lm_head.weight, modifier_rank=0):
                if torch.distributed.get_rank() == 0:
                    # Copy old lm head weights to new lm head
                    if not transposed:
                        new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :]
                    else:
                        new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy]

                    # Copy bias weights to new lm head
                    if has_new_lm_head_bias:
                        new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]
        else:
            # Copy old lm head weights to new lm head
            if not transposed:
                new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :]
            else:
                new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy]

            # Copy bias weights to new lm head
            if has_new_lm_head_bias:
                new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]

which surely can be refactored into a sub-function, but still this is not very easy to read.

I propose breaking down these context managers into separate enter/exit calls (while keeping the context manager as is), so that the code becomes:

        if is_deepspeed_zero3_enabled():
            import deepspeed
            gather_ctx = deepspeed.zero.GatheredParameters(...).start

        x = y

        if is_deepspeed_zero3_enabled():
            gather_ctx.stop()

this is a way easier to read and requires no code duplication. And the reader who isn't interested in deepspeed, can mentally quickly skip over the if deepspeed branches.

Same goes for zero.Init().

Alternatively, there should be a no-op wrapper for when there is no deepspeed, but it would still hinder the readability since the reader might not understand why there is a wrapper there in the first place.

Thoughts?

Start by locating the implementations of deepspeed.zero.Init and deepspeed.zero.GatheredParameters, then review how their current context-manager behavior is used with Zero-3. Define and test separate enter/exit tracking while preserving the existing context-manager interface; done means callers can avoid duplicated branches without changing behavior.