code: def write_tp_pp_split(
      model, splits, app_state, tp_size, pp_rank, write_path
  ):
      """
      Function to write the given TP PP split to NeMo File.
  
      Save each of the TP ranks in reverse order
      This is done so that the last PP rank will save the last TP rank only after all other PP TP ranks are saved
      The final rank will then save a new NeMo file with all other ranks inside.
  
      Args:
          model: The model corresponding to the current TP PP split. Contains partial parameters.
          splits: Nested List of tensors containing the TP splits of the current model given current PP rank.
              Indexed as splits[idx][tp_rank].
          app_state: AppState object.
          tp_size:  The global tensor-parallel size of the final model.
          pp_rank: The local pipeline parallel rank of the final model.
          write_path: The path to save the NeMo file.
      """
      for tp_rank in range(tp_size - 1, -1, -1):
          app_state.pipeline_model_parallel_rank = pp_rank
          app_state.tensor_model_parallel_rank = tp_rank
  
          idx = 0
          for name, param in model.named_parameters():
              split_val = splits[idx][tp_rank].clone()
  
              if param.shape != split_val.shape:
                  logging.info(
                      f"Warning: Shape mismatch for parameter {name} required shape: {param.shape}, split shape: {split_val.shape}. Padding to match required size."
                  )
  
                  if split_val.shape[1:] == param.shape[1:]:
                      pad = [0, 0] * len(split_val.shape)
                      pad[-1] = param.shape[0] - split_val.shape[0]
                      split_val = torch.nn.functional.pad(
                          split_val, pad, "constant"
                      )
                  elif split_val.shape[:-1] == param.shape[:-1]:
                      pad = [0, param.shape[-1] - split_val.shape[-1]]
                      split_val = torch.nn.functional.pad(
                          split_val, pad, "constant"
                      )
                  else:
                      raise RuntimeError(
                          f"Can not handle parameter {name}, required shape: {param.shape}, split shape: {split_val.shape}."
                      )
  
              param.data = split_val
              idx += 1
  
          if write_path is not None:
              logging.info(
                  f"Writing pp rank {pp_rank} tp rank {tp_rank} to file {write_path}"
              )
              model.save_to(write_path)
  
imports: import os
  import shutil
  import tarfile
  import tempfile
  from argparse import ArgumentParser
  from typing import Dict, List
  import torch
  import torch.nn
  from omegaconf import OmegaConf, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.nlp.parts.nlp_overrides import NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE, GradScaler, MegatronHalfPrecisionPlugin, NLPDDPStrategy, NLPSaveRestoreConnector, PipelineMixedPrecisionPlugin
  from nemo.utils import logging, model_utils
  from nemo.utils.app_state import AppState
label_desc: Use image node_id 6 for calling a global function with example usage: write_tp_pp_split(model, splits, app_state, tp_size, pp_rank, write_path) without return types