Revert "Add stopwatch timer to train.py (#208)"

This reverts commit 0eee152.

Author: samsja

scripts/all_reduce.py

@@ -5,7 +5,7 @@
 
 from zeroband.collectives import Compression, all_reduce
 from zeroband.utils.world_info import get_world_info
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 
 from enum import Enum
 

scripts/convert_dl_state.py

@@ -8,7 +8,7 @@
 from zeroband.data import get_dataloader
 from transformers import AutoTokenizer
 from zeroband.train import Config
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 from pydantic_config import parse_argv
 
 COMMON_KEYS = [

scripts/export_dcp.py

@@ -11,7 +11,7 @@
 from zeroband.checkpoint import ModelWrapper
 from zeroband.utils import get_module_signature
 from zeroband.train import Config
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 from pydantic_config import parse_argv
 from transformers import AutoTokenizer
 import math

scripts/skip_data.py

@@ -25,7 +25,7 @@
 from zeroband.data import get_dataloader
 
 from zeroband.utils.world_info import get_world_info
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 
 
 def skip_data(config: Config):

src/zeroband/checkpoint.py

@@ -36,7 +36,7 @@
     send_tensor_and_state_dict,
 )
 from distributed_shampoo import DistributedShampoo
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 from zeroband.config import CkptConfig
 from zeroband.utils.world_info import get_world_info
 
@@ -151,7 +151,7 @@ def non_error_barrier():
     try:
         dist.barrier()
     except Exception as e:
-        from zeroband.utils.logger import get_logger
+        from zeroband.utils.logging import get_logger
         get_logger().info(f"Error in data checkpointing barrier: {e}, continuing training")
 
 

src/zeroband/comms.py

@@ -4,7 +4,7 @@
 import subprocess
 from torch.distributed.device_mesh import init_device_mesh
 from zeroband.utils.world_info import get_world_info
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 import torch.distributed as dist
 from datetime import timedelta
 from typing import List, Tuple, Optional

src/zeroband/data.py

@@ -3,7 +3,7 @@
 from typing import Any, Generator, Optional, List, Dict, TypedDict, Union
 import functools
 
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 from zeroband.config import DataConfig
 
 import torch
