llama2 training has nan

[LZY-the-boys]

System Info

• transformers version: 4.31.0
• Platform: Linux-5.19.0-42-generic-x86_64-with-glibc2.35
• Python version: 3.9.16
• Huggingface_hub version: 0.14.1
• Safetensors version: 0.3.1
• Accelerate version: 0.21.0
• Accelerate config: - compute_environment: LOCAL_MACHINE
  - distributed_type: NO
  - mixed_precision: fp16
  - use_cpu: False
  - num_processes: 1
  - machine_rank: 0
  - num_machines: 1
  - gpu_ids: all
  - rdzv_backend: static
  - same_network: True
  - main_training_function: main
  - downcast_bf16: no
  - tpu_use_cluster: False
  - tpu_use_sudo: False
  - tpu_env: []
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

I am training llama2-7b using huggingface trainer. I find that it will occur nan in forward when training with batchsize > 1 for llama2, however when batchsize=1 there is no error.
I dive into it and find that the nan occurs in layer.31.input_layer_norm, which is caused by inf in layers.30.mlp forward after the post_layer_norm, and this inf may comes from huge value in hidden_size. However, this doesn't explain why llama1 and llama2 with batchsize=1 can work, which also has huge outliners in hidden_size.
The code I use is like this:
The dataset format is: write a xxx. ###sentences: xxxx
The checkpoint of meta-llama/Llama-2-7b-hf and original meta checkpoint converted by transformers code are both tried.

trainer = transformers.Trainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=val_data,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=args.micro_batch,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        warmup_ratio=args.warmup_ratio,
        num_train_epochs=args.num_epoch,
        learning_rate=3e-4,
        fp16=True,
        logging_steps=args.log_steps,
        logging_first_step=True, # convenient
        evaluation_strategy="no",
        save_strategy=args.save_strategy,
        eval_steps=None,
        save_steps=args.save_steps,
        output_dir=args.output_path,
        load_best_model_at_end= False,
        ddp_find_unused_parameters=False if ddp else None,
        report_to="wandb" if args.wandb else [],
        ignore_data_skip=args.ignore_data_skip,
    ),
    data_collator=PROMPT.data_collator()
)
model.config.use_cache = False
if list(pathlib.Path(args.output_path).glob("checkpoint-*")):
    trainer.train(resume_from_checkpoint=True)
else:
    trainer.train()
trainer.save_state()
model.save_pretrained(args.output_path)

Expected behavior

The training should has no nan in forward, thus loss will be normal in backward

[tannonk]

What are you setting as your pad_token? I was getting nan while running batched inference (see post here) and am wondering if this might be related to your issue?

[LZY-the-boys]

I set tokenizer.pad_token_id = 0, which is the same as llama1 (:

[sijeh]

I have the same question!

[hangzhang-nlp]

Me too！

[ArthurZucker]

Hey! Thanks all for reporting this. I would suggest using the following:

        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights + attention_mask
            dtype_min = torch.tensor(
                torch.finfo(attn_weights.dtype).min, device=attn_weights.device, dtype=attn_weights.dtype
            )
            attn_weights = torch.max(attn_weights, dtype_min)

This was removed from:

https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L343-L348

because it is not used in the original model. This should help. Otherwise adding

        # clamp inf values to enable fp16 training
        if hidden_states.dtype == torch.float16:
            max_dtype = torch.finfo(hidden_states.dtype).max
            clamp_value = torch.where(torch.isinf(hidden_states).any(), max_dtype - 1000, max_dtype)
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)

will also help mitigate this

[ArthurZucker]

Another fix is to train with bfloat16

[arazd]

Another fix is to train with bfloat16

+1, this actually worked for me with accelerate FSDP training

[rahul003]

Can you share what config you used for training @arazd I still see nan with bf16

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.