code: def save_checkpoint(
      save_dir: Union[str, os.PathLike],
      booster: Booster,
      model: torch.nn.Module,
      optimizer: Optimizer,
      lr_scheduler: _LRScheduler,
      epoch: int,
      step: int,
      batch_size: int,
      coordinator: DistCoordinator,
  ) -> None:
      """
      Save model checkpoint, optimizer, LR scheduler and intermedidate running states.
      """
  
      save_dir = os.path.join(save_dir, f"epoch-{epoch}_step-{step}")
      os.makedirs(os.path.join(save_dir, "modeling"), exist_ok=True)
  
      booster.save_model(
          model, os.path.join(save_dir, "modeling"), 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"),
          )
  
file_path: ColossalAI/applications/Colossal-LLaMA-2/colossal_llama2/utils/ckpt_io.py
imports: import json
  import os
  from typing import Any, Dict, Tuple, Union
  import torch
  from torch.optim.optimizer import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from colossalai.booster import Booster
  from colossalai.cluster import DistCoordinator
label_desc: Use image node_id 3 for calling a global function with example usage: save_checkpoint(save_dir, booster, model, optimizer, lr_scheduler, epoch, step, batch_size, coordinator) without return types