Repository: EleutherAI/gpt-neox

Issue:
Test set metrics overwrite validation set metrics in TensorBoard and are rejected for logging by Weights and Biases (W&B)
**Describe the bug**

At model train completion, the test set loss is written as iteration 0 to the TensorBoard / W&B chart `validation/lm_loss`, and the test set perplexity is written as iteration 0 to the chart `validation/lm_loss_ppl`. As the validation loss and perplexity has already been written to this chart, this results in TensorBoard deleting all the validation metrics, overwriting them with the test loss and perplexity values. W&B refuses to add the test metrics to the charts at all, throwing a warning that looks like `wandb: WARNING Step must only increase in log calls.  Step 0 < 32000; dropping {'validation/lm_loss': 1.715476632118225}.`

**To Reproduce**
Steps to reproduce the behavior:
1. Pip install and setup TensorBoard and W&B
2. Begin training a model with a train, validation, and test set
3. Observe in both TensorBoard and W&B that validation metrics are being logged
4. Allow the model to train to completion
5. Observe that the TensorBoard validation metrics are now gone, overwritten by the test set metrics
6. Observe the W&B error in the text logs / program output

**Expected behavior**
Test metrics should be written to their own charts.

**Proposed solution**
Test loss and perplexity should be written to their own charts `test/lm_loss` and `test/lm_loss_ppl` respectively.

**Screenshots**
![image](https://user-images.githubusercontent.com/6119143/189752970-3b26dd14-475f-48cb-be84-fae23a99ba10.png)

**Environment (please complete the following information):**
 - GPUs: 4x A100 80 GB
- Configs: (configs that I used to reproduce the bug and test bug fixes are included below)

```
# GPT-2 pretraining setup
{
   # parallelism settings ( you will want to change these based on your cluster setup, ideally scheduling pipeline stages
   # across the node boundaries )
   "pipe-parallel-size": 1,
   "model-parallel-size": 1,

   # model settings
   "num-layers": 24,
   "hidden-size": 1024,
   "num-attention-heads": 16,
   "seq-length": 4096,
   "max-position-embeddings": 4096,
   "norm": "layernorm",
   "pos-emb": "rotary",
   "no-weight-tying": true,

   # these should provide some speedup but takes a while to build, set to true if desired
   "scaled-upper-triang-masked-softmax-fusion": false,
   "bias-gelu-fusion": false,



   # optimizer settings
   "optimizer": {
     "type": "Adam",
     "params": {
       "lr": 0.00003,
       "betas": [0.9, 0.999],
       "eps": 1.0e-8,
     }
   },
   "zero_optimization": {
    "stage": 1,
    "allgather_partitions": True,
    "allgather_bucket_size": 500000000,
    "overlap_comm": True,
    "reduce_scatter": True,
    "reduce_bucket_size": 500000000,
    "contiguous_gradients": True,
    "cpu_offload": False
  },
   # batch / data settings
   "train_micro_batch_size_per_gpu": 16,
   "data-impl": "mmap",
   "split": "949,50,1",

   # activation checkpointing
   "checkpoint-activations": true,
   "checkpoint-num-layers": 1,
   "partition-activations": true,
   "synchronize-each-layer": true,

   # regularization
   "gradient_clipping": 1.0,
   "weight-decay": 0.01,
   "hidden-dropout": 0,
   "attention-dropout": 0,

   # precision settings
   "fp16": {
     "fp16": true,
     "enabled": true,
     "loss_scale": 0,
     "loss_scale_window": 1000,
     "hysteresis": 2,
     "min_loss_scale": 1
   },

   # misc. training settings
   "train-iters": 100,
   "lr-decay-iters": 100,
   "distributed-backend": "nccl",
   "lr-decay-style": "constant",
   "warmup": 0.1,
   "save-interval": 25,
   "eval-interval": 25,
   "eval-iters": 10,

   # Checkpoint
   "finetune": true,

   # logging
   "log-interval": 10,
   "steps_per_print": 10,
   "keep-last-n-checkpoints": 4,
   "wall_clock_breakdown": true,
}
```

```
# Suggested data paths when using GPT-NeoX locally
{
  "train-data-paths": ["/mnt/4TBNVME/gpt-neox/data/preprocessed/train_text_document"],
  "test-data-paths": ["/mnt/4TBNVME/gpt-neox/data/preprocessed/test_text_document"],
  "valid-data-paths": ["/mnt/4TBNVME/gpt-neox/data/preprocessed/val_text_document"],

  "vocab-file": "/mnt/4TBNVME/gpt-neox/data/gpt2-vocab.json",
  "merge-file": "/mnt/4TBNVME/gpt-neox/data/gpt2-merges.txt",

  "save": "/mnt/4TBNVME/checkpoints_test",
  "load": "/mnt/4TBNVME/checkpoints_test",

  "checkpoint_validation_with_forward_pass": False,
  
  "tensorboard-dir": "/mnt/4TBNVME/logs/tensorboard/bug_fix_test",
  "log-dir": "/mnt/4TBNVME/logs/gptneox/bug_fix_test",

  "use_wandb": True,
  "wandb_host": "https://api.wandb.ai",
  "wandb_project": "neox_test"
}
```

```
# Add this to your config for sparse attention every other layer
{
  "attention_config": [[["local", "global"], "all"]],

  # sparsity config:
  # (these are the defaults for local sliding window sparsity, training will work without this here, but it's left in for
  # illustrative purposes)
  # see https://www.deepspeed.ai/tutorials/sparse-attention/#how-to-config-sparsity-structures for
  # more detailed config instructions and available parameters

  "sparsity_config": {
    "block": 16, # block size
    "num_local_blocks": 32,
  }
}
```

**Additional context**

I have a bug fix ready, will follow up with it.


Discussion:

Resolving patch:
diff --git a/megatron/training.py b/megatron/training.py
--- a/megatron/training.py
+++ b/megatron/training.py
@@ -144,9 +144,10 @@ def pretrain(neox_args):
             forward_step_func=forward_step,
             data_iterator=test_data_iterator,
             model=model,
-            iteration=0,  # iteration 0 in order to always use full test data
+            iteration=iteration,
             verbose=True,
             timers=timers,
+            chart_name="test"
         )
 
 
@@ -736,6 +737,7 @@ def evaluate_and_print_results(
     iteration,
     verbose=False,
     timers=None,
+    chart_name="validation"
 ):
     """Helper function to evaluate and dump results on screen."""
     total_loss_dict = evaluate(
@@ -746,14 +748,14 @@ def evaluate_and_print_results(
         verbose=verbose,
         timers=timers,
     )
-    string = f" validation results at {prefix} | "
+    string = f" {chart_name} results at {prefix} | "
     for k, v in total_loss_dict.items():
         if isinstance(v, dict):
             for k2, v2 in v.items():
                 k3 = "_".join([k, k2])
                 string += f"{k3} value: {v2:.6E} | "
                 tb_wandb_log(
-                    f"validation/{k3}",
+                    f"{chart_name}/{k3}",
                     v2,
                     iteration,
                     use_wandb=neox_args.use_wandb,
@@ -762,7 +764,7 @@ def evaluate_and_print_results(
         else:
             string += f"{k} value: {v:.6E} | "
             tb_wandb_log(
-                f"validation/{k}",
+                f"{chart_name}/{k}",
                 v,
                 iteration,
                 use_wandb=neox_args.use_wandb,


Test patch: