code: def train_epoch(
      epoch,
      model,
      optimizer,
      _criterion,
      lr_scheduler,
      dataloader,
      booster,
      coordinator,
  ):
      torch.cuda.synchronize()
  
      use_pipeline = (
          isinstance(booster.plugin, HybridParallelPlugin)
          and booster.plugin.pp_size > 1
      )
      is_pp_last_stage = (
          use_pipeline and booster.plugin.stage_manager.is_last_stage()
      )
      total_step = len(dataloader)
  
      model.train()
      optimizer.zero_grad()
      dataloader = iter(dataloader)
      with tqdm(
          range(total_step),
          desc=f"Epoch [{epoch + 1}]",
          disable=not (coordinator.is_master() or is_pp_last_stage),
      ) as pbar:
          # Forward pass
          for _ in pbar:
              if use_pipeline:
                  outputs = booster.execute_pipeline(
                      dataloader,
                      model,
                      _criterion,
                      optimizer,
                      return_loss=True,
                      return_outputs=True,
                  )
                  # Backward and optimize
                  if is_pp_last_stage:
                      loss = outputs["loss"]
                      pbar.set_postfix({"loss": loss.item()})
              else:
                  data = next(dataloader)
                  data = move_to_cuda(data)
                  outputs = model(**data)
                  loss = _criterion(outputs, None)
                  # Backward
                  booster.backward(loss, optimizer)
                  pbar.set_postfix({"loss": loss.item()})
  
              optimizer.step()
              optimizer.zero_grad()
              lr_scheduler.step()
  
file_path: colossalai/examples/language/opt/opt_train_demo.py
imports: import datasets
  import torch
  import transformers
  from args import parse_demo_args
  from data import NetflixDataset, netflix_collator
  from tqdm import tqdm
  from transformers import AutoConfig, AutoTokenizer, OPTForCausalLM, get_linear_schedule_with_warmup
  from transformers.utils.versions import require_version
  import colossalai
  from colossalai.booster import Booster
  from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin, TorchDDPPlugin
  from colossalai.cluster import DistCoordinator
  from colossalai.logging import disable_existing_loggers, get_dist_logger
  from colossalai.nn.optimizer import HybridAdam
label_desc: Use image node_id 2 for calling a global function with example usage: train_epoch(epoch, model, optimizer, _criterion, lr_scheduler, dataloader, booster, coordinator) without return types