class_docstr: This class mirrors deepspeed functionality. Not all changes are implemented in this version.
  
  See TorchEngineFull for more modifications.
class_docstr_tok: ["This","class","mirrors","deepspeed","functionality",".","Not","all","changes","are","implemented","in","this","version",".","See","TorchEngineFull","for","more","modifications","."]
docstr: Load Engine. The model will be compiled by default.
docstr_tok: ["Load","Engine",".","The","model","will","be","compiled","by","default","."]
code: def __init__(
      self,
      model,
      cfg_train,
      cfg_impl,
      setup=_default_setup,
      seq_length=128,
  ):
      """Load Engine. The model will be compiled by default."""
      super().__init__(model, cfg_train, cfg_impl, setup, seq_length)
  
      # Optional sequence curriculum:
      self.sequence_curriculum = "sequence_curriculum" in cfg_train
      self.data_seq_length = seq_length
      self.current_seq_length = (
          seq_length
          if not self.sequence_curriculum
          else cfg_train.sequence_curriculum.lengths[0]
      )
      self.sequence_unfold = (
          None
          if not self.sequence_curriculum
          else cfg_train.sequence_curriculum.unfold
      )
  
      # Optional EMA/LAWA-type weight averages
      if "weight_averaging" in cfg_train:
          self.weight_averaging_frequency = (
              cfg_train.weight_averaging.frequency
          )
          self.weight_averaging = cfg_train.weight_averaging
          if self.weight_averaging.type == "EMA":
              self.param_store = [
                  p.detach().clone() for p in model.parameters()
              ]  # keep on CPU
              self.buffer_store = [
                  b.detach().clone() for b in model.buffers()
              ]
          else:
              self.store = []
      else:
          self.weight_averaging_frequency = 0
      self.initial_time = time.time()
  
imports: import torch
  import torch._inductor.utils
  import os
  import json
  from omegaconf import OmegaConf
  from functools import partial
  from contextlib import nullcontext
  import time
  import logging
  import transformers
  from safetensors.torch import load_file, save_file
  from transformers.utils.generic import working_or_temp_dir
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from .utils import group_parameters, prepare_pretraining_dataloader, update_ema, updated_latest_weight_average
  from .optimizers.schedulers import get_schedule_fn
  from .optimizers import Adahessian, AdamWScale, Shampoo, LARS, SAM, ProgressiveBatching, AGD, Sophia
  import warnings
label_desc: Use image node_id 1 to create a new TorchEngineFull object from inherited base classes: TorchEngineMinimal with example: obj = TorchEngineFull(model, cfg_train, cfg_impl, setup, seq_length)