code: def test_1_param_group(self, tmpdir, zero_stage):
      # XXX: ideally refactor with the 2_param_group test as 75% is the same
      # force all params to be partitioned by forcing threshold=0
      config_dict = {
          "train_micro_batch_size_per_gpu": 2,
          "gradient_accumulation_steps": 2,
          "steps_per_print": 1,
          "zero_optimization": {
              "stage": zero_stage,
              "stage3_param_persistence_threshold": 0,
          },
          "optimizer": {"type": "Adam", "params": {"lr": 1e-3}},
          "fp16": {"enabled": True, "initial_scale_power": 8},
      }
  
      class MyModel(torch.nn.Module):
          def __init__(self, hidden_dim, n_layers):
              super().__init__()
              # to reproduce https://github.com/microsoft/DeepSpeed/pull/1372 it is important that
              # the number of total elements is uneven:
              # (1) 4 layers of 3*(3+1)=12 elements each, 48 in total
              self.ll = torch.nn.ModuleList(
                  torch.nn.Linear(hidden_dim, hidden_dim)
                  for i in range(n_layers)
              )
              # (2) the following adds 4+1=5 elements
              self.classifier = torch.nn.Linear(4, 1)
              # total 48+5=53 (uneven as desired) elements
              self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
  
          def forward(self, x, y):
              hidden = x
              for l in self.ll:
                  hidden = l(hidden)
              return self.cross_entropy_loss(hidden, y)
  
      hidden_dim = 3  # do not change
  
      world_size = dist.get_world_size()
      # we want at least 2x layers as there are gpus to trigger round_robin_fp16_groups reshuffle in zero2
      n_layers = world_size * 2
      model = MyModel(hidden_dim=hidden_dim, n_layers=n_layers)
  
      model, _, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          model_parameters=model.parameters(),
      )
      data_loader = random_dataloader(
          model=model,
          total_samples=16,
          hidden_dim=hidden_dim,
          device=model.device,
      )
  
      for i, batch in enumerate(data_loader):
          loss = model(batch[0], batch[1])
          model.backward(loss)
          model.step()
  
      model.save_checkpoint(tmpdir)
  
      # make sure all sides saved it
      dist.barrier()
  
      if zero_stage == 3:
          with deepspeed.zero.GatheredParameters(
              list(model.module.parameters(recurse=True)),
              modifier_rank=None,
          ):
              pass  # this forces gathering the model
  
      # dump_state_dict(model)
  
      orig_state_dict = {}
      for name, param in model.module.named_parameters():
          orig_state_dict[name] = param.detach().cpu()
  
      if dist.get_rank() == 0:
          fp32_model = load_state_dict_from_zero_checkpoint(
              model.module, tmpdir
          )
          # dump_state_dict(fp32_model)
  
          fp32_state_dict = fp32_model.state_dict()
          for name in orig_state_dict.keys():
              # float() workaround for torch<1.6
              assert torch.allclose(
                  orig_state_dict[name].float(),
                  fp32_state_dict[name].float(),
              )
  
imports: import math
  from typing import Dict, List, Set
  import pytest
  import deepspeed.comm
  import torch
  from torch import Tensor
  from torch.nn import Linear, Module
  from torch.nn.modules.container import ModuleList
  from torch.nn.modules.loss import L1Loss
  from torch.nn.parameter import Parameter
  from unit.common import DistributedTest
  from unit.simple_model import SimpleModel, random_dataloader
  import deepspeed
  from deepspeed.runtime.engine import DeepSpeedEngine
  from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
  from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
label_desc: Use image node_id 1 for calling the TestZeroToFP32 obj's underlying member method code with example usage: obj.test_1_param_group(tmpdir, zero_stage) and returns: self