code: def get_llama_wrapper():
      """we register our main layer class and use the fsdp transformer wrapping policy
      ensures embedding layers are in the root fsdp unit for shared access and that fsdp units map to transformer layers
      """
      # ====   use new transformer wrapper
  
      llama_auto_wrap_policy = functools.partial(
          transformer_auto_wrap_policy,
          transformer_layer_cls={
              LlamaDecoderLayer,
          },
      )
  
      return llama_auto_wrap_policy
  
imports: import functools
  from transformers.models.llama.modeling_llama import LlamaDecoderLayer
  from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy, size_based_auto_wrap_policy
label_desc: Use image node_id 2 for calling a global function with example usage: get_llama_wrapper() and returns: llama_auto_wrap_policy