当在 SelfAttention Layer 之间插入 adapter 模块时，forward pass 出现 nan 和 inf 异常值导致 loss 为 0

[ZionDoki]

在使用类似 adapter 的方式，在 forward pass过程中出现异常值，初步排查推测应该是混合精度导致的溢出。但是在 HF 的 Trainer 部分我已经给了 fp16=True 的参数。数据的预处理方式参考的这个仓库。最让人头疼的是几乎一样的实现方式 ChatGLM 就没有类似的问题，我使用几乎完全一样的方法进行微调模型没有出现异常值，BP也稳定。查了很多资料都没有找到解决的办法。此前也尝试过使用ChatGLM2，也出现了类似的错误。

2024-01-09 08:17 | model stats:
>    {'Total': 6375790592, 'Trainable': 132206592, 'Percent': '2.0736%'}
2024-01-09 08:17 | init data collator √
2024-01-09 08:17 | Training  new model...
output has infs

Detected inf/nan during batch_number=0
Last 6 forward frames:
abs min  abs max  metadata


                  *** Starting batch number=0 ***
abs min  abs max  metadata
                  transformer.embedding.word_embeddings Embedding
0.00e+00 1.86e-01 weight
0.00e+00 6.48e+04 input[0]
0.00e+00 1.55e-01 output
                  transformer.embedding Embedding
0.00e+00 6.48e+04 input[0]
0.00e+00 1.55e-01 output
                  transformer.rotary_pos_emb RotaryEmbedding
     not a tensor input[0]
0.00e+00 1.00e+00 output
                  transformer.encoder.layers.0.input_layernorm RMSNorm
7.15e-07 8.48e-01 weight
0.00e+00 1.55e-01 input[0]
0.00e+00 2.91e+00 output
                  transformer.encoder.layers.0.self_attention.query_key_value Linear
0.00e+00 6.99e-01 weight
1.01e-06 7.47e+00 bias
0.00e+00 2.91e+00 input[0]
1.55e-06 9.55e+00 output
                  transformer.encoder.layers.0.self_attention.adapterk.wi Linear
0.00e+00 0.00e+00 weight
0.00e+00 3.48e+14 bias
6.14e-06 8.84e+00 input[0]
0.00e+00      inf output

插入 adapter 的代码：

class Adapter(nn.Module):
    def __init__(self, config, in_feats: int = None) -> None:
        super().__init__()
        if self.use_adapter:
            # print(in_feats or config.hidden_size)
            self.wi = nn.Linear(in_feats or config.hidden_size, in_feats or config.hidden_size)

    def forward(self, x):
        if not self.use_adapter:
            return x
        shortcut = x
        x = self.wi(x)
        return x + shortcut

class SelfAttentionWA(SelfAttention): 
    # 类似这种方式重载了所有类
    def __init__(self, config: ChatGLMConfig, layer_number, device=None):
        super().__init__(config, layer_number, device)

        if config.only_adapter_trainable:
            for params in super().parameters():
                params.requires_grad = False

        self.adapterk, self.adapterv = None, None
        if config.use_adapter:
            self.adapterk = Adapter(config, in_feats=self.qkv_hidden_size // 3)
            self.adapterv = Adapter(config, in_feats=self.qkv_hidden_size // 3)

    def forward(self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True):
          # hidden_states: [sq, b, h]
  
          # =================================================
          # Pre-allocate memory for key-values for inference.
          # =================================================
          # =====================
          # Query, Key, and Value
          # =====================
  
          # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
          mixed_x_layer = self.query_key_value(hidden_states)
          hidden_size = mixed_x_layer.shape[-1] // 3
  
          # 修改部分
          q, k, v = mixed_x_layer.permute(1, 0, 2).split(hidden_size, dim=-1)
          if self.adapterk:
              k = self.adapterk(k)
          if self.adapterv:
              v = self.adapterv(v)
          mixed_x_layer = torch.cat((q, k, v), dim=-1).permute(1, 0, 2)

          # ...

微调使用的是 huggingface 的 trainer：

from transformers import (
    AutoConfig,
    EarlyStoppingCallback,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
)

# .... 其他代码逻辑

training_args = Seq2SeqTrainingArguments(
    output_dir=f"output_{args.device}",
    evaluation_strategy="steps",
    logging_strategy="steps",
    logging_steps=LOGGING_STEPS,
    disable_tqdm=True,
    push_to_hub=False,
    per_device_train_batch_size=args.train_batch_size,
    per_device_eval_batch_size=args.eval_batch_size,
    eval_steps=args.save_steps,
    save_total_limit=5,
    fp16=getattr(args, "fp16", False), # 这里的fp16是True
    save_steps=args.save_steps,
    load_best_model_at_end=True,
    overwrite_output_dir=True,
    metric_for_best_model="eval_rouge-2",
    learning_rate=args.lr,
    num_train_epochs=args.epoch,
    report_to=["tensorboard"],
    warmup_steps=args.warmup_steps,
    debug="underflow_overflow" if args.debug else "",
)

#.... 其他代码逻辑

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=the_datasets.dataset["train"],
    eval_dataset=the_datasets.dataset["val"],
    data_collator=data_collator,
    compute_metrics=get_metric(token_handler.tokenizer),
    preprocess_logits_for_metrics=preprocess_logits_for_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=4)],
)

Expected behavior / 期待表现

1. 希望可以搞清楚ChatGLM和之后两个版本间存在什么样的差异导致了前者可以正常训练、后两者则出现了异常值。
2. 希望可以正常训练。