code: def training_function(kwargs: dict):
      print("training_function called")
  
      # Train has a bug somewhere that causes ACCELERATE_TORCH_DEVICE to not be set
      # properly on multi-gpu nodes
      cuda_visible_device = os.environ["CUDA_VISIBLE_DEVICES"].split(
          ","
      )
      local_rank = int(os.environ["LOCAL_RANK"])
      device_id = cuda_visible_device[local_rank]
      os.environ["ACCELERATE_TORCH_DEVICE"] = f"cuda:{device_id}"
  
      config = kwargs["config"]
      args = argparse.Namespace(**kwargs["args"])
      special_tokens = kwargs.get("special_tokens", [])
      model_id = config["model_name"]
  
      # We need to download the model weights on this machine if they don't exit.
      # We need to acquire a lock to ensure that only one process downloads the model
      bucket_uri = get_mirror_link(model_id)
      download_path = get_download_path(model_id)
      base_path = Path(download_path).parent
      base_path.mkdir(parents=True, exist_ok=True)
      lock_file = str(
          base_path / f'{model_id.replace("/",  "--")}.lock'
      )
      with FileLock(lock_file):
          download_model(
              model_id=model_id,
              bucket_uri=bucket_uri,
              s3_sync_args=["--no-sign-request"],
          )
  
      # 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"])
      gradient_accumulation_steps = int(
          config["gradient_accumulation_steps"]
      )
  
      # Get deepspeed config to setup the batch size per device
      ds_plugin = config["ds_plugin"]
      ds_plugin.hf_ds_config.config[
          "train_micro_batch_size_per_gpu"
      ] = batch_size
  
      # Initialize accelerator
      accelerator = Accelerator(
          deepspeed_plugin=ds_plugin,
          gradient_accumulation_steps=gradient_accumulation_steps,
          mixed_precision=args.mx,
      )
  
      set_seed(seed)
  
      # train_ds is the local shard for this model
      train_ds = train.get_dataset_shard("train")
      valid_ds = train.get_dataset_shard("valid")
  
      train_ds_len = len(list(train_ds.iter_batches(batch_size=1)))
  
      _test_tokenizer(args.model_name)
      tokenizer = get_tokenizer(
          model_name=args.model_name, special_tokens=special_tokens
      )
      collate_partial = functools.partial(
          collate_fn,
          tokenizer=tokenizer,
          block_size=config["block_size"],
          device=accelerator.device,
      )
  
      pretrained_path = get_pretrained_path(model_id)
      print(f"Loading model from {pretrained_path} ...")
      s = time.time()
      model = AutoModelForCausalLM.from_pretrained(
          pretrained_path,
          trust_remote_code=True,
          torch_dtype=torch.bfloat16,
          # `use_cache=True` is incompatible with gradient checkpointing.
          use_cache=False,
          use_flash_attention_2=True,
      )
      print(f"Done loading model in {time.time() - s} seconds.")
  
      model.resize_token_embeddings(len(tokenizer))
  
      if config["lora"]:
          # Apply LoRA
          s = time.time()
          lora_config = LoraConfig(**config["lora_config"])
  
          expected_num_parameters = get_expected_lora_num_parameters(
              lora_config=lora_config, model=model
          )
  
          print(f"Attempting to apply LoRA config: {lora_config}")
  
          model.enable_input_require_grads()
          model = get_peft_model(model, lora_config)
  
          num_parameters = get_number_of_params(model)
  
          if num_parameters != expected_num_parameters:
              raise ValueError(
                  f"Expected {expected_num_parameters} parameters, got {num_parameters} "
                  f"parameters. LoRA-ification failed."
              )
  
          print(
              f"LoRA-ification done in {time.time() - s} seconds. Estimated checkpoint "
              f"size (fp16): {num_parameters * 2 / 1e6} MB"
          )
  
      print(
          f"Number of checkpointed parameters: {get_number_of_params(model)}"
      )
  
      print(
          "Model initialized with pretrained weights. Training starting..."
      )
      if not args.no_grad_ckpt:
          model.gradient_checkpointing_enable()
  
      optimizer_cls = (
          torch.optim.AdamW
          if accelerator.state.deepspeed_plugin is None
          or "optimizer"
          not in accelerator.state.deepspeed_plugin.deepspeed_config
          else DummyOptim
      )
  
      optimizer = optimizer_cls(
          model.parameters(),
          lr=lr,
          betas=OPTIM_BETAS,
          weight_decay=OPTIM_WEIGHT_DECAY,
          eps=OPTIM_EPS,
      )
  
      # Instantiate scheduler
      # Creates Dummy Scheduler if `scheduler` was specified in the config file or
      # else, creates `args.lr_scheduler_type` Scheduler
      # get train and valid dataset lengths
  
      num_steps_per_epoch = math.ceil(
          train_ds_len / args.batch_size_per_device
      )
      total_training_steps = (
          num_steps_per_epoch
          * num_epochs
          // gradient_accumulation_steps
      )
  
      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=NUM_WARMUP_STEPS * args.num_devices,
              num_training_steps=total_training_steps
              * args.num_devices,
          )
      else:
          lr_scheduler = DummyScheduler(
              optimizer,
              warmup_num_steps=NUM_WARMUP_STEPS * args.num_devices,
              total_num_steps=total_training_steps * args.num_devices,
          )
  
      # 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.
      s = time.time()
      model, optimizer, lr_scheduler = accelerator.prepare(
          model, optimizer, lr_scheduler
      )
      print(f"Prepare done in {time.time() - s} seconds.")
  
      # Now we train the model
      if accelerator.is_main_process:
          print("Starting training ...")
          print(
              "Number of batches on main process",
              train_ds_len // batch_size,
          )
  
      for epoch in range(num_epochs):
          fwd_time_sum, bwd_time_sum, optim_step_time_sum = 0, 0, 0
          s_epoch = time.time()
          model.train()
          loss_sum = torch.tensor(0.0).to(accelerator.device)
  
          train_dataloader = train_ds.iter_torch_batches(
              batch_size=batch_size,
              collate_fn=collate_partial,
          )
  
          for step, batch in tqdm.tqdm(
              enumerate(train_dataloader),
              total=train_ds_len // batch_size + 1,
          ):
              # We could avoid this line since we set the accelerator with
              # `device_placement=True`.
              with accelerator.accumulate(model):
                  s_fwd = time.time()
                  outputs = model(**batch)
                  loss = outputs.loss
                  loss_sum += loss.item()
                  e_fwd = time.time()
                  fwd_time = e_fwd - s_fwd
                  fwd_time_sum += fwd_time
                  s_bwd = time.time()
                  accelerator.backward(loss)
                  e_bwd = time.time()
                  bwd_time = e_bwd - s_bwd
                  bwd_time_sum += bwd_time
  
                  s_opt_step = time.time()
                  optimizer.step()
                  lr_scheduler.step()
                  optimizer.zero_grad()
                  e_opt_step = time.time()
                  optim_step_time_sum += e_opt_step - s_opt_step
  
              if accelerator.is_main_process:
                  accelerator.print(
                      f"[epoch {epoch} step {step}] "
                      f"loss: {loss.item()} step-time: {e_opt_step - s_fwd}"
                  )
  
              aggregated_loss = torch.mean(
                  accelerator.gather(loss[None])
              ).item()
  
              if config["as_test"]:
                  break
  
              # as long as this is not the last step report here
              if step != (train_ds_len // batch_size - 1):
                  train.report(
                      {
                          "epoch": epoch,
                          "iteration": step,
                          "train_loss_batch": aggregated_loss,
                          "avg_train_loss_epoch": None,
                          "eval_loss": None,
                          "perplexity": None,
                          "num_iterations": step + 1,
                          "train_time_per_epoch": None,
                          "eval_time_per_epoch": None,
                          "fwd_time": fwd_time,
                          "bwd_time": bwd_time,
                          "avg_fwd_time_per_epoch": None,
                          "avg_bwd_time_per_epoch": None,
                          "learning_rate": lr_scheduler.get_lr()[0],
                      }
                  )
  
          e_epoch = time.time()
          accelerator.print("Train time per epoch: ", e_epoch - s_epoch)
  
          eval_s_epoch = time.time()
          print("Running evaluation ...")
          perplex, eloss = evaluate(
              model=model,
              eval_ds=valid_ds,
              accelerator=accelerator,
              bsize=config["eval_batch_size"],
              ds_kwargs={"collate_fn": collate_partial},
              as_test=config["as_test"],
          )
          accelerator.print("Eval result loss", eloss)
          accelerator.print("Eval perplex", perplex)
  
          eval_e_epoch = time.time()
          accelerator.print(
              "Eval time per epoch: ", eval_e_epoch - eval_s_epoch
          )
          accelerator.print("avg fwd time: ", fwd_time_sum / (step + 1))
          accelerator.print("avg bwd time: ", bwd_time_sum / (step + 1))
          accelerator.print(
              "avg opt step time: ", optim_step_time_sum / (step + 1)
          )
  
          metrics = {
              "epoch": epoch,
              "iteration": step,
              "train_loss_batch": aggregated_loss,
              "avg_train_loss_epoch": loss_sum.item() / (step + 1),
              "eval_loss": eloss,
              "perplexity": perplex,
              "num_iterations": step + 1,
              "train_time_per_epoch": e_epoch - s_epoch,
              "eval_time_per_epoch": eval_e_epoch - eval_s_epoch,
              "fwd_time": fwd_time,
              "bwd_time": bwd_time,
              "avg_fwd_time_per_epoch": fwd_time_sum / (step + 1),
              "avg_bwd_time_per_epoch": bwd_time_sum / (step + 1),
              "learning_rate": lr_scheduler.get_lr()[0],
          }
  
          with tempfile.TemporaryDirectory(
              dir=args.output_dir
          ) as temp_checkpoint_dir:
              accelerator.print(
                  f"Saving the model locally at {temp_checkpoint_dir}"
              )
              accelerator.wait_for_everyone()
  
              checkpoint_save_start = time.perf_counter()
  
              if accelerator.is_main_process:
                  print("Saving tokenizer and config.")
                  tokenizer.save_pretrained(temp_checkpoint_dir)
  
              accelerator.wait_for_everyone()
  
              # Checkpointing strategy 1: Distributed checkpointing
              # This checkpointing method makes deepspeed checkpoints on each node
              # and then Ray Train will aggregate them to a central s3 bucket.
              # It should be done on all processes (not just the Rank 0)
              # aggregate_on_rank_0 = False
              # checkpoint_model(
              #     checkpoint_folder=tempdir,
              #     ckpt_id=epoch,
              #     model=model,
              #     epoch=epoch,
              #     last_global_step=step
              # )
  
              # Checkpointing strategy 2: Aggregate model on the rank 0 worker then upload
              aggregate_on_rank_0 = True
              unwrapped_model = accelerator.unwrap_model(model)
              unwrapped_model.save_pretrained(
                  temp_checkpoint_dir,
                  is_main_process=accelerator.is_main_process,
                  save_function=accelerator.save,
                  safe_serialization=True,
                  state_dict=accelerator.get_state_dict(model),
              )
              accelerator.wait_for_everyone()
              print(
                  "Checkpoint save time: ",
                  time.perf_counter() - checkpoint_save_start,
              )
  
              checkpoint_upload_start = time.perf_counter()
  
              # Create the checkpoint object to report to Ray Train and upload to storage.
              # If we aggregated the model on rank 0, we only need to report
              # the checkpoint from the rank 0 worker, since all other checkpoint
              # directories are empty (`save_pretrained` was a noop for other workers).
              if aggregate_on_rank_0:
                  checkpoint = (
                      Checkpoint.from_directory(temp_checkpoint_dir)
                      if accelerator.is_main_process
                      else None
                  )
              else:
                  # Distributed checkpointing should upload shards from each worker.
                  checkpoint = Checkpoint.from_directory(
                      temp_checkpoint_dir
                  )
  
              # Note: After `train.report`, in the case of remote storage,
              # the checkpoint directory will be uploaded to the remote storage.
              train.report(metrics, checkpoint=checkpoint)
  
              print(
                  "Checkpoint upload time: ",
                  time.perf_counter() - checkpoint_upload_start,
              )
              print(
                  "Total checkpointing time: ",
                  time.perf_counter() - checkpoint_save_start,
              )
  
          if perplex < args.stop_perplexity:
              print(
                  f"Perplexity reached {perplex} < {args.stop_perplexity}. Stopping."
              )
              break
  
          if config["as_test"]:
              break
  
imports: import argparse
  from filelock import FileLock
  import functools
  import json
  import math
  import os
  from pathlib import Path
  import re
  import tempfile
  import time
  import tree
  from typing import Tuple
  from accelerate import Accelerator, DeepSpeedPlugin
  from accelerate.utils import DummyOptim, DummyScheduler, set_seed
  import torch
  import torch.nn
  import tqdm
  from transformers import AutoModelForCausalLM, AutoTokenizer, get_linear_schedule_with_warmup
  from peft import LoraConfig, get_peft_model
  import ray
  from ray import train
  import ray.util.scheduling_strategies
  from ray.train.torch import TorchTrainer
  from ray.train import Checkpoint
  from utils import get_checkpoint_and_refs_dir, get_mirror_link, download_model, get_download_path
label_desc: Use image node_id 9 for calling a global function with example usage: training_function(kwargs) without return types