cleanup verbosity a bit

gitbook/README.md

@@ -1,2 +1 @@
 # Page
-

src/axolotl/train.py

@@ -1,6 +1,5 @@
 """Prepare and train a model on a dataset. Can also infer from a model or merge lora"""
 
-import logging
 import os
 import signal
 import sys
@@ -10,6 +9,7 @@
 
 import torch
 import transformers.modelcard
+from accelerate.logging import get_logger
 from datasets import Dataset
 from optimum.bettertransformer import BetterTransformer
 from transformers.deepspeed import is_deepspeed_zero3_enabled
@@ -25,7 +25,7 @@
 sys.path.insert(0, src_dir)
 
 configure_logging()
-LOG = logging.getLogger("axolotl.train")
+LOG = get_logger("axolotl.train")
 
 
 @dataclass
@@ -43,15 +43,21 @@ def train(
  *, cfg: DictDefault, cli_args: TrainerCliArgs, dataset_meta: TrainDatasetMeta
 ):
  # load the tokenizer first
- LOG.info(f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}")
+ LOG.debug(
+ f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}",
+ main_process_only=True,
+ )
  tokenizer = load_tokenizer(cfg)
 
  train_dataset = dataset_meta.train_dataset
  eval_dataset = dataset_meta.eval_dataset
  total_num_steps = dataset_meta.total_num_steps
 
  # Load the model and tokenizer
- LOG.info("loading model and (optionally) peft_config...")
+ msg = "loading model"
+ if cfg.adapter:
+ msg += " and peft_config..."
+ LOG.debug(msg)
  model, peft_config = load_model(cfg, tokenizer, inference=cli_args.inference)
 
  safe_serialization = cfg.save_safetensors is True

src/axolotl/utils/distributed.py

@@ -50,6 +50,17 @@ def get_world_size():
  return int(os.getenv("WORLD_SIZE", "1"))
 
 
+@contextmanager
+def zero_only():
+ """
+ Context manager that only runs the enclosed block on the main rank.
+ """
+ if is_main_process():
+ yield
+ else:
+ yield None
+
+
 @contextmanager
 def zero_first(is_main):
  """

src/axolotl/utils/trainer.py

@@ -1,5 +1,4 @@
 """Module containing the Trainer class and related functions"""
-import logging
 import math
 import os
 from contextlib import contextmanager
@@ -10,6 +9,7 @@
 import torch
 import torch.cuda
 import torch.distributed as dist
+from accelerate.logging import get_logger
 from datasets import set_caching_enabled
 from torch.utils.data import DistributedSampler, RandomSampler
 
@@ -23,7 +23,7 @@
  zero_first,
 )
 
-LOG = logging.getLogger("axolotl")
+LOG = get_logger("axolotl")
 
 
 @torch.jit.script
@@ -153,14 +153,13 @@ def calculate_total_num_steps(cfg, train_dataset, tokenizer):
  # we have to drop anything longer then sequence len otherwise
  # flash attention with position ids fails
  if not cfg.total_num_tokens:
- LOG.info("calculating 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.info(f"total_num_tokens: {total_num_tokens}")
+ LOG.debug(f"total_num_tokens: {total_num_tokens}", main_process_only=True)
  cfg.total_num_tokens = total_num_tokens
 
  if not cfg.total_supervised_tokens:
@@ -170,7 +169,10 @@ def calculate_total_num_steps(cfg, train_dataset, tokenizer):
  .apply(lambda x: np.sum(np.array(x) != -100))
  .sum()
  )
- LOG.info(f"`total_supervised_tokens: {total_supervised_tokens}`")
+ LOG.debug(
+ f"`total_supervised_tokens: {total_supervised_tokens}`",
+ main_process_only=True,
+ )
  cfg.total_supervised_tokens = total_supervised_tokens
 
  if cfg.sample_packing_eff_est:
@@ -189,8 +191,9 @@ def calculate_total_num_steps(cfg, train_dataset, tokenizer):
  )
  * cfg.num_epochs
  )
- LOG.info(
- f"total_num_tokens: {cfg.total_num_tokens}, total_num_steps: {total_num_steps}"
+ LOG.debug(
+ f"total_num_tokens: {cfg.total_num_tokens}, total_num_steps: {total_num_steps}",
+ main_process_only=True,
  )
  else:
  if cfg.world_size > 1 and is_distributed():
@@ -220,7 +223,7 @@ def calculate_total_num_steps(cfg, train_dataset, tokenizer):
  )
  data_loader_len = data_loader.len_w_stats()
  actual_eff = data_loader.efficiency()
- LOG.info(f"data_loader_len: {data_loader_len}")
+ 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))
@@ -237,12 +240,15 @@ def calc_sample_packing_eff_est(estimates: List[float]):
  math.ceil(sample_packing_actual_eff_all * 100.0) / 100.0
  )
  cfg.sample_packing_eff_est = sample_packing_eff_est
- LOG.info(f"sample_packing_eff_est: {cfg.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 / cfg.batch_size)
  )
- LOG.info(f"total_num_steps: {total_num_steps}")
+ LOG.debug(f"total_num_steps: {total_num_steps}", main_process_only=True)
  return total_num_steps