LLM Debugging: Deterministic Replay with FSDP

Distributed LLM training failures plague even the most seasoned ML teams. A single, elusive bug can halt development for days, draining engineering resources and delaying product launches. Pinpointing the root cause is exceptionally difficult when errors appear intermittently, a common symptom of non-deterministic behavior. This post explores deterministic replay GPU training FSDP, a powerful approach to conquer these complex debugging challenges.

What is Deterministic Replay for GPU Training?

Deterministic replay for GPU training involves meticulously capturing the entire execution state of a distributed training run and then faithfully reproducing it. Imagine recording a complex orchestra performance, then playing it back note-for-note, exactly as it happened. That is the essence of deterministic replay. It allows developers to rerun a failed training job, observe the exact sequence of events, and identify the precise point of failure, even across multiple GPUs and nodes.

This technique solves the infuriating problem of “it works on my machine, sometimes” or “the error disappears when I add print statements.” It is particularly valuable for complex, distributed setups. Older debugging methods often relied on speculative fixes or extensive logging, which could alter the very behavior they aimed to observe. Deterministic replay replaces this guesswork with precise, repeatable execution. It is crucial for anyone wrestling with the unpredictable nature of large-scale AI infrastructure.

Why deterministic replay GPU training FSDP Matters in 2026

The rapid scaling of Large Language Models (LLMs) means training often spans hundreds or thousands of GPUs. This complexity introduces new failure modes. Non-deterministic bugs, which manifest differently across identical runs, can consume significant engineering cycles. For example, a major AI lab reported spending over 40% of their debugging time on non-reproducible issues in distributed training. Deterministic replay GPU training FSDP directly addresses this, dramatically reducing time spent on elusive bugs.

Consider a scenario where a multi-billion parameter LLM fails intermittently during its 100th training epoch. Without deterministic replay, engineers might spend days or weeks attempting to recreate the specific conditions. With replay, the exact failure state can be revisited repeatedly. This shortens debugging from weeks to hours. It allows companies like Anthropic or Google DeepMind to maintain aggressive LLM development schedules. Furthermore, by ensuring reproducibility, it improves the integrity of research findings and accelerates model development by 20-30% in some cases, significantly impacting development costs and time-to-market.

Core Concepts and Architecture

Challenges of non-determinism in FSDP and mixed-precision training (gradient accumulation, NCCL atomicity)

Non-determinism in distributed training often arises from subtle interactions within parallel systems. FSDP (Fully Sharded Data Parallel) training, especially with mixed precision, introduces unique complexities. Here, operations like gradient accumulation can become non-deterministic due to floating-point arithmetic ordering. Different GPU execution orders or communication timings can slightly alter summation results. Similarly, NCCL (NVIDIA Collective Communications Library) operations, while generally designed for correctness, can expose non-determinism related to message ordering or network latency, particularly under contention or specific hardware configurations.

To counteract this, one must understand that floating-point addition is not associative. (a + b) + c is not always exactly equal to a + (b + c) due to precision limits. In parallel sum reductions, the order of additions can vary, leading to minor differences. Similarly, in NCCL, the atomic operations at the hardware level might not guarantee a fixed ordering across runs if underlying OS or driver scheduling is variable.

# Example: Non-deterministic gradient accumulation (conceptual)
import torch

def train_step_non_deterministic(model, optimizer, data, labels):
    # Potential for non-determinism if operations order changes
    # across GPUs or runs, especially with mixed precision.
    # Different GPU scheduling might lead to slightly different order of additions
    # for gradients during reduction.
    outputs = model(data)
    loss = some_loss_fn(outputs, labels)
    loss.backward()
    # Gradients are accumulated across micro-batches or across GPUs
    # The sum order can change, leading to non-determinism.
    # optimizer.step() # Would apply these potentially non-deterministic gradients

A common pitfall is assuming that identical input data guarantees identical outputs in distributed, mixed-precision training. The precise order of floating-point operations can change, leading to different numerical results that accumulate over many steps.

Techniques for capturing and replaying GPU memory states and communication primitives (e.g., NCCL operations)

