Question about fp16 dynamic loss scale overflow in DeepSpeed

Hi, I have some issues when using DeepSpeed for mixed precision training (fp16 with dynamic scaling). I would genuinely appreciate it if you could please give me some help on this issue.

Currently, I am using DeepSpeed for mixed precision training (fp16 with dynamic scaling) to reproduce CodeBERT. When setting lr=1e-4 (a relatively small learning rate), everything goes on well: the training loss curve went smoothly, and there only exists very few fp16 dynamic loss scale overflow. However, When setting lr=5e-4 (a relatively large learning rate, following the original paper's setting), after training for 10k steps, the training loss rises dramatically and the validation accuracy after the huge increase drops to a very low level (even 0). Before the increase happens, the dynamic loss scale gradually rises to 32768.0; after the huge increase, the dynamic loss scale is reduced to 1024.0 step by step, but the training loss is still abnormal.

This is my DeepSpeed configurations:

    ds_config = {
        "train_batch_size": args.batch_size,
        "train_micro_batch_size_per_gpu": args.micro_batch_size,
        "steps_per_print": 10000,
        "gradient_clipping": 1.0,
        "wall_clock_breakdown": False,
        "fp16": {
            "enabled": args.fp16,
            "loss_scale": 0.0,
            "loss_scale_window": 1000,
            "hysteresis": 2,
            "min_loss_scale": 1
        },
        "local_rank": args.local_rank,
        "zero_optimization": {
            "stage3_gather_16bit_weights_on_model_save": True
        }
    }

This is the picture that shows the training loss and the learning rate before and after the huge increase:

I've been working on this problem for a while, but I failed to solve it. Could you please give me some advice on this issue? (i.e. what could be the cause of this problem? will it help to enlarge loss_scale_window, limit the dynamic loss scale to a small number, or change other DeepSpeed configurations?)

No source file or test is identified in the issue. Start by reviewing the DeepSpeed fp16 configuration and reproducing the two learning-rate runs around the 10k-step loss increase, including the dynamic loss-scale values. Done means identifying whether the behavior is caused by the configuration or training setup and documenting actionable guidance.