class_docstr: Plugin for Gemini.
  
  ```python
  from colossalai.booster import Booster
  from colossalai.booster.plugin import GeminiPlugin
  
  model, train_dataset, optimizer, criterion = ...
  plugin = GeminiPlugin()
  
  train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8)
  booster = Booster(plugin=plugin)
  model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion)
  ```
  
  Args:
      chunk_config_dict (dict, optional): chunk configuration dictionary.
      chunk_init_device (torch.device, optional): device to initialize the chunk.
      placement_policy (str, optional): "static" and "auto". Defaults to "static".
      enable_gradient_accumulation (bool, optional): Whether to enable gradient accumulation. When set to True, gradient will be stored after doing backward pass. Defaults to False.
      shard_param_frac (float, optional): fraction of parameters to be sharded. Only for "static" placement.
          If `shard_param_frac` is 1.0, it's equal to zero-3. If `shard_param_frac` is 0.0, it's equal to zero-2. Defaults to 1.0.
      offload_optim_frac (float, optional): fraction of optimizer states to be offloaded. Only for "static" placement.
          If `shard_param_frac` is 1.0 and `offload_optim_frac` is 0.0, it's equal to old "cuda" placement. Defaults to 0.0.
      offload_param_frac (float, optional): fraction of parameters to be offloaded. Only for "static" placement.
          For efficiency, this argument is useful only when `shard_param_frac` is 1.0 and `offload_optim_frac` is 1.0.
          If `shard_param_frac` is 1.0, `offload_optim_frac` is 1.0 and `offload_param_frac` is 1.0, it's equal to old "cpu" placement.
          When using static placement, we recommend users to tune `shard_param_frac` first and then `offload_optim_frac`.
          Defaults to 0.0.
      warmup_non_model_data_ratio (float, optional): ratio of expected non-model data memory during warmup. Only for "auto" placement. Defaults to 0.8.
      steady_cuda_cap_ratio (float, optional): ratio of allowed cuda capacity for model data during steady state. Only for "auto" placement. Defaults to 0.9.
      precision (str, optional): precision. Support 'fp16' and 'bf16'. Defaults to 'fp16'.
      master_weights (bool, optional): Whether to keep fp32 master parameter weights in optimizer. Defaults to True.
      pin_memory (bool, optional): use pin memory on CPU. Defaults to False.
      force_outputs_fp32 (bool, optional): force outputs are fp32. Defaults to False.
      strict_ddp_mode (bool, optional): use strict ddp mode (only use dp without other parallelism). Defaults to False.
      search_range_m (int, optional): chunk size searching range divided by 2^20. Defaults to 32.
      hidden_dim (int, optional): the hidden dimension of DNN.
          Users can provide this argument to speed up searching.
          If users do not know this argument before training, it is ok. We will use a default value 1024.
      min_chunk_size_m (float, optional): the minimum chunk size divided by 2^20.
          If the aggregate size of parameters is still smaller than the minimum chunk size,
          all parameters will be compacted into one small chunk.
      memstats (MemStats, optional) the memory statistics collector by a runtime memory tracer.
      gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward)
          which will be used when using hybrid CPU optimizer.
          This argument is meaningless when `placement_policy` of `GeminiManager` is not "auto".
          Defaults to 0.0.
      initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**16.
      min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1.
      growth_factor (float, optional): growth_factor used by DynamicGradScaler. Defaults to 2.
      backoff_factor (float, optional): backoff_factor used by DynamicGradScaler. Defaults to 0.5.
      growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000.
      hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2.
      max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32.
      max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do
          clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm.
      norm_type (float, optional): norm_type used for `clip_grad_norm`.
      tp_size (int, optional): If 'tp_size' is set to be greater than 1, it means using tensor parallelism strategy, which is implemented in Shardformer, 'tp_size' determines the size of the tensor parallel process group. Default to 1.
      extra_dp_size (int, optional): If 'extra_dp_size' is set to be greater than 1, it means creating another group to run with a ddp-like strategy. Default to 1.
      enable_all_optimization (bool, optional): Whether to switch on all the optimizations supported by Shardformer.
                                                  Currently all the optimization methods include fused normalization, flash attention and JIT.
                                                  Defaults to False.
      enable_fused_normalization (bool, optional): Whether to switch on fused normalization in Shardformer. Defaults to False.
      enable_flash_attention (bool, optional): Whether to switch on flash attention in Shardformer. Defaults to False.
      enable_jit_fused (bool, optional): Whether to switch on JIT in Shardformer. Default to False.
      enable_sequence_parallelism (bool): Whether to turn on sequence parallelism in Shardformer. Defaults to False.
      enable_sequence_overlap (bool): Whether to turn on sequence overlap in Shardformer. Defaults to False.
      verbose (bool, optional): verbose mode. Debug info including chunk search result will be printed. Defaults to False.