Capturing GPU memory states and communication primitives forms the backbone of deterministic replay. This involves intercepting specific API calls that modify GPU memory or initiate communication. For memory, one can hook into CUDA memory allocation and deallocation calls, saving the content of allocated buffers at critical points. For communication, intercepting NCCL function calls (like ncclAllReduce, ncclBroadcast) allows logging the arguments, participating ranks, and even the data payloads.

The process typically involves a runtime instrumentation layer. This layer records every relevant event: memory writes, kernel launches, and data sent/received over NCCL. During replay, these recorded events are then “fed” back into the system, ensuring the exact same sequence and data values are presented to the GPU and communication fabric. This ensures bit-for-bit identical execution.

# Conceptual example of a hook for NCCL all_reduce in PyTorch
# This would require deeper framework modifications or LD_PRELOAD
import torch.distributed as dist

# Imagine a custom logging hook for NCCL operations
def instrumented_all_reduce(tensor, op=dist.ReduceOp.SUM, group=None, async_op=False):
    # Log information before the actual all_reduce call
    print(f"[{dist.get_rank()}] Recording all_reduce for tensor of shape {tensor.shape} at step X")
    # Store tensor state or metadata for later replay
    # actual_all_reduce(tensor, op, group, async_op) # The original NCCL call
    dist.all_reduce(tensor, op=op, group=group, async_op=async_op)
    print(f"[{dist.get_rank()}] all_reduce completed for tensor of shape {tensor.shape} at step X")

# In a replayed run, this would instead restore and provide the recorded data
# from the log, bypassing the actual NCCL call.

A common pitfall is trying to capture everything. This generates enormous log files and significantly slows down training. Focus on capturing only the inputs to non-deterministic operations or specific memory regions known to be critical.

Leveraging checkpointing and snapshotting strategies for replay initiation points

Replaying an entire LLM training run from scratch to reproduce a bug that occurs late in training is impractical. This is where checkpointing and snapshotting strategies become vital. Instead of replaying from the very beginning, a full system snapshot (including GPU memory, CPU memory, and process state) can serve as a replay initiation point. This means you can save the complete state of your training job at regular intervals or right before a suspected problematic region.

When a failure occurs, you load a recent snapshot and then initiate deterministic replay from that point forward. This significantly reduces the overhead of replay. PyTorch’s torch.save and FSDP’s state_dict can capture model weights and optimizer states, but a true system snapshot needs to extend to GPU memory contents, CUDA streams, and potentially even NCCL communicator states. This often requires custom tooling that interacts directly with the CUDA driver API.

# Conceptual example: Saving FSDP state and a 'system snapshot' ID
import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

def save_snapshot_and_checkpoint(model, optimizer, step_id, snapshot_path):
    # Save model and optimizer state using FSDP's state_dict
    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
        cpu_state = model.state_dict()
        optimizer_state = FSDP.get_optimizer_state_dict(model, optimizer)

    # For a *true* deterministic replay snapshot, one would also need
    # to capture full GPU memory, CUDA stream states, and potentially
    # process memory of each rank. This is highly tool-dependent.
    # For illustration, we simulate saving a "system snapshot token".
    snapshot_data = {
        "step": step_id,
        "model_state": cpu_state,
        "optimizer_state": optimizer_state,
        "system_snapshot_token": f"snapshot_gpu_{dist.get_rank()}_step_{step_id}",
    }
    torch.save(snapshot_data, f"{snapshot_path}/snapshot_rank_{dist.get_rank()}_{step_id}.pt")
    print(f"[{dist.get_rank()}] Saved checkpoint and simulated system snapshot for step {step_id}")

# In a replay scenario, this snapshot is loaded, and then the
# captured trace from the original run continues from this point.

A common misconception is that standard PyTorch checkpoints are sufficient for deterministic replay. While they capture model weights, they typically do not capture the exact, bit-for-bit GPU memory contents or CUDA runtime state required for perfect replay.

Strategies for isolating non-deterministic components (random seeds, data loading, race conditions)

