You are training your latest AI model, anxiously watching as the loss steadily decreases when suddenly \u2014 boom! Your logs are flooded with NaNs (Not a Number) \u2014 your model is irreparably corrupted and you\u2019re left staring at your screen in despair. To make matters worse, the NaNs don\u2019t appear consistently. Sometimes your model trains just fine; other times, it fails inexplicably. Sometimes it will crash immediately, sometimes after many days of training.<\/p>\n
NaNs in Deep Learning<\/a> workloads are amongst the most frustrating issues to encounter. And because they often appear sporadically \u2014 triggered by a specific combination of model state, input data, and stochastic factors \u2014 they can be incredibly difficult to reproduce and debug.<\/p>\n Given the considerable cost of training AI models and the potential waste caused by NaN failures, it is recommended to have dedicated tools for capturing and analyzing NaN occurrences. In a\u00a0previous post<\/a>, we discussed the challenge of debugging NaNs in a TensorFlow training workload. We proposed an efficient scheme for capturing and reproducing NaNs and shared a sample TensorFlow implementation. In this post, we adopt and demonstrate a similar mechanism for debugging NaNs in PyTorch workloads. The general scheme is as follows:<\/p>\n On each training step:<\/p>\n Although this scheme can be easily implemented in native PyTorch, we will take the opportunity to demonstrate some of the conveniences of\u00a0PyTorch Lightning<\/a>\u00a0\u2014 a powerful open-source framework designed to streamline the development of machine learning (ML) models. Built on PyTorch, Lightning abstracts away many of the boiler-plate components of an ML experiment, such as training loops, data distribution, logging, and more, enabling developers to focus on the core logic of their models.<\/p>\n To implement our NaN capturing scheme, we will use\u00a0Lightning\u2019s callback<\/a>\u00a0interface \u2014 a dedicated structure that enables inserting custom logic at specific points during the flow of execution.<\/p>\n Importantly, please do not view our choice of Lightning or any other tool or technique that we mention as an endorsement of its use. The code that we will share is intended for demonstrative purposes \u2014 please do not rely on its correctness or optimality.<\/p>\n Many thanks to\u00a0Rom Maltser<\/a>\u00a0for his contributions to this post.<\/p>\n To implement our NaN capturing solution, we create a NaNCapture Lightning callback. The constructor receives a directory path for storing\/loading checkpoints and sets up the NaNCapture state. We also define utilities for checking for NaNs, storing checkpoints, and halting the training job.<\/p>\n We begin by implementing the\u00a0on_train_batch_start<\/a>\u00a0hook to store a copy of each input batch. In case of a NaN event, this batch will be stored in the checkpoint.<\/p>\n Next we implement the\u00a0on_before_optimizer_step<\/a>\u00a0hook. Here, we check for NaN entries in all of the gradient tensors. If found, we store a checkpoint with the uncorrupted model weights and halt the training.<\/p>\n To enable reproducibility, we include the NaNCapture state in the checkpoint by appending it to the training state dictionary. Lightning provides dedicated utilities for saving and loading a\u00a0callback state<\/a>:<\/p>\n We have described how our NaNCapture callback can be used to store the training state that resulted in a NaN, but how do we reload this state in order to reproduce the issue and debug it? To accomplish this, we leverage Lightning\u2019s dedicated data loading class,\u00a0LightningDataModule<\/a>.<\/p>\n In the code block below, we extend the\u00a0LightningDataModule<\/a>\u00a0class to allow injecting a fixed training input batch. This is achieved by overriding the\u00a0on_before_batch_transfer<\/a>\u00a0hook, as shown below:<\/p>\n The final step is modifying the\u00a0on_train_start<\/a>\u00a0hook of our NaNCapture callback to inject the stored training batch into the\u00a0LightningDataModule<\/a>.<\/p>\n In the next section we will demonstrate the end-to-end solution using a toy example.<\/p>\n To test our new callback, we create a\u00a0resnet50<\/a>-based image classification model with a loss function deliberately designed to trigger NaN occurrences.<\/p>\n Instead of using the standard\u00a0CrossEntropy<\/a>\u00a0loss, we compute\u00a0binary_cross_entropy_with_logits<\/a>\u00a0for each class independently and divide the result by the number of samples belonging to that class. Inevitably, we will encounter a batch in which one or more classes are missing, leading to a divide-by-zero operation, resulting in NaN values and corrupting the model.<\/p>\n The implementation below follows Lightning\u2019s\u00a0introductory tutorial<\/a>.<\/p>\n We define a synthetic dataset and encapsulate it in our Finally, we initialize a Lightning\u00a0Trainer<\/a>\u00a0with our NaNCapture callback and call trainer.fit with our Lightning module and Lightning DataModule.<\/p>\n After a number of training steps, a NaN event will occur. At this point a checkpoint is saved with the full training state and the training is halted.<\/p>\n When the script is run again the exact state that caused the NaN will be reloaded allowing us to easily reproduce the issue and debug its root cause.<\/p>\n To assess the impact of our NaNCapture callback on runtime performance, we modified our experiment to use\u00a0CrossEntropyLoss<\/a>\u00a0(to avoid NaNs) and measured the average throughput when running with and without NaNCapture callback. The experiments were conducted on an\u00a0NVIDIA L40S GPU<\/a>, with a\u00a0PyTorch 2.5.1 Docker<\/a>\u00a0image.<\/p>\n For our toy model, the NaNCapture callback adds a minimal 1.5% overhead to the runtime performance \u2014 a small price to pay for the valuable debugging capabilities it provides.<\/p>\n Naturally, the actual overhead will depend on the specifics of the model and runtime environment.<\/p>\n The solution we have described henceforth will succeed in reproducing the training state provided that the model does not include any randomness. However, introducing stochasticity into the model definition is often critical for convergence. A common example of a stochastic layer is\u00a0torch.nn.Dropout<\/a>.<\/p>\n You may find that your NaN event depends on the precise state of randomness when the failure occurred. Consequently, we would like to enhance our NaNCapture callback to capture and restore the random state at the point of failure. The random state is determined by a number of libraries. In the code block below, we attempt to capture the full state of randomness:<\/p>\n Importantly, setting the random state may not guarantee full\u00a0reproducibility<\/a>. The GPU owes its power to its massive parallelism. In some GPU operations, multiple threads may read or write concurrently to the same memory locations resulting in nondeterminism. PyTorch allows for some control over this via its\u00a0use_deterministic_algorithms<\/a>, but this may impact the runtime performance. Additionally, there is a possibility that the NaN event will not reproduced once this configuration setting is changed. Please see the PyTorch documentation on\u00a0reproducibility<\/a>\u00a0for more details.<\/p>\n Encountering NaN failures is one of the most discouraging events that can happen in machine learning development. These errors not only waste valuable computation and development resources, but often indicate fundamental issues in the model architecture or experiment design. Due to their sporadic, sometimes elusive nature, debugging NaN failures can be a nightmare.<\/p>\n This post introduced a proactive approach for capturing and reproducing NaN errors using a dedicated Lightning callback. The solution we shared is a proposal which can be modified and extended for your specific use case.<\/p>\n While this solution may not address every possible NaN scenario, it significantly reduces debugging time when applicable, potentially saving developers countless hours of frustration and wasted effort.<\/p>\n The post Debugging the Dreaded NaN<\/a> appeared first on Towards Data Science<\/a>.<\/p>\n<\/div>\n\n
NaNCapture Callback<\/h2>\n
import os\nimport torch\nfrom copy import deepcopy\nimport lightning.pytorch as pl\n\nclass NaNCapture(pl.Callback):\n\n def __init__(self, dirpath: str):\n # path to checkpoint\n self.dirpath = dirpath\n \n # update to True when Nan is identified\n self.nan_captured = False\n \n # stores a copy of the last batch\n self.last_batch = None\n self.batch_idx = None\n\n @staticmethod\n def contains_nan(tensor):\n return torch.isnan(tensor).any().item()\n # alternatively check for finite\n # return not torch.isfinite(tensor).item()\n\n @staticmethod\n def halt_training(trainer):\n trainer.should_stop = True\n # communicate stop command to all other ranks\n trainer.strategy.reduce_boolean_decision(trainer.should_stop,\n all=False)\n\n def save_ckpt(self, trainer):\n os.makedirs(self.dirpath, exist_ok=True)\n # include trainer.global_rank to avoid conflict\n filename = f\"nan_checkpoint_rank_{trainer.global_rank}.ckpt\"\n full_path = os.path.join(self.dirpath, filename)\n print(f\"saving ckpt to {full_path}\")\n trainer.save_checkpoint(full_path, False)<\/code><\/pre>\nCallback Function: on_train_batch_start<\/h3>\n
Callback Function: on_before_optimizer_step<\/h3>\n
Python\"> def on_before_optimizer_step(self, trainer, pl_module, optimizer):\n if not self.nan_captured:\n # Check if gradients contain NaN\n grads = [p.grad.view(-1) for p in pl_module.parameters()\n if p.grad is not None]\n all_grads = torch.cat(grads)\n if self.contains_nan(all_grads):\n print(\"nan found\")\n self.save_ckpt(trainer)\n self.halt_training(trainer)\n<\/code><\/pre>\nCapturing the Training State<\/h3>\n
def state_dict(self):\n d = {\"nan_captured\": self.nan_captured}\n if self.nan_captured:\n d[\"last_batch\"] = self.last_batch\n return d\n\n\n def load_state_dict(self, state_dict):\n self.nan_captured = state_dict.get(\"nan_captured\", False)\n if self.nan_captured:\n self.last_batch = state_dict[\"last_batch\"]<\/code><\/pre>\nReproducing the NaN Occurrence<\/h2>\n
DataModule Function: on_before_batch_transfer<\/h3>\n
from lightning.pytorch import LightningDataModule\n\nclass InjectableDataModule(LightningDataModule):\n\n def __init__(self):\n super().__init__()\n self.cached_batch = None\n\n def set_custom_batch(self, batch):\n self.cached_batch = batch\n\n def on_before_batch_transfer(self, batch, dataloader_idx):\n if self.cached_batch:\n return self.cached_batch\n return batch<\/code><\/pre>\nCallback Function: on_train_start<\/h3>\n
def on_train_start(self, trainer, pl_module):\n if self.nan_captured:\n datamodule = trainer.datamodule\n datamodule.set_custom_batch(self.last_batch)<\/code><\/pre>\nToy Example<\/h2>\n
import lightning.pytorch as pl\nimport torch\nimport torchvision\nimport torch.nn.functional as F\n\nnum_classes = 20\n\n\n# define a lightning module\nclass ResnetModel(pl.LightningModule):\n def __init__(self):\n \"\"\"Initializes a new instance of the MNISTModel class.\"\"\"\n super().__init__()\n self.model = torchvision.models.resnet50(num_classes=num_classes)\n\n def forward(self, x):\n return self.model(x)\n\n def training_step(self, batch, batch_nb):\n x, y = batch\n outputs = self(x)\n # uncomment for default loss\n # return F.cross_entropy(outputs, y)\n \n # calculate binary_cross_entropy for each class individually\n losses = []\n for c in range(num_classes):\n count = torch.count_nonzero(y==c)\n masked = torch.where(y==c, 1., 0.)\n loss = F.binary_cross_entropy_with_logits(\n outputs[..., c],\n masked,\n reduction='sum'\n )\n mean_loss = loss\/count # could result in NaN\n losses.append(mean_loss)\n total_loss = torch.stack(losses).mean()\n return total_loss\n\n def configure_optimizers(self):\n return torch.optim.Adam(self.parameters(), lr=0.02)<\/code><\/pre>\nInjectableDataModule<\/code> class:<\/p>\nimport os\nimport random\nfrom torch.utils.data import Dataset, DataLoader\n\nbatch_size = 128\nnum_steps = 800\n\n# A dataset with random images and labels\nclass FakeDataset(Dataset):\n def __len__(self):\n return batch_size*num_steps\n\n def __getitem__(self, index):\n rand_image = torch.randn([3, 224, 224], dtype=torch.float32)\n label = torch.tensor(random.randint(0, num_classes-1),\n dtype=torch.int64)\n return rand_image, label\n\n\n\n# define a lightning datamodule\nclass FakeDataModule(InjectableDataModule):\n\n def train_dataloader(self):\n dataset = FakeDataset()\n return DataLoader(\n dataset,\n batch_size=batch_size,\n num_workers=os.cpu_count(),\n pin_memory=True\n )<\/code><\/pre>\nimport time\n\nif __name__ == \"__main__\":\n\n # Initialize a lightning module\n lit_module = ResnetModel()\n\n # Initialize a DataModule\n mnist_data = FakeDataModule()\n\n # Train the model\n ckpt_dir = \".\/ckpt_dir\"\n trainer = pl.Trainer(\n max_epochs=1,\n callbacks=[NaNCapture(ckpt_dir)]\n )\n\n ckpt_path = None\n \n # check is nan ckpt exists\n if os.path.isdir(ckpt_dir):\n\n # check if nan ckpt exists\n if os.path.isdir(ckpt_dir):\n dir_contents = [os.path.join(ckpt_dir, f)\n for f in os.listdir(ckpt_dir)]\n ckpts = [f for f in dir_contents\n if os.path.isfile(f) and f.endswith('.ckpt')]\n if ckpts:\n ckpt_path = ckpts[0]\n\n t0 = time.perf_counter()\n trainer.fit(lit_module, mnist_data, ckpt_path=ckpt_path)\n print(f\"total runtime: {time.perf_counter() - t0}\")<\/code><\/pre>\nPerformance Overhead<\/h2>\n
![\"\"](\"https:\/\/i0.wp.com\/towardsdatascience.com\/wp-content\/uploads\/2025\/02\/image-19.png?resize=738%2C106&ssl=1\")
How to Handle Stochasticity<\/h2>\n
import os\nimport torch\nimport random\nimport numpy as np\nfrom copy import deepcopy\nimport lightning.pytorch as pl\n\nclass NaNCapture(pl.Callback):\n\n def __init__(self, dirpath: str):\n # path to checkpoint\n self.dirpath = dirpath\n \n # update to True when Nan is identified\n self.nan_captured = False\n \n # stores a copy of the last batch\n self.last_batch = None\n self.batch_idx = None\n\n # rng state\n self.rng_state = {\n \"torch\": None,\n \"torch_cuda\": None,\n \"numpy\": None,\n \"random\": None\n }\n\n @staticmethod\n def contains_nan(tensor):\n return torch.isnan(tensor).any().item()\n # alternatively check for finite\n # return not torch.isfinite(tensor).item()\n\n @staticmethod\n def halt_training(trainer):\n trainer.should_stop = True\n trainer.strategy.reduce_boolean_decision(trainer.should_stop,\n all=False)\n\n def save_ckpt(self, trainer):\n os.makedirs(self.dirpath, exist_ok=True)\n # include trainer.global_rank to avoid conflict\n filename = f\"nan_checkpoint_rank_{trainer.global_rank}.ckpt\"\n full_path = os.path.join(self.dirpath, filename)\n print(f\"saving ckpt to {full_path}\")\n trainer.save_checkpoint(full_path, False)\n\n def on_train_start(self, trainer, pl_module):\n if self.nan_captured:\n # inject batch\n datamodule = trainer.datamodule\n datamodule.set_custom_batch(self.last_batch)\n\n def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):\n if self.nan_captured:\n # restore random state\n torch.random.set_rng_state(self.rng_state[\"torch\"])\n torch.cuda.set_rng_state_all(self.rng_state[\"torch_cuda\"])\n np.random.set_state(self.rng_state[\"numpy\"])\n random.setstate(self.rng_state[\"random\"])\n else:\n # capture current batch\n self.last_batch= deepcopy(batch)\n self.batch_idx = batch_idx\n \n # capture current random state\n self.rng_state[\"torch\"] = torch.random.get_rng_state()\n self.rng_state[\"torch_cuda\"] = torch.cuda.get_rng_state_all()\n self.rng_state[\"numpy\"] = np.random.get_state()\n self.rng_state[\"random\"] = random.getstate()\n \n def on_before_optimizer_step(self, trainer, pl_module, optimizer):\n if not self.nan_captured:\n # Check if gradients contain NaN\n grads = [p.grad.view(-1) for p in pl_module.parameters()\n if p.grad is not None]\n all_grads = torch.cat(grads)\n if self.contains_nan(all_grads):\n print(\"nan found\")\n self.save_ckpt(trainer)\n self.halt_training(trainer)\n\n def state_dict(self):\n d = {\"nan_captured\": self.nan_captured}\n if self.nan_captured:\n d[\"last_batch\"] = self.last_batch\n d[\"rng_state\"] = self.rng_state\n return d\n\n def load_state_dict(self, state_dict):\n self.nan_captured = state_dict.get(\"nan_captured\", False)\n if self.nan_captured:\n self.last_batch = state_dict[\"last_batch\"]\n self.rng_state = state_dict[\"rng_state\"]<\/code><\/pre>\nSummary<\/h2>\n