code: def assert_dist_model_equal(
      model: torch.nn.Module,
      distributed_model: torch.nn.Module,
      device_mesh: DeviceMesh,
      sharding_spec_dict: dict,
  ) -> None:
      state = model.state_dict()
      distributed_state = distributed_model.state_dict()
  
      assert len(state) == len(
          distributed_state
      ), f"len {len(state)} vs {len(distributed_state)}"
  
      for (n1, t1), (n2, t2) in zip(
          state.items(), distributed_state.items()
      ):
          assert n1 == n2
          t1 = t1.cuda()
          t2 = t2.cuda()
          if n2 in sharding_spec_dict:
              layout = Layout(
                  device_mesh=device_mesh,
                  sharding_spec=sharding_spec_dict[n2],
                  global_shape=t1.shape,
              )
              t2.dist_layout = layout
              t2 = to_global(t2)
          assert torch.equal(t1, t2), f"{n1} {t1} vs {t2}"
  
file_path: ColossalAI/tests/test_lazy/lazy_init_utils.py
imports: import random
  from copy import deepcopy
  from typing import Any, Callable, Optional, Tuple
  import numpy
  import torch
  from packaging import version
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.lazy.lazy_init import LazyInitContext, LazyTensor, _MyTensor
  from colossalai.tensor.d_tensor import to_global
  from colossalai.tensor.d_tensor.layout import Layout
  from tests.kit.model_zoo.registry import ModelAttribute
label_desc: Use image node_id 5 for calling a global function with example usage: assert_dist_model_equal(model, distributed_model, device_mesh, sharding_spec_dict) without return types