Bugs that broke driving: Machine Learning edition

“We have solved self driving, we just need to fix the bugs.”

Rumor has it that the Linux kernel averages around 0.5 bugs per 1,000 lines of code. Some bugs are easy to avoid or catch with tests and metrics. Others are very good at hiding.

This is a sequel to Bugs that broke driving with more Machine Learning related bugs.

Some of these bugs affected correctness, others slowed things down or used up extra memory. They all made training or inference worse.

Looking back at these bugs helps us understand where our code tends to break and how to catch problems earlier.

Output layers need to be run in high precisionPermalink

September 2026.

The world model is a diffusion transformer (DiT) which predicts the next frame and an “expert” driving plan.

During inference, the model runs in mixed precision bf16_fp8_nvfp4:

• Transformer MLP linears use NVFP4 (most FLOPs)
• Attention projections, cached attention Q/K/V, and the KV cache use FP8 (saves VRAM)
• Other layers use BF16 (not too many FLOPs)
• Norm reductions, attention scores, and sinusoidal timestep features use FP32 (standard safe mixed precision)

In particular, the plan output layers ran in BF16. Around 30 m/s, BF16 can only represent speeds in steps of 0.125 m/s. The acceleration calculation from the plans subtracts nearby speed predictions and divides by a short time interval $\Delta t$, magnifying the rounding errors:

The fix is to promote the plan head and its inputs to FP32 before running it, Casting the finished plan to FP32 just gives the staircase more decimal places.

+plan_head.float()
-plan = plan_head(features).float()
+plan = plan_head(features.float())

ConvNeXt FP16Permalink

August 2026.

The driving model’s vision encoder is a ConvNeXt, with a learned multiplier on each residual branch: x + gamma * f(x). At runtime, it runs in FP16.

The pretrained checkpoint we used we pick had a value of almost 54 at stages[2].blocks[29]. After fine-tuning, some inputs push the residual past FP16’s 65,504 limit leading to inf.

The annoying thing is that this only triggered in some seemingly random cases, and didn’t happen in the various checks we run during model exporting to FP16.

The workaround was to divide that stage’s residual stream by four and carry the scale through its biases, multipliers, and LayerNorm epsilons:

 stage = encoder.stages[2]
+stage.downsample[1].weight.div_(4)
+stage.downsample[1].bias.div_(4)
+for block in stage.blocks:
+    block.conv_dw.bias.div_(4)
+    block.gamma.div_(4)
+    block.norm.eps /= 16
+encoder.stages[3].downsample[0].eps /= 16

Missing desire targetsPermalink

July 2026.

The driving model predicts “driving desire”, like turning or changing lanes.

For a while, some training data had no desire targets: we use NaNs to distinguish missing targets, and our loss helpers normally handle this by masking those targets out. But the desire loss took an argmax first, and did not have the NaN handling.

>>> import torch
>>> import torch.nn.functional as F
>>> target = torch.full((1, 8), torch.nan)
>>> target.argmax(-1)
tensor([0])
>>> F.cross_entropy(torch.zeros(1, 8), target.argmax(-1))
tensor(2.0794)

All those NaNs became class 0, “no desire”, with a perfectly normal loss.

The fix replaces NaNs with zeros while keeping a mask, then uses that mask to ignore missing targets:

def num_from_nan(target):
    valid = ~torch.isnan(target)
    return target.masked_fill(~valid, 0).detach(), valid

+target, valid = num_from_nan(target)
 loss = F.cross_entropy(logits, target.argmax(-1), reduction="none")
+loss = valid.all(-1) * loss

Redundant biases drifted apart across distributed ranksPermalink

April 2026.

The driving model has many nn.Linear layers. During training, we use torch.compile, BF16 autocast and Distributed Data Parallel (DDP) as “acceleration” recipes.

Sometimes, harmless redundant biases are written in the architecture:

a = nn.Linear(8, 8192)
b = nn.Linear(4, 8192)
y = a(x) + b(x[:, :4])

Note that the 8/4 is just illustrative, and keeps the weight gradients different while the bias gradients are identical.

a.bias and b.bias will be merged by the compiler into one bias, so they get the same gradient. The compiler went one step further: it made both .grad tensors share the same memory.

DDP synchronized them correctly, then clip_grad_norm_ used its foreach implementation to scale that shared buffer from two CUDA blocks at once.

Depending on timing, an entry could be multiplied by the clipping coefficient once or twice, giving each rank a different update.

In one recorded run, rank 0 clipped all 8,192 entries once; rank 1 clipped 384 of them twice. The drift can reach a magnitude of 0.0056 between ranks after only 50 SGD steps, while the weight matrices stayed identical.

Both biases receive the same faulty update within a rank, their sum drifts too, which can change the model’s outputs. DDP synchronizes gradients during training, so it doesn’t repair those parameter differences: the GPUs are now training different copies of the model.

