code: def test_state_averager(
      offload_optimizer: bool,
      reuse_tensors: bool,
      sync_epoch_when_averaging: bool,
  ):
      dht1 = hivemind.DHT(start=True)
      dht2 = hivemind.DHT(
          initial_peers=dht1.get_visible_maddrs(), start=True
      )
  
      torch.manual_seed(1337)
      torch.use_deterministic_algorithms(True)
      # note: use_deterministic_algorithms does not affect further tests because this test is forked
  
      model1 = nn.Linear(2, 3)
      model2 = nn.Linear(2, 3)
  
      extras1 = (torch.randn(2, 2), -torch.rand(1))
      extras2 = (-torch.randn(2, 2), torch.rand(1))
  
      common_kwargs = dict(
          optimizer=partial(torch.optim.Adam, lr=0.1, betas=(0.9, 0.9)),
          scheduler=partial(
              torch.optim.lr_scheduler.LambdaLR,
              lr_lambda=lambda t: 1.0 / max(1, t),
          ),
          sync_epoch_when_averaging=sync_epoch_when_averaging,
          average_opt_statistics=("exp_avg_sq",),
          offload_optimizer=offload_optimizer,
          reuse_tensors=reuse_tensors,
          target_group_size=2,
          prefix="my_exp",
      )
  
      avgr1 = TrainingStateAverager(
          dht=dht1,
          params=model1.parameters(),
          extra_tensors=extras1,
          start=True,
          **common_kwargs,
      )
      avgr2 = TrainingStateAverager(
          dht=dht2,
          params=model2.parameters(),
          extra_tensors=extras2,
          start=True,
          **common_kwargs,
      )
  
      x = torch.ones(2)
  
      for step in range(20):
          F.mse_loss(model1(x), torch.ones(3)).mul(2).backward()
          avgr1.step(
              optimizer_step=True,
              zero_grad=True,
              averaging_round=(step == 10),
              delay_averaging=True,
          )
  
          F.mse_loss(model2(x), -torch.ones(3)).backward()
          avgr2.step(
              optimizer_step=True,
              zero_grad=True,
              averaging_round=(step == 10),
              delay_averaging=False,
          )
  
      if ZERO_GRAD_SET_TO_NONE_DEFAULT:
          assert (
              model1.weight.grad is None and model2.weight.grad is None
          ), ".zero_grad() wasn't called"
      else:
          assert torch.all(model1.weight.grad == 0) and torch.all(
              model2.weight.grad == 0
          ), ".zero_grad() wasn't called"
      assert (
          model1(x).mean() > 0.5 and model2(x).mean() < -0.5
      ), "models did not train properly"
      assert torch.allclose(
          extras1[0], extras2[0]
      ), "first extra tensors were not averaged"
      assert torch.allclose(
          extras1[1], extras2[1]
      ), "second extra tensors were not averaged"
  
      stats1 = avgr1.optimizer.state_dict()["state"][0][
          "exp_avg_sq"
      ].clone()
      stats2 = avgr2.optimizer.state_dict()["state"][0][
          "exp_avg_sq"
      ].clone()
      assert not torch.allclose(stats1, stats2)
  
      avgr1.step(increment_epoch=True)
  
      avgr1.step(
          increment_epoch=True,
          averaging_round=True,
          delay_averaging=True,
      )
      avgr2.step(
          increment_epoch=True,
          averaging_round=True,
          delay_averaging=True,
      )
  
      avgr1.step(wait_for_delayed_updates=True)
      avgr2.step(wait_for_delayed_updates=True)
  
      assert torch.allclose(
          model1(x), model2(x)
      ), "model parameters were not averaged correctly"
      assert torch.allclose(
          avgr1.optimizer.state_dict()["state"][0]["exp_avg_sq"],
          (stats1 + stats2) / 2,
      )
      assert torch.allclose(
          avgr2.optimizer.state_dict()["state"][0]["exp_avg_sq"],
          (stats1 + stats2) / 2,
      )
      assert avgr1.local_epoch == 2
      assert avgr2.local_epoch == (
          2 if sync_epoch_when_averaging else 1
      )
  
imports: import ctypes
  import multiprocessing
  import time
  from functools import partial
  import numpy
  import pytest
  import torch
  import torch.nn
  import torch.nn.functional
  import hivemind
  from hivemind.averaging.control import AveragingStage
  from hivemind.optim.grad_averager import GradientAverager, GradientAveragerFactory
  from hivemind.optim.optimizer import Optimizer
  from hivemind.optim.power_sgd_averager import PowerSGDGradientAverager
  from hivemind.optim.progress_tracker import ProgressTracker
  from hivemind.optim.state_averager import ZERO_GRAD_SET_TO_NONE_DEFAULT, TrainingStateAverager
  from hivemind.utils.crypto import RSAPrivateKey
label_desc: Use image node_id 3 for calling a global function with example usage: test_state_averager(offload_optimizer, reuse_tensors, sync_epoch_when_averaging) without return types