docstr: Check that local_state_dict can be saved and loaded for a given worker, and that training updates it
docstr_tok: ["Check","that","local_state_dict","can","be","saved","and","loaded","for","a","given","worker",",","and","that","training","updates","it"]
code: def _load_local_and_train(
      self, config, rank, group, d_model=16, d_vocab=23
  ):
      """Check that local_state_dict can be saved and loaded for a given worker, and that training updates it"""
      model = self.get_wrapped_model(
          group,
          cuda_first=False,
          config=config,
          d_vocab=d_vocab,
          d_model=d_model,
          add_bn=False,
      )  # Set bn=True here to show that BN doesn't get updated
      state_1 = model.local_state_dict()
      state_before_training = {
          k: v.cpu().clone() for k, v in state_1.items()
      }
      assert len(state_1) > 0
      model.load_local_state_dict(state_1)
      weight_key = (
          "flat_param_0"
          if model.flatten_parameters
          else "embed_tokens.weight"
      )
  
      state_1_weight = state_1[weight_key]
      assert (
          state_1_weight.dtype == torch.float32
      ), f"got dtype {state_1_weight.dtype} expected torch.float32"
      if not model.flatten_parameters:
          # The weight will be sharded since we access module.state_dict directly
          state_1_module_weight = model.module.state_dict()[weight_key]
          torch.testing.assert_allclose(
              state_1_weight, state_1_module_weight
          )
          torch.testing.assert_allclose(
              state_1_weight, model.module.embed_tokens.weight
          )
      # increasing number of epochs from 1 to 6 for ShardedGradScaler to work properly.
      # test fails for num_epochs < 6 since the updates are skipped due to gradient being inf.
      self._train_for_several_steps(model, 6, model.mixed_precision)
  
      state_2 = model.local_state_dict()
      state_after_training = {
          k: v.cpu().clone() for k, v in state_2.items()
      }
      model.load_local_state_dict(state_2)
  
      assert state_1.keys() == state_2.keys()
  
      # Assert that parameters were updated since before training
      unchanged = []
      unwrapped_model = model.module.module
      buffers = {name for name, _ in unwrapped_model.named_buffers()}
      for k in state_1:
          if (
              state_before_training[k] == state_after_training[k]
          ).all() and (k not in buffers):
              unchanged.append(k)
      if unchanged:
          raise AssertionError(
              f"params {unchanged} not changed after training"
          )
  
imports: import functools
  import unittest
  from parameterized import parameterized
  import pytest
  import torch
  from torch import nn
  from fairscale.fair_dev.testing.testing import dist_init, objects_are_equal, skip_if_cuda, teardown, temp_files_ctx
  from fairscale.internal import torch_version
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from .test_fsdp import CONFIG_OPTIONS, DistributedTest, NestedWrappedModule, TransformerWithSharedParams, rename_test, spawn_and_init
label_desc: Use image node_id 2 for calling the TestLocalStateDict obj's underlying member method code with example usage: obj._load_local_and_train(config, rank, group, d_model, d_vocab) without return types