code: def split_partition(
      model,
      partitions,
      pp_size: int,
      tp_size: int,
      pp_rank: int,
      offset: int,
      pp_split_rank: int = 0,
      write_path: str = None,
      megatron_legacy: bool = False,
  ):
      if len(partitions) != 2:
          raise ValueError(
              "Can only split partitions of model with TP=1. For partitions of models with TP>1, merge first."
          )
  
      if tp_size < 1:
          raise ValueError("TP size must to be >= 1.")
  
      if pp_size < 1:
          raise ValueError("PP size must to be >= 1.")
  
      # Setup app state to mimic current PP and TP ranks with single merged module
      app_state = AppState()
      app_state.data_parallel_rank = 0
      app_state.pipeline_model_parallel_size = pp_size
      app_state.tensor_model_parallel_size = tp_size
      app_state.model_parallel_size = (
          app_state.pipeline_model_parallel_size
          * app_state.tensor_model_parallel_size
      )
  
      # Go in reverse for TP order, as PP 0 TP 0 will merge all preceding files
      app_state.pipeline_model_parallel_rank = pp_rank
      app_state.tensor_model_parallel_rank = tp_size - 1
  
      # Compute reverse offset of parameter index from global map
      num_params = sum(
          [1 for _ in model.parameters()]
      )  # Count number of parameters iteratively
      idx = (
          offset - num_params + 1
      )  # start index of current PP TP rank in global map
  
      assert (
          idx + num_params - 1 == offset
      ), f"idx = {idx}, num_params = {num_params}, sum = {idx + num_params}, offset = {offset}"
  
      # Special case for GPT models - whose last PP TP rank has a duplicate embedding tensor
  
      if "gpt" in model.cfg.target.lower():
          logging.info("Splitting GPT model")
          handler = GPTHandler(megatron_legacy=megatron_legacy)
  
      elif "t5" in model.cfg.target.lower():
          logging.info("Splitting T5 model")
          handler = T5Handler(megatron_legacy=megatron_legacy)
  
      else:
          raise ValueError(
              f"Unsupported model for Pipeline Parallelism change - {model.cfg.target}"
          )
  
      idx = handler.compute_split_index(
          model, idx, 0, pp_rank, pp_split_rank, tp_size, pp_size
      )
  
      # Print some debug info
      logging.info(
          f"Start Layer Idx: {idx} Number of layers in current rank: {num_params} Offset: {offset}"
      )
      logging.info("\n")
  
      # Split the model's parameters according to TP PP ranks
      idx, splits = handler.compute_splits(
          model,
          partitions,
          idx,
          0,
          pp_rank,
          pp_split_rank,
          tp_size,
          pp_size,
      )
  
      # Compute the new offset for the next PP rank in reverse order
      # Add 1 to offset to account for last PP rank's duplicated Embedding
      offset_diff = offset - num_params
      offset_diff = handler.compute_split_offset(
          offset_diff, 0, pp_rank, pp_split_rank, tp_size, pp_size
      )
  
      # Finalize the new offset
      new_offset = offset_diff
  
      # 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.
      write_tp_pp_split(
          model, splits, app_state, tp_size, pp_rank, write_path
      )
  
      return new_offset
  
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 9 for calling a global function with example usage: split_partition(model, partitions, pp_size, tp_size, pp_rank, offset, pp_split_rank, write_path, megatron_legacy) and returns: new_offset