We worked around it by removing the redundant biases.

This was fixed in PyTorch 2.13.0 as a side effect of a compiler optimization change.

It preserved FP32 → BF16 → FP32 cast chains so smaller intermediate tensors could still be materialized: in our backward graph, that happened to keep the two bias gradients in separate FP32 buffers.

Note: The foreach kernel itself still assumes the gradient buffers don’t overlap…

DDP random seed bugPermalink

August 2026.

Torchtitan uses the same seed across DDP (replicate), this is correct because weight initialization should match across replicate ranks. But torchtitan then keeps that rank identical across during for training.

With ordinary tensors (not DTensors), different ranks could get identical dropout masks and diffusion noise.

The fix is to reseed each data-parallel rank after initialization (essentially):

 torch.manual_seed(seed)
 model.init_weights()
+torch.manual_seed(seed + dp_rank)
 train(model)

Flux does this after initialization, setting distinct_seed_mesh_dims=["cp", "dp_shard", "dp_replicate"].

AllNormPermalink

June 2021 to April 2026.

Normalization layers are widely used in Machine Learning. BatchNorm normalizes each channel using the mean and variance across the batch and spatial dimensions. During training, it uses the current batch statistics and updates running averages. During inference, it uses those running averages.

The driving models are multi-task learners, with some tasks being more sparse than others, for example: hard brake events.

One failure mode we encountered was when a channel specializes in a rare feature. Through batches without that feature, the channel stays quiet and its BatchNorm running variance shrinks.

When validation activates that channel again, it gets divided by a stored standard deviation that is too small, blowing up the outputs. Training can still look fine because its statistics come from the current batch: when the feature returns, the denominator grows with it.

LayerNorm fixes this, but needs per input statistics at inference. On device, we want normalization folded into convolutions. Recalibrating BatchNorm running statistics offline also fixes this, but adds overhead and another training pass. Same for lowering momentum which needed extra tuning.

The best fix we came up with was pooling the whole batch, channels, and spatial dimensions into one set of statistics. We called this AllNorm.

class AllNorm(nn.Module):
    def __init__(self):
        super().__init__()
        self.bn = nn.BatchNorm1d(1)

    def forward(self, x):
        return self.bn(x.reshape(-1, 1)).reshape_as(x)

All channels share statistics, making it robust to this failure, and keeping the ability to fold the statistics into the convolutions.

Note: With the extra compute that the chestnut gives, we can now afford using LayerNorm and not deal with this!

Gigashuffle’s reusable CPU buffersPermalink

July 2026.

Gigashuffle returns tensors backed by reusable CPU buffers. Each reader allocates these once and fills them again for every batch:

# Once per reader:
reader_buffer[i][key] = torch.empty(shape, dtype=dtype).share_memory_()

# For every batch:
torch.index_select(shuffle_buffer[i][key], 0, indices,
                   out=reader_buffer[i][key])

Normally, each batch is copied to the GPU before the next read. Our gradient accumulation loop collected several microbatches first, so later reads overwrote batches already in the list.

One way to fix this is to clone both inputs and targets before fetching the next microbatch:

 for _ in range(grad_accum):
     inputs, targets = next(loader)
+    inputs = {k: v.clone() for k, v in inputs.items()}
+    targets = {k: v.clone() for k, v in targets.items()}
     microbatches.append((inputs, targets))

FSDP2 init bugsPermalink

April 2025.

With FSDP2, we construct the model on meta, shard it, then use to_empty() to allocate the weights. Initialization has to run after allocation: not defining an init function for a parameter initializes it to the arbitrary memory value it received during allocation.

>>> import torch
>>> with torch.device("meta"):
...     model = torch.nn.Module()
...     model.weight = torch.nn.Parameter(torch.ones(4))
>>> _ = model.to_empty(device="cpu")
>>> model.weight.detach()
tensor([ 3.7189e-38,  0.0000e+00, -6.3894e-37,  4.5160e-41])

No clear fix here but to be paranoid about not forgetting init methods!

NCCL timeout increase didn’t do anythingPermalink

July 2026.

Our RL training runs collect rollouts from the cluster. If the cluster is busy, some ranks can wait for data longer than PyTorch’s default NCCL timeout of 10 minutes.

To work around this, we added a barrier with a timeout of one hour before each training step, thinking we have fixed it. But that timeout was passed around, only to never be used:

opts = BarrierOptions()
opts.timeout = timeout
opts.asyncOp = async_op
work = group.barrier(opts=opts)

The C++ dispatcher passed it along, but NCCL implemented the barrier as an all-reduce and only forwarded asyncOp:

// C++ dispatcher
BarrierOptions opts;
opts.timeout = std::chrono::milliseconds(timeout);
opts.asyncOp = asyncOp;
backend->barrier(opts);

