Fix incorrect patch in zero.init

[VeryLazyBoy]

The code below has a problem where cls.__init__ in line 525 can be modified before assignment to _old_init. This could lead to an incorrect __init__ being backed up:

DeepSpeed/deepspeed/runtime/zero/partition_parameters.py

Lines 524 to 534 in ffe0af2

	def _enable_class(cls):
	cls._old_init = cls.__init__
	cls.__init__ = partition_after(cls.__init__)
	def _init_subclass(cls, **kwargs):
	cls._old_init = cls.__init__
	cls.__init__ = partition_after(cls.__init__)
	# Replace .__init__() for all existing subclasses of torch.nn.Module recursively
	for subclass in get_all_subclasses(torch.nn.modules.module.Module):
	_enable_class(subclass)

Test Case

import deepspeed
from torch import nn


class ModelA(nn.Module):
    def __init__(self):
        super().__init__()


class ModelB(ModelA):
    pass


original_init = ModelA.__init__


ds_config = {
    'fp16': {'enabled': False},
    'bf16': {'enabled': True},
    'zero_optimization': {
        'stage': 3,
        'offload_optimizer': {
            'device': 'cpu',
            'pin_memory': True
        },
        'offload_param': {
            'device': 'cpu',
            'pin_memory': True
        },
    },
    'gradient_accumulation_steps': 1,
    'gradient_clipping': 1,
    'train_batch_size': 1,
    'train_micro_batch_size_per_gpu': 1
}


with deepspeed.zero.Init(config_dict_or_path=ds_config, enabled=True, mem_efficient_linear=False, mpu=None):
    model_a = ModelA()
    assert ModelA.__init__ != original_init

assert ModelA.__init__ == original_init
assert ModelB.__init__ == original_init   #  Fails here. If not, please try several times since it depends on the order of modifications

[VeryLazyBoy]

@microsoft-github-policy-service agree

[VeryLazyBoy]

A better solution is proposed to handle _init_subclass as well

[tohtana]

Thank you @VeryLazyBoy for the great catch!

I think the issue is that we patch superclass's cls.__init__ when cls doesn't have its __init__. So I try another approach in this branch. Do you think if this works?
This is less intrusive as we do not set __init__ when cls doesn't have it.

[VeryLazyBoy]

@tohtana Yes! Your approach is less intrusive and much better. Let's go ahead with this new method. Should I close this merge request?

[tohtana]

@VeryLazyBoy Thank you for your response!
Let me create a PR using my branch to make sure it works. Let's close this PR after all test pass with the PR.