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: Run data-based initialization search as described in Zhu et al.,
  "GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training"
  
  Depends on functorch!
  
  This is gradinit without gradient aggregation, which allows higher-order derivatives
docstr_tok: ["Run","data-based","initialization","search","as","described","in","Zhu","et","al.",",","''","GradInit",":","Learning","to","Initialize","Neural","Networks","for","Stable","and","Efficient","Training","''","Depends","on","functorch","!","This","is","gradinit","without","gradient","aggregation",",","which","allows","higher-order","derivatives"]
code: def gradinit(self, data_iterable, optim_cfg, gradinit_cfg):
      """Run data-based initialization search as described in Zhu et al.,
      "GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training"
  
      Depends on functorch!
  
      This is gradinit without gradient aggregation, which allows higher-order derivatives
      """
      import functorch
  
      fmodel, params, buffers = functorch.make_functional_with_buffers(
          self.model
      )
  
      scales = [
          torch.tensor(1.0, **self.setup, requires_grad=True)
          for p in params
      ]  # Modify all params by default
      # Prepare for functional optimizer:
  
      exp_avgs = [torch.tensor(0.0, **self.setup) for s in scales]
      exp_avg_sqs = [torch.tensor(0.0, **self.setup) for s in scales]
      state_steps = [torch.tensor(0.0, **self.setup) for s in scales]
  
      adam_fn = partial(
          torch.optim._functional.adam,
          amsgrad=False,
          beta1=0.9,
          beta2=0.98,
          weight_decay=0,
          eps=1e-6,
          maximize=False,
      )
  
      eta = optim_cfg.lr
      for step in range(gradinit_cfg.steps):
          # scale params
          scaled_params = [p * s for p, s in zip(params, scales)]
          # ## Compute first step ##
          data_batch = self.to_device(next(data_iterable)[1])
          with torch.autocast(**self.amp_settings):
              loss0 = fmodel(
                  **data_batch, params=scaled_params, buffers=buffers
              )["loss"]
          grads = torch.autograd.grad(
              loss0,
              scaled_params,
              create_graph=gradinit_cfg.second_order,
              only_inputs=True,
          )
          gnorm = torch.norm(
              torch.stack([torch.norm(g) for g in grads])
          )
          # Take first step
          # p <- p - eta*g
          if gradinit_cfg.step_type == "sign-grad":
              param_step = [
                  p - eta * g.sign()
                  for p, g in zip(scaled_params, grads)
              ]
          elif gradinit_cfg.step_type == "norm-grad":
              param_step = [
                  p - eta * g / gnorm
                  for p, g in zip(scaled_params, grads)
              ]
          else:
              param_step = [
                  p - eta * g for p, g in zip(scaled_params, grads)
              ]
  
          # ## Modify scales ##
          data_batch = self.to_device(next(data_iterable)[1])
          with torch.autocast(**self.amp_settings):
              loss1 = fmodel(
                  **data_batch, params=param_step, buffers=buffers
              )["loss"]
          grads = torch.autograd.grad(
              loss1 / eta + (gnorm - 1).pow(2), scales, only_inputs=True
          )
          [
              g.zero_()
              for (name, _), g in zip(
                  self.model.named_parameters(), grads
              )
              if "pos_embedding" in name
          ]
          # Take adam step:
          with torch.no_grad():
              adam_fn(
                  scales,
                  grads,
                  exp_avgs,
                  exp_avg_sqs,
                  [],
                  state_steps,
                  lr=gradinit_cfg.tau,
              )
              # Project onto constraints and detach
              scales = [
                  s.clamp_(
                      min=gradinit_cfg.min_scale,
                      max=gradinit_cfg.max_scale,
                  )
                  for s in scales
              ]
          log.info(
              f"Gradinit: Loss0: {loss0:2.4f}. Loss1: {loss1:2.4f}. Grad Norm: {gnorm:2.4f}."
          )
          # print([f"{name}:{s.item():2.4f}" for (name, _), s in zip(self.model.named_parameters(), scales)])
  
      # Finally copy scales into the existing model
      with torch.no_grad():
          for param, scale in zip(self.model.parameters(), scales):
              param.mul_(scale)
  
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 9 for calling the TorchEngineFull obj's underlying member method code with example usage: obj.gradinit(data_iterable, optim_cfg, gradinit_cfg) without return types