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.
code: def __init__(
      self,
      chunk_config_dict: Optional[dict] = None,
      chunk_init_device: Optional[torch.device] = None,
      placement_policy: str = "static",
      enable_gradient_accumulation: bool = False,
      shard_param_frac: float = 1.0,  # only for static placement
      offload_optim_frac: float = 0.0,  # only for static placement
      offload_param_frac: float = 0.0,  # only for static placement
      warmup_non_model_data_ratio: float = 0.8,  # only for auto placement
      steady_cuda_cap_ratio: float = 0.9,  # only for auto placement
      precision: str = "fp16",
      master_weights: bool = True,
      pin_memory: bool = False,
      force_outputs_fp32: bool = False,
      strict_ddp_mode: bool = False,
      search_range_m: int = 32,
      hidden_dim: Optional[int] = None,
      min_chunk_size_m: float = 32,
      memstats: Optional[MemStats] = None,
      gpu_margin_mem_ratio: float = 0.0,
      initial_scale: float = 2**16,
      min_scale: float = 1,
      growth_factor: float = 2,
      backoff_factor: float = 0.5,
      growth_interval: int = 1000,
      hysteresis: int = 2,
      max_scale: float = 2**32,
      max_norm: float = 0.0,
      norm_type: float = 2.0,
      tp_size: int = 1,
      extra_dp_size: int = 1,
      enable_all_optimization: bool = False,
      enable_fused_normalization: bool = False,
      enable_flash_attention: bool = False,
      enable_sequence_parallelism: bool = False,
      enable_jit_fused: bool = False,
      enable_sequence_overlap: bool = False,
      verbose: bool = False,
  ) -> None:
      super().__init__()
      assert (
          precision in SUPPORTED_PRECISION
      ), f"precision {precision} is not supported"
      if IS_NPU_AVAILABLE:
          assert (
              placement_policy == "static"
          ), "NPU only supports static placement policy"
      self.gemini_config = dict(
          chunk_config_dict=chunk_config_dict,
          chunk_init_device=(chunk_init_device or get_current_device()),
          placement_policy=placement_policy,
          enable_gradient_accumulation=enable_gradient_accumulation,
          shard_param_frac=shard_param_frac,
          offload_optim_frac=offload_optim_frac,
          offload_param_frac=offload_param_frac,
          warmup_non_model_data_ratio=warmup_non_model_data_ratio,
          steady_cuda_cap_ratio=steady_cuda_cap_ratio,
          pin_memory=pin_memory,
          force_outputs_fp32=force_outputs_fp32,
          strict_ddp_mode=strict_ddp_mode,
          search_range_m=search_range_m,
          hidden_dim=hidden_dim,
          min_chunk_size_m=min_chunk_size_m,
          memstats=memstats,
          mixed_precision=PRECISION_STR_TO_DTYPE[precision],
          master_weights=master_weights,
      )
      self.zero_optim_config = dict(
          gpu_margin_mem_ratio=gpu_margin_mem_ratio,
      )
      self.optim_kwargs = dict(
          initial_scale=initial_scale,
          growth_factor=growth_factor,
          backoff_factor=backoff_factor,
          growth_interval=growth_interval,
          hysteresis=hysteresis,
          min_scale=min_scale,
          max_scale=max_scale,
          max_norm=max_norm,
          norm_type=norm_type,
      )
      self.enable_tensor_parallelism = tp_size > 1
      self.enable_all_optimization = enable_all_optimization
      self.enable_fused_normalization = enable_fused_normalization
      self.enable_flash_attention = enable_flash_attention
      self.enable_sequence_parallelism = (
          enable_sequence_parallelism
          if self.enable_tensor_parallelism
          else False
      )
      self.enable_jit_fused = enable_jit_fused
      self.enable_sequence_overlap = enable_sequence_overlap
      self.verbose = verbose
  
      self.tp_size = tp_size
      self.extra_dp_size = extra_dp_size
      world_size = dist.get_world_size()
      self.zero_size = world_size // (self.tp_size * self.extra_dp_size)
      assert (
          world_size
          == (self.tp_size * self.extra_dp_size) * self.zero_size
      ), f"The global group size can't be evenly divided by the subgroup size."
  
      self.pg_mesh = ProcessGroupMesh(
          self.zero_size, self.extra_dp_size, self.tp_size
      )
      self.zero_group = (
          self.pg_mesh.get_group_along_axis(ZERO_AXIS)
          if self.zero_size < world_size
          else _get_default_group()
      )
      self.extra_dp_group = (
          self.pg_mesh.get_group_along_axis(DP_AXIS)
          if self.extra_dp_size > 1
          else None
      )
      self.tp_group = (
          self.pg_mesh.get_group_along_axis(TP_AXIS)
          if self.tp_size > 1
          else None
      )
  
      self.shard_config = ShardConfig(
          tensor_parallel_process_group=self.tp_group,
          enable_tensor_parallelism=self.enable_tensor_parallelism,
          enable_all_optimization=self.enable_all_optimization,
          enable_fused_normalization=self.enable_fused_normalization,
          enable_flash_attention=self.enable_flash_attention,
          enable_jit_fused=self.enable_jit_fused,
          enable_sequence_parallelism=self.enable_sequence_parallelism,
          enable_sequence_overlap=self.enable_sequence_overlap,
      )
  
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 1 to create a new GeminiPlugin object from inherited base classes: DPPluginBase with example: obj = GeminiPlugin(chunk_config_dict, chunk_init_device, placement_policy, enable_gradient_accumulation, shard_param_frac, offload_optim_frac, offload_param_frac, warmup_non_model_data_ratio, steady_cuda_cap_ratio, precision, master_weights, pin_memory, force_outputs_fp32, strict_ddp_mode, search_range_m, hidden_dim, min_chunk_size_m, memstats, gpu_margin_mem_ratio, initial_scale, min_scale, growth_factor, backoff_factor, growth_interval, hysteresis, max_scale, max_norm, norm_type, tp_size, extra_dp_size, enable_all_optimization, enable_fused_normalization, enable_flash_attention, enable_sequence_parallelism, enable_jit_fused, enable_sequence_overlap, verbose)