Before resorting to full deterministic replay, it is often productive to isolate and eliminate known sources of non-determinism. Random seeds are a primary culprit; ensure every random operation, across CPU, CUDA, and library calls (e.g., NumPy, torchvision), is seeded deterministically. Data loading can introduce non-determinism if shuffled data is not consistent across workers or if augmentation pipelines are not fixed. Race conditions, particularly in multi-threaded data preprocessing or custom CUDA kernels, are another source.

The strategy involves systematically making components deterministic one by one. Start by fixing all random seeds. Then, fix data loading orders and transformations. If non-determinism persists, investigate potential race conditions within custom code or interactions between libraries. Logging the hashes of tensors before and after suspected operations can reveal where values diverge.

# Example: Setting seeds for determinism
import torch
import numpy as np
import random
import os

def set_all_seeds(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # Additional: Set environment variables for cuDNN determinism
    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8' # For deterministic matmul
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    print(f"All seeds set to {seed}")

# Usage: Call at the very beginning of your script
# set_all_seeds(42)

A common pitfall is forgetting to seed all sources of randomness. Many libraries have their own random number generators, and even seemingly deterministic operations can depend on an unseeded torch.cuda.manual_seed_all or numpy.random.seed.

Tooling and framework extensions (e.g., custom PyTorch hooks, instrumenting communication libraries) for deterministic replay

Achieving true deterministic replay often requires sophisticated tooling. This goes beyond simple torch.autograd.Function hooks. It involves instrumenting the low-level CUDA driver API (libcuda.so), the NCCL library (libnccl.so), and PyTorch’s internals. Tools might use LD_PRELOAD to inject custom libraries that intercept calls to these underlying systems. These custom libraries would then log the function arguments, return values, and memory changes.

For PyTorch, specific hooks can be registered with torch.autograd.graph.saved_tensors_hooks or by subclassing torch.autograd.Function. More generally, frameworks like “Determinism for ML” (research project) or custom extensions built on top of NVIDIA’s Nsight tools offer a path. These tools record system calls, memory allocations, kernel launches, and inter-process communications. During replay, they ensure these same calls are made with identical inputs, effectively recreating the original execution environment.

# Conceptual example: Using a custom PyTorch autograd hook for debugging
# This does not enable full deterministic replay, but shows how hooks work.
def forward_hook(module, input, output):
    # Log or inspect inputs/outputs for debugging
    print(f"[{dist.get_rank()}] Forward hook on {type(module).__name__}: input shape {input[0].shape}, output shape {output.shape}")

def backward_hook(module, grad_input, grad_output):
    # Log or inspect gradients for debugging
    print(f"[{dist.get_rank()}] Backward hook on {type(module).__name__}: grad_output shape {grad_output[0].shape}")

# Example usage:
# model = MyLLMModel()
# for name, module in model.named_modules():
#     # Attach hooks to specific modules for focused inspection
#     if isinstance(module, torch.nn.Linear):
#         module.register_forward_hook(forward_hook)
#         module.register_backward_hook(backward_hook)

A common pitfall is that simply adding PyTorch hooks provides insufficient information for full deterministic replay. These hooks operate at a higher abstraction level than the low-level CUDA and NCCL operations, which are the true source of many non-deterministic behaviors.

Getting Started with deterministic replay GPU training FSDP: Step-by-Step

Implementing full-scale deterministic replay is complex, but you can build a proof-of-concept by focusing on critical components. Here’s how to set up a minimal FSDP training loop and then introduce a conceptual logging mechanism.

Prerequisites:
* Python 3.8+
* PyTorch 2.0+
* NVIDIA GPUs with CUDA support
* torch.distributed configured (e.g., via torchrun)

Step 1: Prepare Your FSDP Training Script

Start with a basic FSDP-enabled training script for a small model.

# train_fsdp_poc.py
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
from torch.distributed.fsdp.api import ShardState
import os

# 1. Initialize Distributed Environment
def setup_distributed():
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    master_addr = os.environ.get("MASTER_ADDR", "localhost")
    master_port = os.environ.get("MASTER_PORT", "29500")
    dist.init_process_group("nccl", rank=rank, world_size=world_size,
                            init_method=f"tcp://{master_addr}:{master_port}")
    torch.cuda.set_device(rank)
    print(f"Rank {rank}/{world_size} initialized.")

def cleanup_distributed():
    dist.destroy_process_group()

# 2. Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(10, 20)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(20, 2) # Output 2 classes

    def forward(self, x):
        return self.layer2(self.relu(self.layer1(x)))

# 3. Main training loop
def main_worker():
    setup_distributed()
    rank = dist.get_rank()

    model = SimpleModel().cuda(rank)
    # Wrap model with FSDP
    model = FSDP(model)

    optimizer = optim.SGD(model.parameters(), lr=0.01)

    # Dummy data
    data = torch.randn(64, 10, device=f"cuda:{rank}")
    labels = torch.randint(0, 2, (64,), device=f"cuda:{rank}")

    for epoch in range(5):
        optimizer.zero_grad()
        outputs = model(data)
        loss = nn.CrossEntropyLoss()(outputs, labels)
        loss.backward()
        optimizer.step()
        if rank == 0:
            print(f"Epoch {epoch}: Loss = {loss.item():.4f}")

    cleanup_distributed()

if __name__ == "__main__":
    main_worker()

Step 2: Run the FSDP Script

Execute the script using torchrun (or torch.distributed.launch for older versions). This example assumes 2 GPUs.

torchrun --nproc_per_node=2 train_fsdp_poc.py

Expected Output:
You should see output similar to this (loss values will vary):

Rank 0/2 initialized.
Rank 1/2 initialized.
Epoch 0: Loss = 0.6974
Epoch 1: Loss = 0.6923
Epoch 2: Loss = 0.6874
Epoch 3: Loss = 0.6826
Epoch 4: Loss = 0.6779

The key is that both ranks initialize, and rank 0 prints the loss.

Step 3: Integrate Conceptual Replay Logging

Now, let’s conceptually add a “logging” mechanism. For real deterministic replay, this would intercept CUDA/NCCL calls. Here, we’ll simulate logging a critical tensor’s hash before an all_reduce operation.

# train_fsdp_replay_poc.py (modified from previous)
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
from torch.distributed.fsdp.api import ShardState
import os
import hashlib
import json

# Global log for conceptual replay
replay_log = []

def get_tensor_hash(tensor):
    return hashlib.sha256(tensor.cpu().numpy().tobytes()).hexdigest()

# 1. Initialize Distributed Environment (same as before)
def setup_distributed():
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    master_addr = os.environ.get("MASTER_ADDR", "localhost")
    master_port = os.environ.get("MASTER_PORT", "29500")
    dist.init_process_group("nccl", rank=rank, world_size=world_size,
                            init_method=f"tcp://{master_addr}:{master_port}")
    torch.cuda.set_device(rank)
    print(f"Rank {rank}/{world_size} initialized.")

def cleanup_distributed():
    dist.destroy_process_group()

# 2. Define a simple model (same as before)
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(10, 20)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(20, 2)

    def forward(self, x):
        return self.layer2(self.relu(self.layer1(x)))

# 3. Main training loop with logging
def main_worker_with_replay_logging():
    setup_distributed()
    rank = dist.get_rank()

    model = SimpleModel().cuda(rank)
    model = FSDP(model)

    optimizer = optim.SGD(model.parameters(), lr=0.01)

    data = torch.randn(64, 10, device=f"cuda:{rank}")
    labels = torch.randint(0, 2, (64,), device=f"cuda:{rank}")

    for epoch in range(5):
        optimizer.zero_grad()
        outputs = model(data)
        loss = nn.CrossEntropyLoss()(outputs, labels)
        loss.backward()

        # --- Conceptual Replay Logging ---
        # Before optimizer step, capture gradient info (pre-all_reduce/reduce_scatter)
        # For true replay, you'd intercept NCCL calls directly

        # This is a simplification; FSDP handles gradient reduction internally.
        # Here, we're just conceptually logging a tensor related to model state.
        # In a real scenario, you'd hook into the specific all_reduce calls FSDP makes.

        # Example: Log hash of a parameter's gradient after backward (before FSDP's all_reduce implicitly)
        # This parameter's gradient would be sharded
        if model.parameters().__next__().grad is not None and rank == 0:
            first_param_grad = model.parameters().__next__().grad
            # Note: For FSDP, actual gradients are sharded. This is a simplified example.
            # You'd typically want to log the full unflattened gradient or the input to NCCL.
            replay_log.append({
                "step": epoch,
                "rank": rank,
                "event": "pre_optimizer_step_grad_hash",
                "grad_hash": get_tensor_hash(first_param_grad.float().cpu()), # Ensure deterministic hash
                "grad_shape": list(first_param_grad.shape)
            })

        optimizer.step()
        if rank == 0:
            print(f"Epoch {epoch}: Loss = {loss.item():.4f}")

    # Save the conceptual replay log
    log_filename = f"replay_log_rank_{rank}.json"
    with open(log_filename, 'w') as f:
        json.dump(replay_log, f, indent=4)
    print(f"Rank {rank} saved replay log to {log_filename}")

    cleanup_distributed()

if __name__ == "__main__":
    main_worker_with_replay_logging()

Step 4: Run with Replay Logging and Verify

torchrun --nproc_per_node=2 train_fsdp_replay_poc.py

Expected Output/Verification:
You will see the usual training output. Additionally, replay_log_rank_0.json and replay_log_rank_1.json files will be created. Inspect replay_log_rank_0.json:

[
    {
        "step": 0,
        "rank": 0,
        "event": "pre_optimizer_step_grad_hash",
        "grad_hash": "...", // A SHA256 hash
        "grad_shape": [10]
    },
    // ... more entries for subsequent epochs
]

The hashes should be consistent across multiple runs if all sources of non-determinism (especially random seeds) were perfectly controlled. If you run this script twice and the grad_hash values for the same epoch are different, it indicates non-determinism is still present.

Common Error and Fix:
* Error: “RuntimeError: Address already in use” or torchrun hangs.
* Cause: A previous torchrun process did not clean up properly, or the default MASTER_PORT (29500) is in use.
* Fix:
1. Ensure no lingering python or torchrun processes. Use pgrep -lf python or lsof -i :29500 to find and kill them.
2. Try a different MASTER_PORT environment variable: MASTER_PORT=29501 torchrun ....

Real-World Example

A leading autonomous vehicle company was struggling with intermittent failures in their vision transformer training pipeline, which used FSDP across 128 GPUs. Every few days, a job would silently diverge or produce NaN losses after hours of training. Debugging involved restarting the 12-hour job, adding more print statements, and hoping to catch the bug – a time-consuming, frustrating process.

By integrating a custom deterministic replay framework (built using LD_PRELOAD to intercept CUDA and NCCL calls) they captured a 15-minute trace leading up to a failure. The trace, though large (several terabytes), allowed them to replay the exact conditions of the failure on a smaller, single-node cluster. They discovered a rare race condition in a custom CUDA kernel that processed sparse point cloud data. This race condition led to an uninitialized memory read only when specific GPU task timings aligned, causing an infrequent NaN.

Before: Weeks of speculative debugging, lost GPU hours, delayed model updates.
After: Root cause identified in 3 days, reducing overall debugging time by approximately 90% for similar issues, saving millions in operational costs annually. The deterministic replay capability became a standard part of their MLOps toolkit.

Deterministic Replay vs. Alternatives

Common Pitfalls and Best Practices

Further Learning and Next Steps

Implementing deterministic replay GPU training FSDP is a significant undertaking, but the benefits for debugging complex AI systems are substantial. Start by enhancing your understanding of FSDP’s internal workings and low-level CUDA/NCCL interactions.

1. Deep Dive into FSDP: Explore the official PyTorch FSDP documentation to understand how it shards parameters, gradients, and optimizer states, as well as its communication patterns.
   PyTorch FSDP Documentation
2. Explore NCCL Internals: Understand how NCCL primitives like all_reduce operate at a low level, including their atomic properties and potential for non-determinism under specific conditions.
   NVIDIA NCCL Documentation
3. Study Research Papers on Deterministic Replay: Investigate academic work and open-source projects focused on deterministic replay for distributed systems and machine learning to grasp the state-of-the-art. Look for projects exploring “record and replay” for CUDA/MPI/NCCL.
   Example: “Determinism in Deep Learning: An Empirical Study” – arXiv