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,
      already_elapsed_time=0.0,
      setup=_default_setup,
      seq_length=128,
  ):
      """Load Engine. The model will be compiled by default."""
      super().__init__()
  
      self.cfg_train = cfg_train
      self.cfg_impl = cfg_impl
      if self.cfg_impl.microbatch_size is None:
          self.cfg_impl.microbatch_size = self.cfg_train.batch_size
      if self.cfg_impl.microbatch_size > self.cfg_train.batch_size:
          raise ValueError(
              f"MBS is {self.cfg_impl.microbatch_size}, but BS is only {self.cfg_train.batch_size}."
          )
      self.current_seq_length = seq_length
  
      # Mixed Precision:
      enabled = (
          self.cfg_impl.mixed_precision
          if setup["device"].type != "cpu"
          else False
      )
      # Modules like LN are unsupported on CPU amp, so mixed precision args are disregarded on CPU
      # See https://pytorch.org/docs/stable/amp.html#cpu-op-specific-behavior and check for layer_norm
      enable_scaling = (
          self.cfg_impl.grad_scaling
          and self.cfg_impl.mixed_precision
          and setup["device"].type != "cpu"
      )
      self.scaler = torch.cuda.amp.GradScaler(enabled=enable_scaling)
      amp_dtype = (
          getattr(torch, self.cfg_impl.mixed_precision_target_dtype)
          if setup["device"].type != "cpu"
          else torch.bfloat16
      )
      self.amp_settings = dict(
          device_type=setup["device"].type,
          enabled=enabled,
          dtype=amp_dtype,
      )
  
      # Choose setup and move model
      self.setup = setup
      model.to(**self.setup)
  
      from ..utils import flatten
  
      model = torch.compile(
          model,
          mode=self.cfg_impl.mode,
          dynamic=self.cfg_impl.dynamic,
          fullgraph=self.cfg_impl.fullgraph,
          backend=self.cfg_impl.backend,
          disable=not cfg_impl.compile_torch,
          # detailed options; cannot be given at the same time as mode:
          options=flatten(
              cfg_impl._inductor_vars, parent_key="", sep="."
          )
          if cfg_impl._inductor_vars is not None
          else None,
      )
  
      if torch.distributed.is_initialized():
          self.model = self._init_distributed(model)
          self.num_machines = torch.distributed.get_world_size()
      else:
          self.model = model
          self.model.no_sync = nullcontext
          self.num_machines = 1
  
      # Microbatch accumulation settings and counters
      self.effective_mbs = (
          self.cfg_impl.microbatch_size * self.num_machines
      )  # across machines
      self.current_batch_size = (
          self.cfg_train.batch_size
          if self.cfg_train.batch_size_ramp == 0
          else self.effective_mbs
      )
      self.accumulation_steps_expected = (
          self.current_batch_size // self.effective_mbs
      )
      self.accumulated_samples = 0  # Record the number of samples seen, reset after triggering gradient update
      self.steps = (
          0  # Record the number of times "step" has been triggered
      )
  
      self.initial_time = time.time() - already_elapsed_time
      self.optimizer, self.scheduler = _load_optimizer(
          model, cfg_train, cfg_impl, self.initial_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 TorchEngineMinimal object from inherited base classes: torch.nn with example: obj = TorchEngineMinimal(model, cfg_train, cfg_impl, already_elapsed_time, setup, seq_length)