Downgrade to jax[cuda12]==0.4.31. I think there was some incompatibility between whatever wheels they are built with, and the drivers we have on g2.

What didn’t work:

• Clear LD_LIBRARY_PATH.
• This was not super helpful.

Minor complaint: Never have I managed to install Jax without pain, local/on g2/on TPU nodes, perhaps with the exception of Colab. It seems like a shared experience for others, others, and others.

Solution: Use LaTeX backend for matplotlib, because I reallly don’t want to use tikz. I’m not sure why but sys.path.append didn’t work in this case.

import os
import matplotlib.pyplot as plt

os.environ['PATH'] = '.../texlive/2024/bin/x86_64-linux:' + \
    os.environ['PATH']


plt.rcParams['text.usetex'] = True
plt.rcParams['font.family'] = 'serif'
plt.rcParams['text.latex.preamble'] = r'\usepackage{times}'

Solution: Add more LoRA adapters.

Solution: Filter val samples. Remove unreliable ones.

Solution:

1. Plot val/loss. My model was massively overfitted. I should have trained for 10 epochs instead of 30 epochs.
2. Add regularization by label smoothing, or reducing model complexity (e.g. LoRA target modules).
3. Read about model calibration (sequence probability aligned with actual accuracy), plotting reliability diagram. On Calibration of Modern Neural Networks
4. Overfitting to loss doesn’t always lead to overfitting to accuracy. How is it possible that validation loss is increasing while validation accuracy is increasing as well

Solution: Avg pool image latents https://huggingface.co/HuggingFaceM4/idefics2-8b/discussions/18. Reducing image_seq_length from 64 to 32 allows me to feed 3 examples in one batch.

Solution: Idefics2 had funny ways of expanding image tokens. I’m not sure why but using a longer max_length for sequence helped.

Solution: Set num_processes=null in default accelerate config. This will pass to xmp.spawn for multi-process scripts, which signals null for all chips available.

How: Stepping through tutorials like this. And really understand the mapping of one chip per process.

Solution: Do not use_cache if gradient checkpointing is enabled. PR

How:

1. Omer pointed out this shape error is due to gradient checkpointing.
2. Notice that previously Idefics had a similar patch to disable use_cache.

Solution:

1. Use the recommended runtime tpu-ubuntu2204-base.
2. Follow the correct tutorial. I was following Run a calculation on a Cloud TPU VM using PyTorch but actually v4-8 is a VM instance but v4-32 is a pod. So Run PyTorch code on TPU Pod slices is the one to go.
3. A pod and a single VM are difference because a pod needs multiprocessing (kudos Yoav), and xm.xla_device() cannot stand alone in global scope.
4. So really, follow the correct tutorial solved the problem.

How:

1. sudo lsof -w /dev/accel0
2. TPU logs cat /tmp/tpu_logs/tpu_driver.INFO and TPU_STDERR_LOG_LEVEL=0 TPU_MIN_LOG_LEVEL=0
3. XLA logs PT_XLA_DEBUG=1 and XLA trouble-shooting guidelines

Solution:

1. Set seed before model initialization.
2. Disable dropout layers (in the base model as well as the adapters). BatchNorm is also grad-accum unsafe.
3. CrossEntropyLoss (with reduce=’mean’ by default) may not be averaged if the -100 masks are applied to the targets.
4. Play with dtype. Floating point errors arise when I try to compute the loss in float32 vs float16. It makes things more complicated in gradients because I also used QLoRA.
5. Give up because training is only tractable with QLoRA in my case.

How: Observe numerical values. Set up two debuggers side by side.

Solution:

1. Instead of a fixed max sequence length, use dynamic padding.
2. After getting the attention mask from the processor, reset attention mask of padding tokens to zero.

How: Observe numerical values.

Solution: set_seed(42) before model initialization.

How: Control experiments and Google search

Solution: Downgrade transformers library.

How: Talk to kind experts

Solution: Online and offline models used different pre-processors.

How: Observe numerical values flowing into serving vs evaluation pipeline.

Solution: I put one battery in reverse.

How: Ask for help from TA