code: def save(
      booster: Booster,
      model: nn.Module,
      optimizer: Optimizer,
      lr_scheduler: _LRScheduler,
      epoch: int,
      step: int,
      batch_size: int,
      coordinator: DistCoordinator,
      save_dir: str,
  ):
      save_dir = os.path.join(save_dir, f"epoch{epoch}-step{step}")
      os.makedirs(os.path.join(save_dir, "model"), exist_ok=True)
  
      booster.save_model(
          model, os.path.join(save_dir, "model"), shard=True
      )
      booster.save_optimizer(
          optimizer, os.path.join(save_dir, "optimizer"), shard=True
      )
      booster.save_lr_scheduler(
          lr_scheduler, os.path.join(save_dir, "lr_scheduler")
      )
      running_states = {
          "epoch": epoch,
          "step": step,
          "sample_start_index": step * batch_size,
      }
      if coordinator.is_master():
          save_json(
              running_states,
              os.path.join(save_dir, "running_states.json"),
          )
  
imports: import argparse
  import os
  import resource
  from contextlib import nullcontext
  from functools import partial
  from typing import Optional, Tuple
  import torch
  import torch.distributed
  import torch.nn
  from attn import SUPPORT_XFORMERS, replace_xformers
  from data_utils import load_json, prepare_dataloader, save_json
  from datasets import load_dataset
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.tensorboard import SummaryWriter
  from tqdm import tqdm
  from transformers.models.llama.configuration_llama import LlamaConfig
  from transformers.models.llama.modeling_llama import LlamaForCausalLM
  from transformers.models.llama.tokenization_llama import LlamaTokenizer
  import colossalai
  from colossalai.booster import Booster
  from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin
  from colossalai.cluster import DistCoordinator
  from colossalai.lazy import LazyInitContext
  from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
  from colossalai.nn.optimizer import HybridAdam
  from colossalai.utils import get_current_device
label_desc: Use image node_id 5 for calling a global function with example usage: save(booster, model, optimizer, lr_scheduler, epoch, step, batch_size, coordinator, save_dir) without return types