Method test

tests/unit/v1/zero/test_zero.py:1288–1376

(self, dtype)

    world_size = 2

    def test(self, dtype):

        if not dtype in get_accelerator().supported_dtypes():
            pytest.skip("{dtype} is not supported")

        config_dict = {
            "train_batch_size": 4,
            "steps_per_print": 1,
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 1e-4
                }
            },
            "zero_optimization": {
                "stage": 3
            },
        }

        if dtype == torch.bfloat16:
            if get_accelerator().is_bf16_supported():
                config_dict["bf16"] = {"enabled": True}
            else:
                pytest.skip("bfloat16 is not supported on this accelerator")
        elif dtype == torch.float16:
            if get_accelerator().is_fp16_supported():
                config_dict["fp16"] = {"enabled": True}
            else:
                pytest.skip("fp16 is not supported on this accelerator")
        hidden_dim = 10

        class SubModel(torch.nn.Module):

            def __init__(self, input_size, output_size, dropout_prob=0.5, device=None):
                super(SubModel, self).__init__()
                self.linear = torch.nn.Linear(input_size, output_size, device=device)
                self.dropout = torch.nn.Dropout(dropout_prob)
                self.module_list = torch.nn.ModuleList([torch.nn.Linear(input_size, output_size, device=device)])

            def forward(self, x):
                x = self.linear(x)
                x = self.dropout(x)
                x = self.module_list[0](x)
                return x

        class MyModel(torch.nn.Module):

            def __init__(self, hidden_dim):
                super(MyModel, self).__init__()
                self.l1 = skip_init(Linear, hidden_dim, hidden_dim)
                self.l2 = skip_init(SubModel, hidden_dim, hidden_dim)
                self.l3 = torch.nn.Linear(hidden_dim, hidden_dim)
                self.cel = torch.nn.CrossEntropyLoss()
                self.l4 = skip_init(SubModel, hidden_dim, hidden_dim)

            def forward(self, x, y):
                x = self.l1(x)
                x = self.l2(x)