class_docstr: A single transformer layer.
  
  Transformer layer takes input with size [s, b, h] and returns an
  output of the same size.
      
class_docstr_tok: ["A","single","transformer","layer",".","Transformer","layer","takes","input","with","size","[","s",",","b",",","h","]","and","returns","an","output","of","the","same","size","."]
code: def sharded_state_dict(self, prefix=""):
      state_dict = self.state_dict(keep_vars=True)
  
      tensor_parallel_layers_axis_map = {
          "self_attention.linear_qkv.weight": 0,
          "self_attention.linear_qkv.bias": 0,
          "self_attention.linear_proj.weight": 1,
          "mlp.linear_fc1.weight": 0,
          "mlp.linear_fc1.bias": 0,
          "mlp.linear_fc2.weight": 1,
      }
  
      offset = self._get_layer_offset()
      num_layers = self.config.num_layers
  
      sharded_state_dict = {}
  
      for layer_name in state_dict.keys():
          tensor = state_dict[layer_name]
          global_layer_offset = (
              self.layer_number - 1
          )  # self.layer_number starts at 1
          layer_key = f"{prefix}{global_layer_offset - offset}.{layer_name}"  # module list index in TransformerBlock
          sharded_offsets = [
              (0, global_layer_offset, num_layers)
          ]  # PP sharding
  
          if layer_name in tensor_parallel_layers_axis_map:
              tp_axis = tensor_parallel_layers_axis_map[layer_name]
              # TP sharding
              sharded_offsets.append(
                  [
                      tp_axis + 1,  # +1 for PP dimension
                      parallel_state.get_tensor_model_parallel_rank(),
                      parallel_state.get_tensor_model_parallel_world_size(),
                  ]
              )
              replica_id = parallel_state.get_data_parallel_rank()
          else:
              replica_id = (
                  parallel_state.get_data_parallel_rank()
                  * parallel_state.get_data_parallel_world_size()
                  + parallel_state.get_tensor_model_parallel_rank()
              )
  
          if layer_name.endswith("._extra_state"):
              sharded_state_dict[layer_key] = ShardedObject(
                  f"{prefix}{layer_name}",
                  tensor,
                  (num_layers,),
                  (global_layer_offset,),
                  replica_id,
              )
  
          else:
              sharded_state_dict[
                  layer_key
              ] = ShardedTensor.from_rank_offsets(
                  f"{prefix}{layer_name}",
                  tensor,
                  *sharded_offsets,
                  replica_id=replica_id,
                  prepend_axis_num=1,  # for PP sharding
              )
  
      return sharded_state_dict
  
imports: from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults
label_desc: Use image node_id 3 for calling the FalconTransformerLayer obj's underlying member method code with example usage: obj.sharded_state_dict(prefix) and returns: sharded_state_dict