Numerical instability in truncated normal initialization with bfloat16

[francesco-bertolotti]

Bug description

I am experiencing numerical instability issues in training with qwen3. I tracked down the issue to the initialization of the model. In particular, to the usage of torch.nn.init.trunc_normal_. The issue is illustrated in the following code snippet. It is apparent that this function does not produce the intended behavior and it is skewed heavily toward the argument of a, in this case -2. I see that this function has an history of being problematic (see pytorch/pytorch#155588 and pytorch/pytorch#145498)

In [20]: import torch
    ...: x = torch.zeros(10,500, dtype=torch.bfloat16)
    ...: torch.nn.init.trunc_normal_(x, mean=0, std=0.02, a=-2, b=2)
    ...: print("min", x.min(-1)[0])
    ...: print("max", x.max(-1)[0])
    ...: print("mean", x.mean(-1)[0])
    ...: print("sum", x.sum(-1)[0])

min tensor([-2.0000, -2.0000, -2.0000, -2.0000, -2.0000, -0.0532, -2.0000, -2.0000, -2.0000, -2.0000], dtype=torch.bfloat16)
max tensor([0.0532, 0.0532, 0.0532, 0.0532, 0.0532, 0.0532, 0.0532, 0.0532, 0.0532, 0.0483], dtype=torch.bfloat16)
mean tensor(-0.0084, dtype=torch.bfloat16)
sum tensor(-4.1875, dtype=torch.bfloat16)

My current workaround is to cast to float32, use truncated_init_, and cast back to bfloat16. This is the related code-snippet

@torch.no_grad()
def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
    orig_dtype = tensor.dtype

    if isinstance(tensor, torch.distributed.tensor.DTensor):
        # Convert DTensor -> float32 DTensor
        tensor_fp32 = tensor.to(dtype=torch.float32)

        # Local init (safe because trunc_normal_ is elementwise)
        torch.nn.init.trunc_normal_(tensor_fp32._local_tensor, mean=mean, std=std, a=a, b=b)

        # Cast back to original dtype and copy into original tensor
        tensor.copy_(tensor_fp32.to(dtype=orig_dtype))
    else:
        # Regular Tensor
        tensor_fp32 = tensor.to(dtype=torch.float32)
        torch.nn.init.trunc_normal_(tensor_fp32, mean=mean, std=std, a=a, b=b)
        tensor.copy_(tensor_fp32.to(dtype=orig_dtype))

    del tensor_fp32

    return tensor

---

Please, let me know if you have any suggestion regarding this issue. I can also submit a PR with the proposed workaround.

Versions

torch: 2.11.0.dev20260107+cu126

[tianyu-l]

@albanD what's your latest take on this issue? Do you think this is not a bug, but expected with bf16?

[tianyu-l]

@francesco-bertolotti Curious: are you training qwen3 from scratch, instead of loading a checkpoint? Would love to hear more about your use case.

[francesco-bertolotti]

Hi @tianyu-l,

Yes, we’re essentially iterating on the qwen3 architecture and pretraining from scratch. I’m trying to ensure the training process runs as smoothly as possible, so I’m also taking a close look at the activations.

[francesco-bertolotti]

Small followup.

I found that the suggested workaround was affected by the same problem. I find this implementation to be more stable, though it could loop indefinitely when a is close to b

@torch.no_grad()
def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
    if isinstance(tensor, torch.distributed.tensor.DTensor):
        tensor = tensor.to_local()

    tensor.normal_(mean, std)
    while tensor.lt(a).any() or tensor.gt(b).any():
        tensor[tensor.lt(a).logical_or(tensor.gt(b))] = tensor[tensor.lt(a).logical_or(tensor.gt(b))].normal_(mean, std)
    
    return tensor

[francesco-bertolotti]

Hi @tianyu-l, I see this has been confirmed as a bug. If the last proposed workaround is an acceptable fix, I’m happy to put together a PR for it. Let me know if you’d like me to proceed!

[albanD]

I found that the suggested workaround was affected by the same problem.

Which workaround? Sampling in fp32 then down-casting?

what's your latest take on this issue? Do you think this is not a bug, but expected with bf16?

It is a pretty subtle one for sure.
Yes you should expect the distribution to be worse when sampling in bf16.
But I also agree that it shouldn't be that bad!
The question is where do we want to put the tradeoff between extra compute and quality of the result.

For torchtitan, if you know you always want good quality sampling. I think upscaling to fp32 to do the sampling is a good fix.
For core, we would need to study the sampling function more closely to try and generate a better distribution given the limited randomness and without a significant slowdown (if this is memory bound, doing fp32 in-kernel might be enough). Someone would have to prepare such a PR and carefully validate the accuracy of it.

[francesco-bertolotti]

Which workaround? Sampling in fp32 and then down-casting?

Yes. Using init.trunc_normal_ in fp32 and casting to bf16 exhibits the same issue---specifically, anomalous sampling at the left bound.

The question is where we want to place the tradeoff between extra compute and result quality.

For PyTorch, it’s probably preferable to provide a fast implementation that is as accurate as possible, with any lower-precision instabilities documented. For TorchTitan, however, I think it makes sense to trade some performance for a more stable initialization.

My reasoning is that this issue surfaced when early-training activations in a Qwen3 model became unusually large. Replacing the unstable function with a stable one immediately reduced activation magnitudes and lowered the loss. Additionally, with standard initialization parameters, we only observed 2-3 resampling events, and only for the largest tensors (the embeddings).

I’ll follow up in the next few days with more thorough documentation of the issue, potentially including plots.