code: def training_function(config, args):
      # Initialize accelerator
      accelerator = Accelerator()
  
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      lr = config["lr"]
      num_epochs = int(config["num_epochs"])
      seed = int(config["seed"])
      batch_size = int(config["batch_size"])
      model_name = args.model_name_or_path
  
      set_seed(seed)
      train_dataloader, eval_dataloader = get_dataloaders(
          accelerator, batch_size, model_name
      )
  
      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      model = AutoModelForSequenceClassification.from_pretrained(
          model_name, return_dict=True
      )
  
      # Instantiate optimizer
      optimizer_cls = (
          AdamW
          if accelerator.state.deepspeed_plugin is None
          or "optimizer"
          not in accelerator.state.deepspeed_plugin.deepspeed_config
          else DummyOptim
      )
      optimizer = optimizer_cls(params=model.parameters(), lr=lr)
  
      if accelerator.state.deepspeed_plugin is not None:
          gradient_accumulation_steps = (
              accelerator.state.deepspeed_plugin.deepspeed_config[
                  "gradient_accumulation_steps"
              ]
          )
      else:
          gradient_accumulation_steps = 1
      max_training_steps = (
          len(train_dataloader) * num_epochs
      ) // gradient_accumulation_steps
  
      # Instantiate scheduler
      if (
          accelerator.state.deepspeed_plugin is None
          or "scheduler"
          not in accelerator.state.deepspeed_plugin.deepspeed_config
      ):
          lr_scheduler = get_linear_schedule_with_warmup(
              optimizer=optimizer,
              num_warmup_steps=0,
              num_training_steps=max_training_steps,
          )
      else:
          lr_scheduler = DummyScheduler(
              optimizer,
              total_num_steps=max_training_steps,
              warmup_num_steps=0,
          )
  
      # Prepare everything
      # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
      # prepare method.
      (
          model,
          optimizer,
          train_dataloader,
          eval_dataloader,
          lr_scheduler,
      ) = accelerator.prepare(
          model,
          optimizer,
          train_dataloader,
          eval_dataloader,
          lr_scheduler,
      )
  
      # We need to keep track of how many total steps we have iterated over
      overall_step = 0
      # We also need to keep track of the stating epoch so files are named properly
      starting_epoch = 0
      metric = evaluate.load("glue", "mrpc")
      ending_epoch = num_epochs
  
      if args.partial_train_epoch is not None:
          ending_epoch = args.partial_train_epoch
  
      if args.resume_from_checkpoint:
          accelerator.load_state(args.resume_from_checkpoint)
          epoch_string = args.resume_from_checkpoint.split("epoch_")[1]
          state_epoch_num = ""
          for char in epoch_string:
              if char.isdigit():
                  state_epoch_num += char
              else:
                  break
          starting_epoch = int(state_epoch_num) + 1
          accuracy = evaluation_loop(
              accelerator, model, eval_dataloader, metric
          )
          accelerator.print("resumed checkpoint performance:", accuracy)
          accelerator.print(
              "resumed checkpoint's scheduler's lr:",
              lr_scheduler.get_lr()[0],
          )
          accelerator.print(
              "resumed optimizers's lr:",
              optimizer.param_groups[0]["lr"],
          )
          with open(
              os.path.join(
                  args.output_dir, f"state_{starting_epoch-1}.json"
              ),
              "r",
          ) as f:
              resumed_state = json.load(f)
              assert (
                  resumed_state["accuracy"] == accuracy
              ), "Accuracy mismatch, loading from checkpoint failed"
              assert (
                  resumed_state["lr"] == lr_scheduler.get_lr()[0]
              ), "Scheduler learning rate mismatch, loading from checkpoint failed"
              assert (
                  resumed_state["optimizer_lr"]
                  == optimizer.param_groups[0]["lr"]
              ), "Optimizer learning rate mismatch, loading from checkpoint failed"
              assert (
                  resumed_state["epoch"] == starting_epoch - 1
              ), "Epoch mismatch, loading from checkpoint failed"
              return
  
      # Now we train the model
      state = {}
      for epoch in range(starting_epoch, ending_epoch):
          model.train()
          for step, batch in enumerate(train_dataloader):
              outputs = model(**batch)
              loss = outputs.loss
              loss = loss / gradient_accumulation_steps
              accelerator.backward(loss)
              if step % gradient_accumulation_steps == 0:
                  optimizer.step()
                  lr_scheduler.step()
                  optimizer.zero_grad()
  
              overall_step += 1
          output_dir = f"epoch_{epoch}"
          output_dir = os.path.join(args.output_dir, output_dir)
          accelerator.save_state(output_dir)
          accuracy = evaluation_loop(
              accelerator, model, eval_dataloader, metric
          )
          state["accuracy"] = accuracy
          state["lr"] = lr_scheduler.get_lr()[0]
          state["optimizer_lr"] = optimizer.param_groups[0]["lr"]
          state["epoch"] = epoch
          state["step"] = overall_step
          accelerator.print(f"epoch {epoch}:", state)
  
          accelerator.wait_for_everyone()
          if accelerator.is_main_process:
              with open(
                  os.path.join(args.output_dir, f"state_{epoch}.json"),
                  "w",
              ) as f:
                  json.dump(state, f)
  
file_path: accelerate/src/accelerate/test_utils/scripts/external_deps/test_checkpointing.py
imports: import argparse
  import json
  import os
  import evaluate
  import torch
  from datasets import load_dataset
  from torch.optim import AdamW
  from torch.utils.data import DataLoader
  from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
  from accelerate import Accelerator, DistributedType
  from accelerate.utils.deepspeed import DummyOptim, DummyScheduler
label_desc: Use image node_id 3 for calling a global function with example usage: training_function(config, args) without return types