code: def calculate_total_num_steps(cfg, train_dataset, update=True):
      if not cfg.total_num_tokens:
          total_num_tokens = np.sum(
              train_dataset.data.column("input_ids")
              .to_pandas()
              .apply(
                  lambda x: len(x)
              )  # pylint: disable=unnecessary-lambda
              .values
          )
          LOG.debug(
              f"total_num_tokens: {total_num_tokens}",
              main_process_only=True,
          )
          if update:
              cfg.total_num_tokens = total_num_tokens
  
      skip_estimates = cfg.model_config_type == "mamba"
  
      if not skip_estimates and not cfg.total_supervised_tokens:
          total_supervised_tokens = (
              train_dataset.data.column("labels")
              .to_pandas()
              .apply(lambda x: np.sum(np.array(x) != -100))
              .sum()
          )
          LOG.debug(
              f"`total_supervised_tokens: {total_supervised_tokens}`",
              main_process_only=True,
          )
          if update:
              cfg.total_supervised_tokens = total_supervised_tokens
  
      if not skip_estimates and cfg.sample_packing:
          # we have to drop anything longer then sequence len otherwise
          # flash attention with position ids fails
  
          if cfg.sample_packing_eff_est:
              total_num_steps = (
                  # match count to len est in dataloader
                  (
                      math.floor(
                          0.99
                          * cfg.total_num_tokens
                          / cfg.sample_packing_eff_est
                          / cfg.sequence_len
                          // cfg.batch_size
                          // int(os.environ.get("WORLD_SIZE", 1))
                      )
                      - 1
                  )
                  * cfg.num_epochs
              )
              LOG.debug(
                  f"total_num_tokens: {cfg.total_num_tokens}, total_num_steps: {total_num_steps}",
                  main_process_only=True,
              )
          else:
              sampler = MultipackBatchSampler(
                  sampler=RandomSampler(train_dataset),
                  batch_size=cfg.micro_batch_size,
                  drop_last=True,
                  batch_max_len=cfg.micro_batch_size
                  * (cfg.max_packed_sequence_len or cfg.sequence_len),
                  lengths=get_dataset_lengths(train_dataset),
              )
  
              data_loader = DataLoader(
                  train_dataset.remove_columns(["length"]),
                  batch_sampler=sampler,
              )
              data_loader_len = len(data_loader)
              actual_eff = sampler.efficiency()
              LOG.debug(
                  f"data_loader_len: {data_loader_len}",
                  main_process_only=True,
              )
              # FIXME: is there a bug here somewhere? the total num steps depends
              # on the agreed on value for sample_packing_eff_est
              total_num_steps = int(
                  math.floor(
                      data_loader_len
                      * cfg.num_epochs
                      / int(os.environ.get("WORLD_SIZE", 1))
                  )
              )
  
              def calc_sample_packing_eff_est(estimates: List[float]):
                  LOG.info(
                      f"sample_packing_eff_est across ranks: {repr(estimates)}"
                  )
                  return max(estimates)
  
              sample_packing_actual_eff_all = reduce_and_broadcast(
                  lambda: actual_eff,
                  calc_sample_packing_eff_est,
              )
              sample_packing_eff_est = (
                  math.ceil(sample_packing_actual_eff_all * 100.0)
                  / 100.0
              )
              if update:
                  cfg.sample_packing_eff_est = sample_packing_eff_est
              LOG.debug(
                  f"sample_packing_eff_est: {cfg.sample_packing_eff_est}",
                  main_process_only=True,
              )
      else:
          total_num_steps = int(
              math.ceil(
                  len(train_dataset)
                  * cfg.num_epochs
                  / int(os.environ.get("WORLD_SIZE", 1))
                  / cfg.batch_size
              )
          )
      LOG.debug(
          f"total_num_steps: {total_num_steps}", main_process_only=True
      )
      return total_num_steps
  
imports: import math
  import os
  from contextlib import contextmanager
  from functools import partial
  from typing import List
  import numpy
  import torch
  import torch.cuda
  from accelerate.logging import get_logger
  from datasets import set_caching_enabled
  from torch.utils.data import DataLoader, RandomSampler
  from axolotl.core.trainer_builder import HFCausalTrainerBuilder, HFDPOTrainerBuilder
  from axolotl.utils.distributed import is_main_process, reduce_and_broadcast, zero_first
  from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
label_desc: Use image node_id 10 for calling a global function with example usage: calculate_total_num_steps(cfg, train_dataset, update) and returns: total_num_steps, max