code: def __init__(
      self,
      config,
      env,
      path,
      policy,
      pipeline_rank,
      num_pipeline_stages,
      comm_device,
      num_inner_iterations=None,
      async_comm=False,
  ):
      self.config = config
      self.env = env
      self.path = path
      self.policy = policy
      self.num_gpu_batches = self.policy.num_gpu_batches
      self.pipeline_rank = pipeline_rank
      self.num_pipeline_stages = num_pipeline_stages
      self.num_inner_iterations = (
          num_inner_iterations
          if num_inner_iterations is not None
          else num_pipeline_stages
      )
      self.async_comm = async_comm
      if comm_device == "cpu":
          self.comm_device = self.env.cpu
      elif comm_device == "gpu":
          self.comm_device = self.env.gpu
      else:
          raise ValueError(f"Invalid comm_device: {comm_device}")
  
      layers = []
      if pipeline_rank == 0:
          layers.append(InputEmbed(self.config, self.env, self.policy))
      pipeline_stage_sizes = [
          config.num_hidden_layers // num_pipeline_stages
          + int(i < config.num_hidden_layers % num_pipeline_stages)
          for i in range(num_pipeline_stages)
      ]
      layer_start_ids = [0]
      for stage_size in pipeline_stage_sizes:
          layer_start_ids.append(layer_start_ids[-1] + stage_size)
      for i in range(
          layer_start_ids[pipeline_rank],
          layer_start_ids[pipeline_rank + 1],
      ):
          if self.policy.sep_layer:
              layers.append(
                  SelfAttention(self.config, self.env, self.policy, i)
              )
              layers.append(MLP(self.config, self.env, self.policy, i))
          else:
              layers.append(
                  TransformerLayer(
                      self.config, self.env, self.policy, i
                  )
              )
      if pipeline_rank == num_pipeline_stages - 1:
          layers.append(OutputEmbed(self.config, self.env, self.policy))
      self.layers = layers
      self.num_layers = len(layers)
  
      if self.policy.act_gpu_percent == 100:
          self.act_home = self.env.gpu
      elif self.policy.act_cpu_percent == 100:
          self.act_home = self.env.cpu
      elif self.policy.act_disk_percent == 100:
          self.act_home = self.env.disk
      else:
          raise NotImplementedError()
  
      # CUDA streams
      self.load_weight_stream = torch.cuda.Stream()
      self.load_cache_stream = torch.cuda.Stream()
      self.store_cache_stream = torch.cuda.Stream()
  
      self.task = None
      self.init_all_weights()
  
imports: import argparse
  from itertools import count
  import os
  import pickle
  import traceback
  from typing import Union, List, Optional
  import numpy
  import torch
  import torch.distributed
  from transformers import AutoTokenizer
  from flexgen.compression import CompressionConfig
  from flexgen.dist_utils import initialize_distributed
  from flexgen.flex_opt import Policy, InputEmbed, OutputEmbed, SelfAttention, MLP, TransformerLayer, OptLM, get_filename, add_parser_arguments, get_test_inputs, DUMMY_WEIGHT
  from flexgen.opt_config import get_opt_config
  from flexgen.pytorch_backend import TorchDevice, TorchDisk, TorchLink, TorchMixedDevice, TorchTensor
  from flexgen.timer import timers
  from flexgen.utils import Task, ExecutionEnv, GB, T, ValueHolder, array_1d, array_2d, array_3d, array_4d, str2bool, project_decode_latency
label_desc: Use image node_id 1 to create a new DistOptLM object from inherited base classes: OptLM with example: obj = DistOptLM(config, env, path, policy, pipeline_rank, num_pipeline_stages, comm_device, num_inner_iterations, async_comm)