using deepspeed initialize huggingface GPT2LMHeadModel, it occured RuntimeError: "Function 'LogSoftmaxBackward' returned nan values in its 0th output."

i try to use DeepSpeed to initialize huggingface GPT2LMHeadModel, but it raised RuntimeError. I can't solve it, please help me fix this error . Any information you need, please let me add.

the error log

error 1 and error 2 only change torch.autograd.set_detect_anomaly option.

error 1 (with torch.autograd.detect_anomaly)

torch.autograd.set_detect_anomaly(True)
torch.backends.cudnn.enabled = False

error 2 (without torch.autograd.detect_anomaly)

torch.autograd.set_detect_anomaly(False)
torch.backends.cudnn.enabled = False

Version and env

env

NVIDIA-SMI 440.33.01 Driver Version: 440.33.01 CUDA Version: 10.2 Ubuntu 16.04 pytorch 1.5.1

Deepspeed version and some info

[2020-11-24 11:07:39,793] [INFO] [logging.py:60:log_dist] [Rank 0] DeepSpeed info: version=0.3.0+4c37d70, git-hash=4c37d70, git-branch=master
[2020-11-24 11:07:44,706] [INFO] [engine.py:553:_configure_optimizer] DeepSpeed Final Optimizer = <deepspeed.runtime.zero.stage2.FP16_DeepSpeedZeroOptimizer object at 0x7f29e699aef0>
[2020-11-24 11:07:44,718] [INFO] [engine.py:393:_configure_lr_scheduler] DeepSpeed using client LR scheduler
[2020-11-24 11:07:44,718] [INFO] [logging.py:60:log_dist] [Rank 0] DeepSpeed LR Scheduler = <learning_rates.AnnealingLR object at 0x7f29e7676748>

ds_config

{
"train_batch_size":1,
"train_micro_batch_size_per_gpu": 1,
"steps_per_print": 1,
"zero_allow_untested_optimizer": true,
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.015,
"weight_decay": 1e-2
}
},
"zero_optimization": {
"stage": 2,
"cpu_offload": true,
"reduce_bucket_size": 50000000
},
"gradient_clipping": 1.0,
"tensorboard": {
"enabled": true,
"output_path": "path",
"job_name": "job"
},
"fp16": {
"enabled": true,
"init_scale": 1048576,
"loss_scale": 0,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},
"activation_checkpointing": {
"partition_activations": true,
"contiguous_memory_optimization": true,
"cpu_checkpointing": true,
"profile": false
},
"wall_clock_breakdown": true
}

transformers version and some info

transformers 3.1.0
use transformers/modeling_gpt2.py: GPT2LMHeadModel as model
use transformers/trainer.py: Trainer as train script
use examples/language-modeling/run_language_modeling.py as main script

model_config

{
  "activation_function": "gelu_new",
  "architectures": [
  "GPT2LMHeadModel"
  ],
"attn_pdrop": 0.1,
 "bos_token_id": 50256,
 "embd_pdrop": 0.1,
 "eos_token_id": 50256,
 "initializer_range": 0.02,
 "layer_norm_epsilon": 1e-03,
 "model_type": " gpt2",
 "n_ctx": 1024,
 "n_embd": 768,
 "n_head": 12,
  "n_layer": 2,
  "n_positions": 1024,
  "resid_pdrop": 0.1,
  "summary_activation": null,
  "summary_first_dropout": 0.1,
  "summary_proj_to_labels": true,
   "summary_type": "cls_index",
  "summary_use_proj": true,
  "task_specific_params": {
     "text-generation": {
        "do_sample": true,
         "max_length": 50
       }
    },
   "vocab_size": 50257
  }

The code

1. add DeepSpeedArguments

2. In run_language_modeling.py:main add initialize_distributed(deepspeed_args)

3. modify Trainer init:

    optimizer, param_groups = get_optimizer(model, deepspeed_args)
    lr_scheduler = get_learning_rate_scheduler(optimizer, deepspeed_args)
    model, optimizer, _, lr_scheduler = deepspeed.initialize(
       model=model.cuda().half(),
       model_parameters=param_groups,
       args=deepspeed_args,
       lr_scheduler=lr_scheduler,
       mpu=mpu,
       dist_init_required=False
    )
   

4. Trainer: training_step:

     # forward_step
     model.train()
     token_inputs = self._prepare_inputs(inputs) # in _prepare_inputs function: fp32_to_fp16
     outputs = model(**token_inputs)
     lm_loss = outputs[0]
    # Backward pass
    model.backward(lm_loss)
    model.step()
   

5. the above function some come from DeepspeedExample/Megatron-LM

some tries

1. It looks like vanishing gradient, i use torch.clamp to GPT2LMHeadModel:forward: shift_logits and shift_labels ，it did't work.

Start by reviewing examples/language-modeling/run_language_modeling.py, transformers/trainer.py, and transformers/modeling_gpt2.py, then reproduce the LogSoftmaxBackward NaN with the supplied DeepSpeed and fp16 configuration. Trace whether the failure occurs during model initialization, the forward loss, or model.backward(); done means the cause is isolated and the reported configuration no longer produces the error.