docstr: Prepare a dataloader for distributed training. The dataloader will be wrapped by
  `torch.utils.data.DataLoader` and `torch.utils.data.DistributedSampler`.
  
  
  Args:
      dataset (`torch.utils.data.Dataset`): The dataset to be loaded.
      shuffle (bool, optional): Whether to shuffle the dataset. Defaults to False.
      seed (int, optional): Random worker seed for sampling, defaults to 1024.
      add_sampler: Whether to add ``DistributedDataParallelSampler`` to the dataset. Defaults to True.
      drop_last (bool, optional): Set to True to drop the last incomplete batch, if the dataset size
          is not divisible by the batch size. If False and the size of dataset is not divisible by
          the batch size, then the last batch will be smaller, defaults to False.
      pin_memory (bool, optional): Whether to pin memory address in CPU memory. Defaults to False.
      num_workers (int, optional): Number of worker threads for this dataloader. Defaults to 0.
      kwargs (dict): optional parameters for ``torch.utils.data.DataLoader``, more details could be found in
              `DataLoader <https://pytorch.org/docs/stable/_modules/torch/utils/data/dataloader.html#DataLoader>`_.
  
  Returns:
      :class:`torch.utils.data.DataLoader`: A DataLoader used for training or testing.
docstr_tok: ["Prepare","a","dataloader","for","distributed","training",".","The","dataloader","will","be","wrapped","by","`","torch.utils.data.DataLoader","`","and","`","torch.utils.data.DistributedSampler","`",".","Args",":","dataset","(","`","torch.utils.data.Dataset","`",")",":","The","dataset","to","be","loaded",".","shuffle","(","bool",",","optional",")",":","Whether","to","shuffle","the","dataset",".","Defaults","to","False",".","seed","(","int",",","optional",")",":","Random","worker","seed","for","sampling",",","defaults","to","1024.","add_sampler",":","Whether","to","add","``","DistributedDataParallelSampler","``","to","the","dataset",".","Defaults","to","True",".","drop_last","(","bool",",","optional",")",":","Set","to","True","to","drop","the","last","incomplete","batch",",","if","the","dataset","size","is","not","divisible","by","the","batch","size",".","If","False","and","the","size","of","dataset","is","not","divisible","by","the","batch","size",",","then","the","last","batch","will","be","smaller",",","defaults","to","False",".","pin_memory","(","bool",",","optional",")",":","Whether","to","pin","memory","address","in","CPU","memory",".","Defaults","to","False",".","num_workers","(","int",",","optional",")",":","Number","of","worker","threads","for","this","dataloader",".","Defaults","to","0.","kwargs","(","dict",")",":","optional","parameters","for","``","torch.utils.data.DataLoader","``",",","more","details","could","be","found","in","`","DataLoader","<","https",":","\/\/pytorch.org\/docs\/stable\/_modules\/torch\/utils\/data\/dataloader.html","#","DataLoader",">","`","_",".","Returns",":",":","class",":","`","torch.utils.data.DataLoader","`",":","A","DataLoader","used","for","training","or","testing","."]
code: def prepare_dataloader(
      self,
      dataset,
      batch_size,
      shuffle=False,
      seed=1024,
      drop_last=False,
      pin_memory=False,
      num_workers=0,
      **kwargs,
  ):
      r"""
      Prepare a dataloader for distributed training. The dataloader will be wrapped by
      `torch.utils.data.DataLoader` and `torch.utils.data.DistributedSampler`.
  
  
      Args:
          dataset (`torch.utils.data.Dataset`): The dataset to be loaded.
          shuffle (bool, optional): Whether to shuffle the dataset. Defaults to False.
          seed (int, optional): Random worker seed for sampling, defaults to 1024.
          add_sampler: Whether to add ``DistributedDataParallelSampler`` to the dataset. Defaults to True.
          drop_last (bool, optional): Set to True to drop the last incomplete batch, if the dataset size
              is not divisible by the batch size. If False and the size of dataset is not divisible by
              the batch size, then the last batch will be smaller, defaults to False.
          pin_memory (bool, optional): Whether to pin memory address in CPU memory. Defaults to False.
          num_workers (int, optional): Number of worker threads for this dataloader. Defaults to 0.
          kwargs (dict): optional parameters for ``torch.utils.data.DataLoader``, more details could be found in
                  `DataLoader <https://pytorch.org/docs/stable/_modules/torch/utils/data/dataloader.html#DataLoader>`_.
  
      Returns:
          :class:`torch.utils.data.DataLoader`: A DataLoader used for training or testing.
      """
      _kwargs = kwargs.copy()
      zero_world_size = self.pg_mesh.size(ZERO_AXIS)
      extra_dp_world_size = self.pg_mesh.size(DP_AXIS)
      zero_rank = self.pg_mesh.coordinate(ZERO_AXIS)
      extra_dp_rank = self.pg_mesh.coordinate(DP_AXIS)
      sampler = DistributedSampler(
          dataset,
          num_replicas=zero_world_size * extra_dp_world_size,
          rank=zero_rank * extra_dp_world_size + extra_dp_rank,
          shuffle=shuffle,
      )
  
      # Deterministic dataloader
      def seed_worker(worker_id):
          worker_seed = seed
          np.random.seed(worker_seed)
          torch.manual_seed(worker_seed)
          random.seed(worker_seed)
  
      return DataLoader(
          dataset,
          batch_size=batch_size,
          sampler=sampler,
          worker_init_fn=seed_worker,
          drop_last=drop_last,
          pin_memory=pin_memory,
          num_workers=num_workers,
          **_kwargs,
      )
  
imports: import gc
  import logging
  import os
  import random
  from pathlib import Path
  from typing import Callable, Iterator, List, Optional, Tuple
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed.distributed_c10d import _get_default_group
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  from torch.utils.data.distributed import DistributedSampler
  from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO, GeneralCheckpointIO
  from colossalai.checkpoint_io.utils import get_model_base_filenames, get_optimizer_base_filenames, load_shard_state_dict, save_config_file, save_state_dict, save_state_dict_shards
  from colossalai.cluster import DistCoordinator, ProcessGroupMesh
  from colossalai.interface import ModelWrapper, OptimizerWrapper
  from colossalai.shardformer import ShardConfig, ShardFormer
  from colossalai.utils import get_current_device
  from colossalai.utils.device import IS_NPU_AVAILABLE
  from colossalai.zero import GeminiDDP, GeminiOptimizer
  from colossalai.zero.gemini.memory_tracer import MemStats
  from .dp_plugin_base import DPPluginBase
label_desc: Use image node_id 8 for calling the GeminiPlugin obj's underlying member method code with example usage: obj.prepare_dataloader(dataset, batch_size, shuffle, seed, drop_last, pin_memory, num_workers) and returns: DataLoader