@@ -222,7 +222,7 @@ class InterleaveDataset(IterableDataset, Stateful):
     The state can be saved and restored. Under the hood we just fast forward the random generator to the current position.
     """
 
-    def __init__(self, datasets: List[ParquetDataset], probabilities: List[float], seed: int = 42):
+    def __init__(self, datasets: List[ParquetDataset], probabilities: Optional[List[float]] = None, seed: int = 42):
         assert len(datasets) > 0, "At least one dataset is required"
         assert len(datasets) == len(probabilities), "The number of datasets and probabilities must be the same"
 

src/zeroband/diloco.py

@@ -5,7 +5,7 @@
 from zeroband.comms import ElasticDeviceMesh
 from zeroband.collectives import Compression, all_reduce
 from zeroband.utils.world_info import get_world_info
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 from zeroband.config import DilocoConfig
 import torch.distributed as dist
 from torch.distributed._tensor.api import DTensor

src/zeroband/train.py

@@ -33,8 +33,7 @@
 from zeroband.utils.activation_ckpt import apply_ac_ckpt
 from zeroband.utils.profiler import MemoryProfiler
 from zeroband.utils.world_info import get_world_info
-from zeroband.utils.logger import get_logger
-from zeroband.utils.stopwatch import Stopwatch
+from zeroband.utils.logging import get_logger
 
 from transformers import AutoTokenizer
 from pydantic_config import parse_argv
@@ -94,11 +93,6 @@ def train(config: Config):
             config.ckpt.interval % config.diloco.inner_steps == 0
         ), "ckpt interval must be a multiple of diloco inner steps as we only save at the end of an outer step"
 
-    sw = Stopwatch(config)
-    sw.start("train()")
-
-    # Load tokenizer
-    sw.start_block()
     if config.data.fake and config.name_model == "debugmodel":
         tokenizer = FakeTokenizer()
     elif config.type_model == "llama2":
@@ -107,10 +101,11 @@ def train(config: Config):
         tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B", use_fast=True)
     else:
         raise ValueError(f"Model type {config.type_model} not supported")
-    sw.end_block("tokenizer loaded")
+
+    logger.debug("tokenizer loaded")
 
     with record_function("Get dataloader"):
-        sw.start_block()
+        logger.debug("Getting dataloader")
         train_dataloader = get_dataloader(
             tokenizer=tokenizer,
             world_size=world_info.world_size,
@@ -119,16 +114,13 @@ def train(config: Config):
             data_config=config.data,
         )
         train_dataloader_iterator = iter(train_dataloader)
-        sw.end_block("dataloader loaded")
 
     with record_function("Get model"):
-        sw.start_block("Constructing model")
+        logger.debug("Constructing model")
         model, model_config = get_model(
             config,
             vocab_size=len(tokenizer) if config.name_model != "debugmodel" or not config.data.fake else TEST_VOCAB_SIZE,
         )
-        sw.end_block("Constructed model")
-
 
     gpu_peak_flops = get_peak_flops(torch.cuda.get_device_name(torch.device("cuda")))
     logger.info(f"Peak FLOPS used for computing MFU: {gpu_peak_flops:.3e}")
@@ -142,7 +134,6 @@ def train(config: Config):
     )
 
     with record_function("Shard model"):
-        sw.start_block("Sharding model")
         if config.train.ac_ckpt:
             num = 1 if isinstance(config.train.ac_ckpt, bool) else config.train.ac_ckpt
             apply_ac_ckpt(model, num)
@@ -178,11 +169,10 @@ def train(config: Config):
             reshard_after_forward=config.train.reshard_after_forward,
             offload_policy=offload_policy,
         )
-        sw.end_block()
+        logger.debug("model fsdped")
 
     # Setup optimizers
     with record_function("Set up Optimizers"):
-        sw.start_block()
         inner_optimizer = get_optimizer(config, model.parameters())
 
         diloco = Diloco(config.diloco, model, elastic_device_mesh) if config.diloco is not None else None
@@ -209,7 +199,7 @@ def train(config: Config):
             diloco_offloaded_param_list=diloco.param_list_cpu if config.diloco is not None else None,  # type: ignore
         )
 
-        sw.end_block("Optimizers set up")
+        logger.debug("Optimizers set up.")
 
     if world_info.rank == 0:
         logger_cls = WandbMetricLogger if config.metric_logger_type == "wandb" else DummyMetricLogger
@@ -223,15 +213,12 @@ def train(config: Config):
 
     with record_function("Compile model"):
         if config.train.torch_compile:
-            sw.start_block()
             # we need to compile AFTER creating the CKPT manager, DON'T ASK ME WHY
             model = torch.compile(model) if not TYPE_CHECKING else model
-            sw.end_block("model compiled")
+            logger.debug("model compiled")
 
     with record_function("Resume checkpoint"):
         if config.ckpt.resume is not None:
-            sw.start_block("Resuming checkpoint")
-
             # all is inplace
             ckpt_manager.load(
                 resume_ckpt_path=config.ckpt.resume,
@@ -242,8 +229,6 @@ def train(config: Config):
                 config, model, inner_optimizer, diloco, metric_logger, step=training_progress.step, id="resume"
             )
 
-            sw.end_block("Checkpoint resumed")
-
     if config.train.memory_profiler is not None:
         memory_profiler = MemoryProfiler(config.train.memory_profiler.freq, config.train.memory_profiler.snapshot_dir)
 
@@ -254,7 +239,7 @@ def train(config: Config):
     num_inner_steps = config.diloco.inner_steps if config.diloco is not None else 1
     perf_counter = PerfCounter(window_size=10)
 
-    logger.debug("Finished setup in %f seconds", sw.elapsed())
+    logger.info("starting training")
 
     need_live_recovery = config.ckpt.live_recovery_rank_src is not None
     while True:
@@ -312,21 +297,19 @@ def train(config: Config):
 
         for inner_step in range(num_inner_steps):
             logger.debug("Starting inner step.")
-            sw.start("inner_step")
 
             loss_batch = 0
             z_loss_batch = 0
 
-            sw.start_block("Running grad acc steps")
             for grad_acc_step in range(gradient_accumulation_steps):
-                sw.start("grad_acc_step")
+                logger.debug("Starting gradient accumulation step.")
 
                 is_accumulating = grad_acc_step < gradient_accumulation_steps - 1
                 # no sync if we are accumulating gradients
                 model.set_requires_gradient_sync(not is_accumulating)
 
                 with record_function("Load batch"):
-                    sw.start_block()
+                    logger.debug("Loading batch")
                     # TODO/NOTE: We could overlap sending the batch with communication
                     #            although to be honest the perf impact is minimal
                     batch = next(train_dataloader_iterator)
@@ -338,17 +321,15 @@ def train(config: Config):
                     else:
                         seqlens = None
                         block_mask = None
-                    sw.end_block("batch loaded")
 
                 with record_function("Run model"):
-                    sw.start_block()
+                    logger.debug("Running forward()")
                     logits = model(tokens=input_ids, block_mask=block_mask).contiguous()
-                    flatten_logits = logits.reshape(-1, logits.size(-1))  # b seq vocab -> (b * seq) vocab
-                    flatten_labels = labels.reshape(-1)                   # b seq -> (b * seq)
-                    sw.end_block("Ran forward()")
+                    flatten_logits = logits.reshape(-1, logits.size(-1))  # b seq vocab -> (b seq) vocab
+                    flatten_labels = labels.reshape(-1)                   # b seq -> (b seq)
 
                 with record_function("Loss calculation"):
-                    sw.start_block()
+                    logger.debug("Computing loss")
                     ce_loss, z_loss = compute_cross_entropy_loss(
                         flatten_logits,
                         flatten_labels,
@@ -368,28 +349,22 @@ def train(config: Config):
                         loss = ce_loss + z_loss
                     else:
                         loss = ce_loss / gradient_accumulation_steps
-                    sw.end_block("Loss computed")
 
                 with record_function("Backward"):
-                    sw.start_block()
+                    logger.debug("Running backward()")
                     loss.backward()
-                    sw.end_block("Ran backward()")
 
                 with record_function("Clone loss"):
-                    # No need to time, takes 0 seconds
+                    logger.debug("Cloning loss")
                     if config.optim.z_loss:
                         assert z_loss is not None
                         loss_batch += ce_loss.detach().clone()
                         z_loss_batch += z_loss.detach().clone()
                     else:
                         loss_batch += loss.detach().clone()
 
-                elapsed = sw.stop("grad_acc_step")
-                logger.debug(f"Grad acc step {grad_acc_step} completed in {elapsed:.2f} seconds")
-            sw.end_block("Finished grad acc steps")
-
-            with record_function("Loss allreduce"):
-                sw.start_block()
+            with record_function("Inner allreduce"):
+                logger.debug("loss allreduce()")
                 # Launch both allreduces at the same time to hide latency
                 loss_allreduce = dist.all_reduce(tensor=loss_batch, op=dist.ReduceOp.AVG, group=elastic_device_mesh.local_pg, async_op=True)
                 if config.optim.z_loss:
@@ -400,22 +375,18 @@ def train(config: Config):
                 if config.optim.z_loss:
                     assert isinstance(z_loss_allreduce, torch.distributed.Work)
                     z_loss_allreduce.wait()
-                sw.end_block("loss allreduced")
 
             with record_function("Clip grad"):
-                sw.start_block()
-                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0).full_tensor() # type: ignore (is a dtensor)
-                sw.end_block("Clipped grad")
+                logger.debug("clipping grad")
+                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0).full_tensor()
+                # full tensor needed because grad_norm is a DTensor
 
             with record_function("Optimizer step"):
-                sw.start_block()
+                logger.debug("inner optimizer step()")
                 inner_optimizer.step()
                 scheduler.step()
-                sw.end_block("Inner optimizer step()")
-
-                sw.start_block()
+                logger.debug("inner optimizer zero_grad()")
                 inner_optimizer.zero_grad()
-                sw.end_block("inner optimizer zero_grad()")
 
             # logging
             training_progress.step += 1
@@ -472,9 +443,6 @@ def train(config: Config):
             if config.train.memory_profiler is not None:
                 memory_profiler.step()
 
-            elapsed = sw.stop("inner_step")
-            logger.debug(f"Inner step {inner_step} completed in {elapsed:.2f} seconds")
-
         if config.diloco is not None:
             assert diloco is not None
             if world_info.rank == 0 and config.monitor is not None:

src/zeroband/utils/activation_ckpt.py

@@ -2,7 +2,7 @@
 
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper
 
-from zeroband.utils.logger import get_logger
+from zeroband.utils.logging import get_logger
 
 
 def apply_ac_ckpt(model: Transformer, num: int):
@@ -21,4 +21,4 @@ def apply_ac_ckpt(model: Transformer, num: int):
             model.layers.register_module(layer_id, transformer_block)
             layers_ckpt += 1
 
-    logger.debug(f"Applied activation checkpointing to {layers_ckpt} layers")
+    logger.info(f"Applied activation checkpointing to {layers_ckpt} layers")