code: def training_function(config, args):
      # Initialize accelerator
      if args.with_tracking:
          accelerator = Accelerator(
              cpu=args.cpu,
              mixed_precision=args.mixed_precision,
              log_with="all",
              logging_dir=args.logging_dir,
          )
      else:
          accelerator = Accelerator(
              cpu=args.cpu, mixed_precision=args.mixed_precision
          )
  
      if hasattr(args.checkpointing_steps, "isdigit"):
          if args.checkpointing_steps == "epoch":
              checkpointing_steps = args.checkpointing_steps
          elif args.checkpointing_steps.isdigit():
              checkpointing_steps = int(args.checkpointing_steps)
          else:
              raise ValueError(
                  f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
              )
      else:
          checkpointing_steps = None
      # 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"])
  
      # We need to initialize the trackers we use, and also store our configuration
      if args.with_tracking:
          run = os.path.split(__file__)[-1].split(".")[0]
          accelerator.init_trackers(run, config)
  
      tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
      metric = evaluate.load("glue", "mrpc")
  
      # If the batch size is too big we use gradient accumulation
      gradient_accumulation_steps = 1
      if (
          batch_size > MAX_GPU_BATCH_SIZE
          and accelerator.distributed_type != DistributedType.TPU
      ):
          gradient_accumulation_steps = batch_size // MAX_GPU_BATCH_SIZE
          batch_size = MAX_GPU_BATCH_SIZE
  
      def collate_fn(examples):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          if accelerator.distributed_type == DistributedType.TPU:
              return tokenizer.pad(
                  examples,
                  padding="max_length",
                  max_length=128,
                  return_tensors="pt",
              )
          return tokenizer.pad(
              examples, padding="longest", return_tensors="pt"
          )
  
      # load datasets
      train_dataset = load_from_disk(args.training_dir)
      validation_dataset = load_from_disk(args.validation_dir)
  
      accelerator.print(
          f" loaded train_dataset length is: {len(train_dataset)}"
      )
      accelerator.print(
          f" loaded test_dataset length is: {len(validation_dataset)}"
      )
  
      # Instantiate dataloaders.
      train_dataloader = DataLoader(
          train_dataset,
          shuffle=True,
          collate_fn=collate_fn,
          batch_size=batch_size,
      )
      eval_dataloader = DataLoader(
          validation_dataset,
          shuffle=False,
          collate_fn=collate_fn,
          batch_size=EVAL_BATCH_SIZE,
      )
  
      set_seed(seed)
  
      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      model = AutoModelForSequenceClassification.from_pretrained(
          "bert-base-cased", return_dict=True
      )
  
      # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
      # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
      # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
      model = model.to(accelerator.device)
  
      # Instantiate optimizer
      optimizer = AdamW(params=model.parameters(), lr=lr)
  
      # Instantiate scheduler
      lr_scheduler = get_linear_schedule_with_warmup(
          optimizer=optimizer,
          num_warmup_steps=100,
          num_training_steps=(len(train_dataloader) * num_epochs)
          // gradient_accumulation_steps,
      )
  
      # 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
  
      # Potentially load in the weights and states from a previous save
      if args.resume_from_checkpoint:
          if (
              args.resume_from_checkpoint is not None
              or args.resume_from_checkpoint != ""
          ):
              accelerator.print(
                  f"Resumed from checkpoint: {args.resume_from_checkpoint}"
              )
              accelerator.load_state(args.resume_from_checkpoint)
              path = os.path.basename(args.resume_from_checkpoint)
          else:
              # Get the most recent checkpoint
              dirs = [
                  f.name for f in os.scandir(os.getcwd()) if f.is_dir()
              ]
              dirs.sort(key=os.path.getctime)
              path = dirs[
                  -1
              ]  # Sorts folders by date modified, most recent checkpoint is the last
          # Extract `epoch_{i}` or `step_{i}`
          training_difference = os.path.splitext(path)[0]
  
          if "epoch" in training_difference:
              starting_epoch = (
                  int(training_difference.replace("epoch_", "")) + 1
              )
              resume_step = None
          else:
              resume_step = int(
                  training_difference.replace("step_", "")
              )
              starting_epoch = resume_step // len(train_dataloader)
              resume_step -= starting_epoch * len(train_dataloader)
  
      # Now we train the model
      for epoch in range(starting_epoch, num_epochs):
          model.train()
          if args.with_tracking:
              total_loss = 0
          for step, batch in enumerate(train_dataloader):
              # We need to skip steps until we reach the resumed step
              if (
                  args.resume_from_checkpoint
                  and epoch == starting_epoch
              ):
                  if resume_step is not None and step < resume_step:
                      overall_step += 1
                      continue
              # We could avoid this line since we set the accelerator with `device_placement=True`.
              batch.to(accelerator.device)
              outputs = model(**batch)
              loss = outputs.loss
              loss = loss / gradient_accumulation_steps
              # We keep track of the loss at each epoch
              if args.with_tracking:
                  total_loss += loss.detach().float()
              accelerator.backward(loss)
              if step % gradient_accumulation_steps == 0:
                  optimizer.step()
                  lr_scheduler.step()
                  optimizer.zero_grad()
  
              overall_step += 1
  
              if isinstance(checkpointing_steps, int):
                  output_dir = f"step_{overall_step}"
                  if overall_step % checkpointing_steps == 0:
                      if args.output_dir is not None:
                          output_dir = os.path.join(
                              args.output_dir, output_dir
                          )
                      accelerator.save_state(output_dir)
  
          model.eval()
          for step, batch in enumerate(eval_dataloader):
              # We could avoid this line since we set the accelerator with `device_placement=True`.
              batch.to(accelerator.device)
              with torch.no_grad():
                  outputs = model(**batch)
              predictions = outputs.logits.argmax(dim=-1)
              predictions, references = accelerator.gather_for_metrics(
                  (predictions, batch["labels"])
              )
              metric.add_batch(
                  predictions=predictions,
                  references=references,
              )
  
          eval_metric = metric.compute()
          # Use accelerator.print to print only on the main process.
          accelerator.print(f"epoch {epoch}:", eval_metric)
          if args.with_tracking:
              accelerator.log(
                  {
                      "accuracy": eval_metric["accuracy"],
                      "f1": eval_metric["f1"],
                      "train_loss": total_loss.item()
                      / len(train_dataloader),
                      "epoch": epoch,
                  },
                  step=epoch,
              )
  
          if checkpointing_steps == "epoch":
              output_dir = f"epoch_{epoch}"
              if args.output_dir is not None:
                  output_dir = os.path.join(args.output_dir, output_dir)
              accelerator.save_state(output_dir)
  
      accelerator.save(
          accelerator.get_state_dict(model),
          os.path.join(args.output_dir, "model.pt"),
      )
      if args.with_tracking:
          accelerator.end_training()
  
file_path: hf-notebooks/sagemaker/22_accelerate_sagemaker_examples/src/text-classification/train_using_s3_data.py
imports: import argparse
  import os
  import torch
  from torch.optim import AdamW
  from torch.utils.data import DataLoader
  import evaluate
  from accelerate import Accelerator, DistributedType
  from datasets import load_from_disk
  from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
label_desc: Use image node_id 1 for calling a global function with example usage: training_function(config, args) and returns: tokenizer, tokenizer