[BUG] `deepspeed.zero.Init` leaks

Describe the bug
deepspeed.zero.Init leaks. The last batch norm initialization after the deepspeed.zero.Init in the snippet below is offloaded to disk.

This is making testing code with deepspeed extremely complicated...

I can't explain to myself why the linear layer (tested also for a conv) works, but a batch norm doesn't...

deepspeed works for us in the normal use case (running inference, etc.).

To Reproduce
Run the script below. I tested this in a Jupyter Notebook as well as as a Python test.

import os
import deepspeed
import torch
import time

envvars = {
        "RANK": "0",
        "WORLD_SIZE": "1",
        "MASTER_ADDR": "127.0.0.1",
        "MASTER_PORT": "8088",
        "LOCAL_RANK": "0",
        "ACCELERATE_USE_DEEPSPEED": "true",
    }
os.environ.update(envvars)

ds_config = {
    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "cpu",
        },
    },
    "train_batch_size": 1,
    "comms_logger": {
      "enabled": "true",
      "verbose": "true",
      "prof_all": "true",
      "debug": "true"
    }
}

bn = torch.nn.BatchNorm2d(10)
assert not hasattr(bn.weight, "ds_shape")
lin = torch.nn.Linear(10, 10)
assert not hasattr(lin.weight, "ds_shape")

with deepspeed.zero.Init(
    config_dict_or_path=ds_config, enabled=True
):
    bn_ds = torch.nn.BatchNorm2d(10)
    assert hasattr(bn_ds.weight, "ds_shape")
    lin_ds = torch.nn.Linear(10, 10)
    assert hasattr(lin_ds.weight, "ds_shape")

# Works
lin = torch.nn.Linear(10, 10)
assert not hasattr(lin.weight, "ds_shape")
# Doesn't work
bn = torch.nn.BatchNorm2d(10)
assert not hasattr(bn.weight, "ds_shape")

Expected behavior
No assertion fails.

ds_report output
Please run ds_report to give us details about your setup.

Screenshots

System info (please complete the following information):

• OS: Ubuntu 22.04 in Docker
• 1x T4
• Interconnects (if applicable) ❌
• Python version: 3.9.16
• Any other relevant info about your setup

Launcher context
Python.

Docker context
I'm using docker but I can't share the image.

Additional context
Is there any other way to test deepspeed in Python test?

I'm looking for a way to use deepspeed functionally without altering some global state or for a way to reset that inside a Python script/test.

I'd also appreciate some insight into at what time exactly deepspeed patches the modules. I've noticed you're patching the modules' init functionality with a parameter offloading one.

Start with the reproducer in the issue and run it in the stated Python, Ubuntu, and Docker setup, then use ds_report for environment details. Trace deepspeed.zero.Init around the BatchNorm2d construction and its module init patching or parameter offloading. Done means the post-context BatchNorm2d assertion passes without affecting later module construction or requiring global-state cleanup.