All layers translators with nan loss while training the lens

[AmeenAli]

Hello!

Thanks for sharing this amazing work!

I am trying to train the lens over a new dataset HF Dataset (note that the original "the pile" dataset was removed from the internet because of a DMCA claim - as i understood it. as follows :

tuned-lens train --model.name EleutherAI/pythia-160m-deduped --output ./output/ --per_gpu_batch_size=6 --data.name HuggingFaceH4/ultrachat_200k --split train_gen --text_column prompt --wandb test

However, once I look into wandb logs i notice that loss/translator_1 - 12 are all NaNs also the weight_norms_1-12.
I have also tried different datasets, but the issue still exists. Any idea what is off?

Thank you!

[AmeenAli]

• Getting nans logits while also evaluating on the provided pretrained lens

Traceback (most recent call last): File "/home/ameenali/anaconda3/envs/myenv2/bin/tuned-lens", line 8, in <module> sys.exit(main()) File "/home/ameenali/anaconda3/envs/myenv2/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper return f(*args, **kwargs) File "/home/ameenali/anaconda3/envs/myenv2/lib/python3.10/site-packages/tuned_lens/__main__.py", line 30, in main prog.execute() File "/home/ameenali/anaconda3/envs/myenv2/lib/python3.10/site-packages/tuned_lens/__main__.py", line 20, in execute self.command.execute() File "/home/ameenali/anaconda3/envs/myenv2/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 16, in decorate_autocast return func(*args, **kwargs) File "/home/ameenali/anaconda3/envs/myenv2/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context return func(*args, **kwargs) File "/home/ameenali/anaconda3/envs/myenv2/lib/python3.10/site-packages/tuned_lens/scripts/eval_loop.py", line 269, in execute assert not th.isnan(output.logits).any(), "Logits are NaN" AssertionError: Logits are NaN

[levmckinney]

To your first question have you tried setting the --precision to float32?