add grad clipping, use wandb run name for save dir

configs/pile_megatron_dataset.yaml

@@ -3,7 +3,6 @@
   # because we load it from yaml and then feed to NeoXArgs.from_dict().
   # Use _ instead of - in the key names
 
-  "global_num_gpus": 8,
   "pipe_parallel_size": 1,
   "model_parallel_size": 1,
 

peft_pretraining/args_utils.py

@@ -35,15 +35,6 @@ def check_args_torchrun_main(args):
         logger.error("Are you sure? Not training LN is a bad idea.")
         raise ValueError("Are you sure? Not training LN is a bad idea.")
 
-    if args.save_dir is None:
-        if args.model_config is not None:
-            # use checkpoints / model name, date and time as save directory
-            args.save_dir = f"checkpoints/{args.model_config.split('/')[-1].rstrip('.json')}-{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}"
-        elif args.model_name_or_path is not None:
-            args.save_dir = f"checkpoints/{args.model_name_or_path.split('/')[-1]}-{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}"
-        else:
-            raise ValueError("Either --args.save_dir or --model_config or --model_name_or_path must be specified")
-
     if args.tags is not None:
         args.tags = args.tags.split(",")
 

peft_pretraining/training_utils.py

@@ -380,14 +380,12 @@ def print_optimizer_state_size(optimizer):
         first_moment_count += torch.numel(state['exp_avg'])
         second_moment_count += torch.numel(state['exp_avg_sq'])
 
-    logger.info(f'Number of floats in the first moment: {first_moment_count / 1_000_000:.2f}M')
-    logger.info(f'Number of floats in the second moment: {second_moment_count / 1_000_000:.2f}M')
     global_rank = 0
     if dist.is_initialized():
         global_rank = dist.get_rank()
-    if 0 < global_rank < 8:
-        print(f"(Rank {global_rank}) Number of floats in the first moment: {first_moment_count / 1_000_000:.2f}M")
-        print(f"(Rank {global_rank}) Number of floats in the second moment: {second_moment_count / 1_000_000:.2f}M")
+
+    print(f"(Rank {global_rank}) Number of floats in the first moment: {first_moment_count / 1_000_000:.2f}M")
+    print(f"(Rank {global_rank}) Number of floats in the second moment: {second_moment_count / 1_000_000:.2f}M")
 
 
 def check_lr_and_alert(optimizer, max_lr):