Different batch sizes result in large numerical differences in transformer models

[gilkeren1]

Describe the bug:
I was debugging some issue where inference results are completely diverging, while training looks good. It seems to be stemming from numerical differences due to smaller batch size during inference.

Describe how to reproduce:
I do not have a fully reproducible example, but I see this on my end:

> embeded_maps_top.shape
torch.Size([15, 56, 4096])

> self.llama_decoder_top(
    embeded_maps_top[:1],
    padding_mask=None,
)[0][0, :4]
tensor([[-0.0197,  0.1034,  0.2976,  ...,  0.1380, -0.0656, -0.1645],
        [-0.2249, -0.4262, -0.9900,  ..., -0.2404,  0.7360,  0.1237],
        [ 1.2366,  0.4611, -1.8283,  ..., -0.4055,  0.3350,  1.5457],
        [-0.0480, -0.7974, -0.4915,  ..., -0.1564, -0.2716,  0.2660]],
       device='cuda:0', grad_fn=<SliceBackward0>)

> self.llama_decoder_top(
     embeded_maps_top,
     padding_mask=None,
)[0][0, :4]
tensor([[-0.0190,  0.1034,  0.2991,  ...,  0.1373, -0.0656, -0.1645],
        [-0.2310, -0.4302, -0.9899,  ..., -0.2414,  0.7400,  0.1247],
        [ 1.2365,  0.4610, -1.8141,  ..., -0.4020,  0.3368,  1.5527],
        [-0.0459, -0.7936, -0.4916,  ..., -0.1583, -0.2678,  0.2602]],
       device='cuda:0', grad_fn=<SliceBackward0>)

note that the difference between the two calls is the batch size, and the results are different up to 1e-2 is some places. compounded over multiple calls during autoregressive inference, this lead to a complete divergent results.

Describe the expected behavior:
see above.

Environment:
fs2 branch that is off v0.4.6 (the mms one)
torch 2.4.0+cu121
single A100 gpu on a devserver.

Additional Context:

[gilkeren1]

More context about the model:

self.llama_decoder_top
StandardTransformerLMDecoder(
model_dim=4096, self_attn_mask_factory=CausalAttentionMaskFactory(), norm_order=PRE
(layers): ModuleList(
(0-11): 12 x StandardTransformerLMDecoderLayer(
model_dim=4096, norm_order=PRE
(self_attn_layer_norm): RMSNorm(normalized_shape=(4096,), eps=1E-05, elementwise_affine=True)
(self_attn): StandardMultiheadAttention(
num_heads=8, model_dim=4096
(q_proj): Linear(input_dim=4096, output_dim=4096, bias=False, init_fn=init_projection)
(k_proj): Linear(input_dim=4096, output_dim=4096, bias=False, init_fn=init_projection)
(v_proj): Linear(input_dim=4096, output_dim=4096, bias=False, init_fn=init_projection)
(pos_encoder): RotaryEncoder(encoding_dim=512, max_seq_len=512)
(sdpa): TorchSDPA(attn_dropout_p=0)
(output_proj): Linear(input_dim=4096, output_dim=4096, bias=False, init_fn=init_projection)
)
(self_attn_dropout): None
(self_attn_residual): StandardResidualConnect()
(ffn_layer_norm): RMSNorm(normalized_shape=(4096,), eps=1E-05, elementwise_affine=True)
(ffn): GLUFeedForwardNetwork(
model_dim=4096, inner_dim_scale=0.666667, inner_dim_to_multiple=256
(gate_proj): Linear(input_dim=4096, output_dim=2816, bias=False, init_fn=init_projection)
(gate_activation): SiLU()
(inner_proj): Linear(input_dim=4096, output_dim=2816, bias=False, init_fn=init_projection)
(inner_dropout): None
(output_proj): Linear(input_dim=2816, output_dim=4096, bias=False, init_fn=init_projection)
)
(ffn_dropout): None
(ffn_residual): StandardResidualConnect()
)
)
(layer_norm): RMSNorm(normalized_shape=(4096,), eps=1E-05, elementwise_affine=True)
(dropout): None
)

[cirquit]

Responded offline. I except this to be an issue with the TorchSDPA backend which produces slightly different attention outputs given different sizes inputs (batch size modulation).

[gilkeren1]

Thanks! This could be a real silent killer for models. I'm trying on my end to make sure activations are in the 1e-1 scale and not smaller, so that those numeric issue won't make a huge difference. Not sure if it's possible to solve it on the fs2 side to make sure a consistent kernel is used. I'm surprised again that different kernels make such a big difference!

[cirquit]

I think the general consensus is that these differences in numerical reproducibility are mitigated through dropout and other regularization techniques (mostly softmax afaik), such that the final output does not specifically depend on them. Basically, a non-overfit model should not have different outputs with B=1 vs B=32, while internally not having complete bit-wise logit parity.

This type of numerical issues are also expected from PyTorch's perspective: https://docs.pytorch.org/docs/stable/notes/numerical_accuracy.html#batched-computations-or-slice-computations.