Loss is `nan` while training LLama4 via LoRA using torchtune

[vgoklani]

Using the default parameters:

tune run --nproc_per_node 8 lora_finetune_distributed --config /home/jovyan/torchtune/recipes/configs/llama4/scout_17B_16E_lora.yaml \
  checkpointer.checkpoint_dir=/data/hub/models--meta-llama--Llama-4-Scout-17B-16E/snapshots/14d516bdff6ac06cec40678529222f193386189c \
  tokenizer.path=/data/hub/models--meta-llama--Llama-4-Scout-17B-16E/snapshots/14d516bdff6ac06cec40678529222f193386189c/tokenizer.model \
  output_dir=/data/output/llama4a \
  metric_logger._component_=torchtune.training.metric_logging.WandBLogger \
  metric_logger.project="llama4_lora" \
  log_every_n_steps=5

right after the first step:

|Loss: nan: 0%| | 3/3073 [00:15<4:04:50, 4.79s/it]

Does this implementation actually work? Has it been tested?

[vgoklani]

This is the minimum possible:

!tune run --nproc_per_node 8 lora_finetune_distributed --config /home/jovyan/torchtune/recipes/configs/llama4/scout_17B_16E_lora.yaml \
  checkpointer.checkpoint_dir=/data/hub/models--meta-llama--Llama-4-Scout-17B-16E-Instruct/snapshots/92f3b1597a195b523d8d9e5700e57e4fbb8f20d3 \
  tokenizer.path=/data/hub/models--meta-llama--Llama-4-Scout-17B-16E/snapshots/14d516bdff6ac06cec40678529222f193386189c/tokenizer.model \
model.fusion_trainable="frozen" \
model.apply_lora_to_mlp=true \
  output_dir=/data/output/llama4a \
    optimizer.lr=1e-7 \
    optimizer.weight_decay=0.0 \
  clip_grad_norm=0.001 \
  model.lora_attn_modules="""['q_proj',]""" \
  metric_logger._component_=torchtune.training.metric_logging.WandBLogger \
  metric_logger.project="llama4_lora" \
  log_every_n_steps=5

Are you sure this implementation works? The 2nd training step is nan

[krammnic]

Thanks for the issue! Will try to reproduce and write you back

[ebsmothers]

Hi @vgoklani thanks for creating the issue and sorry for the delay here, I missed this one over the weekend. Yes, we have tested our Llama4 LoRA configs -- you can see some loss curves in the initial PR #2583. I just reran tune run --nproc_per_node 8 lora_finetune_distributed --config llama4/scout_17B_16E_lora on main and can confirm that I do not see any NaN loss. Can you share more details about your environment? Namely if you can (a) paste your full config, (b) share the outputs of pip list (or similar), (c) share the kind of hardware you're running on, that would be helpful.

[krammnic]

Sorry for the delay, I don't see any nan losses with LoRA either. Maybe your computational resources are insufficient for the Llama4 lora procedure and you see nan loss before the OOM?

[vgoklani]

I tried this during the llama hackathon in NYC, and had access to an 8xH200. Unfortunately I no longer have access to that machine. I used their base image and just installed the latest stable version of pytorch, torchtune, and huggingface_hub. I then downloaded the weights and ran your base LoRA recipe for Scout. The nan showed up consistently after the first step.

We are hoping our blackwell gpus arrive by August, and then I will try this again... :(

[ebsmothers]

@vgoklani ah sorry to hear! I suspect the latest stable of PyTorch and/or torchtune may have caused you issues here. We have landed some fixes that are only available in the nightlies so would recommend to use that (we will soon have a stable torchtune release, but next PyTorch release is still several more weeks away). I am gonna close the issue, but feel free to reopen if you need more support for anything. Sorry again!

[vgoklani]

Thanks @ebsmothers - could you be more specific here:

I suspect the latest stable of PyTorch and/or torchtune may have caused you issues her

[ebsmothers]

@vgoklani to clarify, I mean that there were changes landed in both PyTorch and torchtune since the previous stable release to enable correct training of MoEs. So if you are on the stable releases, these changes would not be available to you. Hence why I'd recommend the nightlies. Actually, we assert this for Llama4 full finetuning here, but did not add the corresponding assertion for the LoRA variant. So I can open a quick PR to add that and save future people the debugging headache. Sorry about that!