Return tqdm to benchmark_dht

Author: mryab

benchmarks/benchmark_dht.py

@@ -2,6 +2,8 @@
 import random
 import time
 
+from tqdm import trange
+
 import hivemind
 from hivemind.moe.server import declare_experts, get_experts
 from hivemind.utils.limits import increase_file_limit
@@ -31,7 +33,7 @@ def benchmark_dht(
 
     logger.info("Creating peers...")
     peers = []
-    for _ in range(num_peers):
+    for _ in trange(num_peers):
         neighbors = [f"0.0.0.0:{node.port}" for node in random.sample(peers, min(initial_peers, len(peers)))]
         peer = hivemind.DHT(initial_peers=neighbors, start=True, wait_timeout=wait_timeout, listen_on=f"0.0.0.0:*")
         peers.append(peer)
@@ -52,7 +54,7 @@ def benchmark_dht(
     benchmark_started = time.perf_counter()
     endpoints = []
 
-    for start in range(0, num_experts, expert_batch_size):
+    for start in trange(0, num_experts, expert_batch_size):
         store_start = time.perf_counter()
         endpoints.append(random_endpoint())
         store_ok = declare_experts(
@@ -76,7 +78,7 @@ def benchmark_dht(
 
     successful_gets = total_get_time = 0
 
-    for start in range(0, len(expert_uids), expert_batch_size):
+    for start in trange(0, len(expert_uids), expert_batch_size):
         get_start = time.perf_counter()
         get_result = get_experts(get_peer, expert_uids[start : start + expert_batch_size])
         total_get_time += time.perf_counter() - get_start

examples/albert/arguments.py

@@ -109,7 +109,7 @@ class AlbertTrainingArguments(TrainingArguments):
     gradient_accumulation_steps: int = 2
     seq_length: int = 512
 
-    max_steps: int = 125_000 # please note: this affects both number of steps and learning rate schedule
+    max_steps: int = 125_000  # please note: this affects both number of steps and learning rate schedule
     learning_rate: float = 0.00176
     warmup_steps: int = 5000
     adam_epsilon: float = 1e-6

examples/albert/run_trainer.py

@@ -10,8 +10,15 @@
 import transformers
 from datasets import load_from_disk
 from torch.utils.data import DataLoader
-from transformers import (set_seed, HfArgumentParser, TrainingArguments,
-                          DataCollatorForLanguageModeling, AlbertTokenizerFast, AlbertConfig, AlbertForPreTraining)
+from transformers import (
+    set_seed,
+    HfArgumentParser,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+    AlbertTokenizerFast,
+    AlbertConfig,
+    AlbertForPreTraining,
+)
 from transformers.optimization import get_linear_schedule_with_warmup
 from transformers.trainer_utils import is_main_process
 from transformers.trainer import Trainer
@@ -23,7 +30,7 @@
 
 
 logger = logging.getLogger(__name__)
-LRSchedulerBase = getattr(torch.optim.lr_scheduler, '_LRScheduler', None)
+LRSchedulerBase = getattr(torch.optim.lr_scheduler, "_LRScheduler", None)
 
 
 def setup_logging(training_args):
@@ -50,13 +57,13 @@ def get_model(training_args, config, tokenizer):
     # Find latest checkpoint in output_dir
     output_dir = Path(training_args.output_dir)
     logger.info(f'Checkpoint dir {output_dir}, contents {list(output_dir.glob("checkpoint*"))}')
-    latest_checkpoint_dir = max(output_dir.glob('checkpoint*'), default=None, key=os.path.getctime)
+    latest_checkpoint_dir = max(output_dir.glob("checkpoint*"), default=None, key=os.path.getctime)
 
     if latest_checkpoint_dir is not None:
-        logger.info(f'Loading model from {latest_checkpoint_dir}')
+        logger.info(f"Loading model from {latest_checkpoint_dir}")
         model = AlbertForPreTraining.from_pretrained(latest_checkpoint_dir)
     else:
-        logger.info(f'Training from scratch')
+        logger.info(f"Training from scratch")
         model = AlbertForPreTraining(config)
         model.resize_token_embeddings(len(tokenizer))
 
@@ -87,17 +94,21 @@ def get_optimizer_and_scheduler(training_args, model):
     )
 
     scheduler = get_linear_schedule_with_warmup(
-        opt,
-        num_warmup_steps=training_args.warmup_steps,
-        num_training_steps=training_args.max_steps
+        opt, num_warmup_steps=training_args.warmup_steps, num_training_steps=training_args.max_steps
     )
 
     return opt, scheduler
 
 
 class CollaborativeCallback(transformers.TrainerCallback):
-    def __init__(self, dht: hivemind.DHT, optimizer: hivemind.CollaborativeOptimizer,
-                 model: torch.nn.Module, local_public_key: bytes, statistics_expiration: float):
+    def __init__(
+        self,
+        dht: hivemind.DHT,
+        optimizer: hivemind.CollaborativeOptimizer,
+        model: torch.nn.Module,
+        local_public_key: bytes,
+        statistics_expiration: float,
+    ):
         super().__init__()
         self.model = model
         self.dht, self.collaborative_optimizer = dht, optimizer
@@ -110,21 +121,23 @@ def __init__(self, dht: hivemind.DHT, optimizer: hivemind.CollaborativeOptimizer
         self.loss = 0
         self.total_samples_processed = 0
 
-    def on_train_begin(self, args: TrainingArguments, state: transformers.TrainerState,
-                       control: transformers.TrainerControl, **kwargs):
-        logger.info('Loading state from peers')
+    def on_train_begin(
+        self, args: TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs
+    ):
+        logger.info("Loading state from peers")
         self.collaborative_optimizer.load_state_from_peers()
 
-    def on_step_end(self, args: TrainingArguments, state: transformers.TrainerState,
-                    control: transformers.TrainerControl, **kwargs):
+    def on_step_end(
+        self, args: TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs
+    ):
         control.should_log = True
         if not self.params_are_finite():
             self.load_from_state(self.previous_state)
             return control
         self.previous_state = self.get_current_state()
 
         if state.log_history:
-            self.loss += state.log_history[-1]['loss']
+            self.loss += state.log_history[-1]["loss"]
             self.steps += 1
             if self.collaborative_optimizer.local_step != self.last_reported_collaboration_step:
                 self.last_reported_collaboration_step = self.collaborative_optimizer.local_step
@@ -135,7 +148,8 @@ def on_step_end(self, args: TrainingArguments, state: transformers.TrainerState,
                     samples_per_second=samples_per_second,
                     samples_accumulated=self.samples,
                     loss=self.loss,
-                    mini_steps=self.steps)
+                    mini_steps=self.steps,
+                )
                 logger.info(f"Step {self.collaborative_optimizer.local_step}")
                 logger.info(f"Your current contribution: {self.total_samples_processed} samples")
                 if self.steps:
@@ -144,26 +158,26 @@ def on_step_end(self, args: TrainingArguments, state: transformers.TrainerState,
                 self.loss = 0
                 self.steps = 0
                 if self.collaborative_optimizer.is_synchronized:
-                    self.dht.store(key=self.collaborative_optimizer.prefix + "_metrics",
-                                   subkey=self.local_public_key, value=statistics.dict(),
-                                   expiration_time=hivemind.get_dht_time() + self.statistics_expiration,
-                                   return_future=True)
+                    self.dht.store(
+                        key=self.collaborative_optimizer.prefix + "_metrics",
+                        subkey=self.local_public_key,
+                        value=statistics.dict(),
+                        expiration_time=hivemind.get_dht_time() + self.statistics_expiration,
+                        return_future=True,
+                    )
 
         self.samples = self.collaborative_optimizer.local_samples_accumulated
 
         return control
 
     @torch.no_grad()
     def get_current_state(self) -> Dict[str, Any]:
-        return {
-            'model': self.model.state_dict(),
-            'opt': self.collaborative_optimizer.opt.state_dict()
-        }
+        return {"model": self.model.state_dict(), "opt": self.collaborative_optimizer.opt.state_dict()}
 
     @torch.no_grad()
     def load_from_state(self, state):
-        self.model.load_state_dict(state['model'])
-        self.collaborative_optimizer.opt.load_state_dict(state['opt'])
+        self.model.load_state_dict(state["model"])
+        self.collaborative_optimizer.opt.load_state_dict(state["opt"])
 
     @torch.no_grad()
     def params_are_finite(self):
@@ -174,10 +188,10 @@ def params_are_finite(self):
 
 
 class NoOpScheduler(LRSchedulerBase):
-    """ Dummy scheduler for transformers.Trainer. The real scheduler is defined in CollaborativeOptimizer.scheduler """
+    """Dummy scheduler for transformers.Trainer. The real scheduler is defined in CollaborativeOptimizer.scheduler"""
 
     def get_lr(self):
-        return [group['lr'] for group in self.optimizer.param_groups]
+        return [group["lr"] for group in self.optimizer.param_groups]
 
     def print_lr(self, *args, **kwargs):
         if self.optimizer.scheduler:
@@ -219,53 +233,65 @@ def main():
 
     opt, scheduler = get_optimizer_and_scheduler(training_args, model)
 
-    validators, local_public_key = metrics_utils.make_validators(
-        collaboration_args_dict['experiment_prefix'])
+    validators, local_public_key = metrics_utils.make_validators(collaboration_args_dict["experiment_prefix"])
     dht = hivemind.DHT(
-        start=True, initial_peers=collaboration_args_dict.pop('initial_peers'),
-        listen=not collaboration_args_dict['client_mode'],
-        listen_on=collaboration_args_dict.pop('dht_listen_on'),
-        endpoint=collaboration_args_dict.pop('endpoint'), record_validators=validators)
+        start=True,
+        initial_peers=collaboration_args_dict.pop("initial_peers"),
+        listen=not collaboration_args_dict["client_mode"],
+        listen_on=collaboration_args_dict.pop("dht_listen_on"),
+        endpoint=collaboration_args_dict.pop("endpoint"),
+        record_validators=validators,
+    )
 
     total_batch_size_per_step = training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps
     if torch.cuda.device_count() != 0:
         total_batch_size_per_step *= torch.cuda.device_count()
 
-    statistics_expiration = collaboration_args_dict.pop('statistics_expiration')
-    adjusted_target_batch_size = collaboration_args_dict.pop('target_batch_size') \
-                                 - collaboration_args_dict.pop('batch_size_lead')
+    statistics_expiration = collaboration_args_dict.pop("statistics_expiration")
+    adjusted_target_batch_size = collaboration_args_dict.pop("target_batch_size") - collaboration_args_dict.pop(
+        "batch_size_lead"
+    )
 
     collaborative_optimizer = hivemind.CollaborativeOptimizer(
-        opt=opt, dht=dht, scheduler=scheduler, prefix=collaboration_args_dict.pop('experiment_prefix'),
-        compression_type=hivemind.utils.CompressionType.Value(collaboration_args_dict.pop('compression')),
-        batch_size_per_step=total_batch_size_per_step, throughput=collaboration_args_dict.pop('bandwidth'),
-        target_batch_size=adjusted_target_batch_size, client_mode=collaboration_args_dict.pop('client_mode'),
-        verbose=True, start=True, **collaboration_args_dict
+        opt=opt,
+        dht=dht,
+        scheduler=scheduler,
+        prefix=collaboration_args_dict.pop("experiment_prefix"),
+        compression_type=hivemind.utils.CompressionType.Value(collaboration_args_dict.pop("compression")),
+        batch_size_per_step=total_batch_size_per_step,
+        throughput=collaboration_args_dict.pop("bandwidth"),
+        target_batch_size=adjusted_target_batch_size,
+        client_mode=collaboration_args_dict.pop("client_mode"),
+        verbose=True,
+        start=True,
+        **collaboration_args_dict,
     )
 
     class TrainerWithIndependentShuffling(Trainer):
         def get_train_dataloader(self) -> DataLoader:
-            """ Shuffle data independently for each peer to avoid duplicating batches [important for quality] """
+            """Shuffle data independently for each peer to avoid duplicating batches [important for quality]"""
             torch.manual_seed(hash(local_public_key))
             return super().get_train_dataloader()
 
     trainer = TrainerWithIndependentShuffling(
-        model=model, args=training_args, tokenizer=tokenizer, data_collator=data_collator,
+        model=model,
+        args=training_args,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
         train_dataset=tokenized_datasets["train"] if training_args.do_train else None,
         eval_dataset=tokenized_datasets["validation"] if training_args.do_eval else None,
         optimizers=(collaborative_optimizer, NoOpScheduler(collaborative_optimizer)),
-        callbacks=[CollaborativeCallback(
-            dht, collaborative_optimizer, model, local_public_key, statistics_expiration)]
+        callbacks=[
+            CollaborativeCallback(dht, collaborative_optimizer, model, local_public_key, statistics_expiration)
+        ],
     )
     trainer.remove_callback(transformers.trainer_callback.PrinterCallback)
     trainer.remove_callback(transformers.trainer_callback.ProgressCallback)
 
     # Training
     if training_args.do_train:
         latest_checkpoint_dir = max(
-            Path(training_args.output_dir).glob('checkpoint*'),
-            default=None,
-            key=os.path.getctime
+            Path(training_args.output_dir).glob("checkpoint*"), default=None, key=os.path.getctime
         )
 
         trainer.train(model_path=latest_checkpoint_dir)