code: def partition_activations(
      args, cpu_checkpoint, contiguous_checkpoint
  ):
      global contiguous_data_buffers, data_offsets
  
      inputs = []
      num_non_fp_tensors = 0
  
      for arg_index, item in enumerate(args):
          if not is_activation_to_checkpoint(item):
              inputs.append(item)
              num_non_fp_tensors += 1
              continue
  
          i = arg_index - num_non_fp_tensors
          partition_size = get_partition_size(item)
          partition = (
              item.detach()
              .contiguous()
              .view(-1)
              .narrow(0, get_partition_start(item), partition_size)
              .clone()
          )
  
          buffer_device = (
              torch.device("cpu")
              if cpu_checkpoint
              else partition.device
          )
  
          if contiguous_checkpoint:
              if i >= len(contiguous_data_buffers):
                  tensor_list = [
                      torch.tensor(()).new_empty(
                          [partition_size],
                          dtype=partition.dtype,
                          device=buffer_device,
                      )
                      for _ in range(num_layers)
                  ]
                  contiguous_data_buffers.append(tensor_list)
                  data_offsets.append(0)
              elif contiguous_data_buffers[i] is None:
                  tensor_list = [
                      torch.tensor(()).new_empty(
                          [partition_size],
                          dtype=partition.dtype,
                          device=buffer_device,
                      )
                      for _ in range(num_layers)
                  ]
                  contiguous_data_buffers[i] = tensor_list
                  data_offsets[i] = 0
  
              # Because the 'new_empty' returns uninitialized pages,
              # the pages need to be populated during the cudaMemcpy time
              # which increases the data copy time. To avoid this, we
              # pre-populate these pages by simply writing 0 ahead of
              # the actual cudaMemcpy operation time. Due to the
              # previously launched GPU kernels, there is a small
              # window of time here for CPUs to populate pages asynchronously.
              contiguous_data_buffers[i][data_offsets[i]].data[
                  range(
                      0,
                      contiguous_data_buffers[i][
                          data_offsets[i]
                      ].data.shape[0],
                      int(
                          mmap.PAGESIZE
                          / contiguous_data_buffers[i][
                              data_offsets[i]
                          ].data.element_size()
                      ),
                  )
              ] = 0
  
              contiguous_partition = contiguous_data_buffers[i][
                  data_offsets[i]
              ].data.copy_(partition.data)
              data_offsets[i] = data_offsets[i] + 1
              inputs.append(contiguous_partition)
          else:
              partition = (
                  partition.cpu() if CPU_CHECKPOINT else partition
              )
              inputs.append(partition)
  
      return inputs
  
imports: import copy
  import torch
  import contextlib
  from deepspeed import comm
  import weakref
  import mmap
  from torch import _C
  from deepspeed.runtime.config import DeepSpeedConfig
  from deepspeed.utils import logger
  from deepspeed.runtime.utils import copy_to_device, move_to_device, see_memory_usage, bwc_tensor_model_parallel_rank
  from deepspeed.utils.timer import SynchronizedWallClockTimer, FORWARD_GLOBAL_TIMER
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 12 for calling a global function with example usage: partition_activations(args, cpu_checkpoint, contiguous_checkpoint) and returns: inputs