// ProcessGroupNCCL::barrier
AllreduceOptions arOpts;
arOpts.asyncOp = opts.asyncOp;
auto work = allreduce_impl(barrierTensor, "nccl:all_reduce_barrier", arOpts);
// opts.timeout goes nowhere.

That all-reduce got its deadline from a different options object: the process group’s configuration. The watchdog checked that deadline; our barrier(timeout=...) argument never entered the calculation:

// Creating the collective's work: options_ belongs to the process group.
assignTimeoutToWork(work, options_);

// Inside assignTimeoutToWork(work, option), omitting timeout extensions:
work->opTimeout_ = option->timeout;

// Watchdog calls work.checkTimeout() with no timeout override:
if (timeElapsed >= opTimeout_) {
  // Report a collective timeout.
}

The fix was to give the wait for data its own Gloo group, with the timeout set on that group:

+# Once at startup, on every rank:
+ready = dist.new_group(backend="gloo", timeout=timedelta(hours=1))

 # Before each training step:
-dist.barrier(timeout=timedelta(hours=1))
+dist.barrier(group=ready)

Memory leak with AdamW’s foreach updates on DTensorsPermalink

March 2026.

AdamW batches updates into calls like torch._foreach_mul_(params, 1 - lr * weight_decay). FSDP2 exposes its parameters as DTensors, so those calls go through DTensor’s dispatcher to work out how the operation applies to the shards. For foreach’s tensor lists, the dispatcher built a fresh OpSchema on every call, then constructed a cache key from the tensors’ shape and sharding metadata.

In the C++ key builder, PyTuple_Pack increments its arguments’ reference counts, but .release() stopped the C++ wrapper from decrementing its own reference afterward. That left an extra reference to the argument tuple on every call, so the tuples and their metadata kept piling up in CPU RAM.

The fix was to simply not release:

 // torch/csrc/autograd/python_variable.cpp
 // Inside DTensor_OpSchema_recompute_comparison_key_impl (simplified):
 py::tuple args_to_hash_tup(args_to_hash.size());
 // ... fill the tuple with argument metadata ...
 comparison_key = PyTuple_Pack(
     2, op,
-    args_to_hash_tup.release().ptr());
+    args_to_hash_tup.ptr());

The bug affected PyTorch 2.10.0 through 2.12.1; the same fix shipped via a later PR in 2.13.0.

CUDA cache allocationPermalink

February 2026.

When doing inference in our distributed cluster, we use an inference server with dynamic batching: triton inference server. We usually benchmark at the largest batch size in isolation, and assume this will be the max VRAM the model will use. But dynamic batching reserved more VRAM, and we didn’t know why.

We traced it to how PyTorch caches freed GPU allocations, so the order of batch sizes matters. With its default native allocator, this takes just one tensor at a time:

import torch

def reserved_after(batch_sizes):
    torch.cuda.empty_cache()  # Reset between experiments.
    for bs in batch_sizes:
        x = torch.empty(bs, 1024, 1024, device="cuda", dtype=torch.float32)
        del x
    return torch.cuda.memory_reserved() // 2**20

print(reserved_after([3, 4]))     # 28 MiB
print(reserved_after([4, 3, 4]))  # 16 MiB

Batch size 3 leaves a cached 12 MiB block, batch size 4 needs 16 MiB, so the allocator keeps both. Starting with 4 lets batch 3 use 12 MiB of the larger 16 MiB block.

We flipped the warmup order to start with the largest batch:

-for bs in range(1, max_batch_size + 1):
+for bs in range(max_batch_size, 0, -1):
     warmup(bs)

CUDA graphs don’t capture Adam’s CPU workPermalink

June 2022.

When training with CUDA graphs, we used to capture the whole training step, optimizer step included. CUDA graphs replay only GPU work, so CPU code inside the capture runs just once!

By default, PyTorch’s eager Adam increments its step counter and computes its bias corrections, 1 - beta1**step and 1 - beta2**step, on the CPU:

# torch/optim/adam.py — Adam._init_group() (simplified)
state = self.state[p]
if len(state) == 0:  # Initialize this parameter's optimizer state once.
    state["step"] = (
        torch.zeros((), device=p.device)
        if group["capturable"] or group["fused"]
        else torch.tensor(0.0, device="cpu") # cpu?!?!
    )

Back then, replay kept updating the GPU moment estimates and weights with those frozen corrections, silently reusing the value of step from the capture-time value. We avoided this by running the optimizer outside the graph.

PyTorch 1.12.0 later added Adam(..., capturable=True), putting the step counter and bias-correction math on the GPU so replay advances them too. Today, manual capture with capturable=False raises an error.

Looking backPermalink

The recurring themes were precision, leaked or reused memory, and distributed training. Some bugs were hard to spot but preventable with better code practices. Others became obvious once we visualized model outputs during inference or inspected model state during training.

Come fix the bugs that are still hiding (https://comma.ai/jobs).

Yassine.