code: def test(self, zero_stage=2):
      config_dict = {
          "train_micro_batch_size_per_gpu": 2,
          "gradient_accumulation_steps": 2,
          "steps_per_print": 1,
          "zero_optimization": {"stage": zero_stage},
          "optimizer": {"type": "Adam", "params": {"lr": 1e-3}},
          "fp16": {"enabled": True, "initial_scale_power": 8},
      }
      hidden_dim = 4
  
      model = SimpleModel(hidden_dim=hidden_dim)
      model, _, _, _ = deepspeed.initialize(
          config=config_dict,
          model=model,
          model_parameters=model.parameters(),
      )
  
      # get nccl all-gather send buffers alignment factor
      nccl_start_alignment_factor = (
          model.optimizer.nccl_start_alignment_factor
      )
  
      parallel_partitioned_bit16_groups = (
          model.optimizer.parallel_partitioned_bit16_groups
          if zero_stage == 2
          else model.optimizer.parallel_partitioned_fp16_groups
      )
      for data_parallel_partitions in parallel_partitioned_bit16_groups:
          for partition_id, partitioned_data in enumerate(
              data_parallel_partitions
          ):
              # verify that data partition start locations are 4-byte aligned
              assert (
                  partitioned_data.data_ptr()
                  % (2 * nccl_start_alignment_factor)
                  == 0
              )
  
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 TestPartitionNcclAlignment obj's underlying member method code with example usage: obj.test(zero_stage) without return types