code: def auto_wrap_ln(
      module, fsdp_config=None, wrap_it=True, assert_on_collision=True
  ):
      """
      Auto wrap all LayerNorm instances with a safer FSDP, esp. when convert
      to LayerNorm is used and the outer FSDP is flattening.
      We put KN in is own full precision, unflatten, single GPU group FSDP.  Note, LNs still have
      a group size == world_size. The input and output for LN are still FP16 in mixed precision mode.
      See ``keep_batchnorm_fp32`` here: https://nvidia.github.io/apex/amp.html
      This needs to be done at each rank, like models being wrapped by FSDP at each rank.
      Args:
          module (nn.Module):
              The model (or part of the model) in which BN to be pre-wrapped.
          process_group (ProcessGroup):
              Optional process group to be used.
          fsdp_config (Dict):
              Optional fsdp_config to be used.
          wrap_it (bool):
              Whether or not wrap the module after setting the config.
              Default: True
          assert_on_collision (bool):
              Whether or not assert if a wrapper_config already exists on the module.
              Default: True
      Returns:
          Processed module, where BNs are wrapped with a special FSDP instance.
      """
      if fsdp_config is None:
          fsdp_config = {
              "mixed_precision": False,  # Keep the weights in FP32.
              "flatten_parameters": False,  # Do not flatten.
              # Reshard==False is good for performance. When FSDP(checkpoint(FSDP(bn))) is used, this
              # **must** be False because BN's FSDP wrapper's pre-backward callback isn't called
              # within the checkpoint's outer backward when multiple forward passes are used.
              "reshard_after_forward": False,
              # No bucketing or small bucketing should be enough for BNs.
              "bucket_cap_mb": 0,
              # Setting this for SyncBatchNorm. This may have a performance impact. If
              # SyncBatchNorm is used, this can be enabled by passing in the `fsdp_config` argument.
              "force_input_to_fp32": False,
          }
  
      # Assign the config dict to BNs.
      for m in module.modules():
          if isinstance(m, torch.nn.LayerNorm):
              if assert_on_collision:
                  assert not hasattr(
                      m, "wrapper_config"
                  ), "Module shouldn't already have a wrapper_config. Is it tagged already by another policy?"
              m.wrapper_config = fsdp_config
  
      # Wrap it.
      with (
          enable_wrap(config_auto_wrap_policy, wrapper_cls=FSDP)
          if wrap_it
          else contextlib.suppress()
      ):
          return auto_wrap(module)
  
imports: import contextlib
  import os
  import random
  from copy import deepcopy
  import numpy
  import pycls.core.benchmark
  import pycls.core.builders
  import pycls.core.checkpoint
  import pycls.core.config
  import pycls.core.distributed
  import pycls.core.logging
  import pycls.core.meters
  import pycls.core.net
  import pycls.core.optimizer
  import pycls.datasets.loader
  import torch
  import torch.cuda.amp
  from fairscale.nn import auto_wrap, config_auto_wrap_policy, enable_wrap
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.optim.grad_scaler import ShardedGradScaler
  from pycls.core.config import cfg
  from pycls.core.io import cache_url, pathmgr
label_desc: Use image node_id 1 for calling a global function with example usage: auto_wrap_ln(module, fsdp_config, wrap_it, assert_on_collision) and returns: auto_wrap