# Different output when we inference through packing with flash attention in bf16

**URL:** <https://discuss.huggingface.co/t/different-output-when-we-inference-through-packing-with-flash-attention-in-bf16/174030>\
**Category:** 🤗Transformers\
**Created:** [March 5, 2026, 6:20pm UTC](https://discuss.huggingface.co/t/different-output-when-we-inference-through-packing-with-flash-attention-in-bf16/174030 "2026-03-05T18:20:55Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Jaswanth-0821](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/jaswanth-0821/32/55721_2.png) [@Jaswanth-0821](https://discuss.huggingface.co/u/Jaswanth-0821)\
**Post date:** [March 5, 2026, 6:20pm UTC](https://discuss.huggingface.co/t/different-output-when-we-inference-through-packing-with-flash-attention-in-bf16/174030/1 "2026-03-05T18:20:55Z")

</div>

i am trying to pack the tokens and see if the output remain same with and without packing

when i used f16 precision results for with packing and without packing remain same, but when i tried with bf16 i found slight difference in the output

Below is the code

```python
import torch
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B', padding_side='left')
model = AutoModel.from_pretrained("Qwen/Qwen3-Embedding-0.6B" , attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16).cuda()

# ignore the example text :)
texts = ["hi how are you this is what so that is that what is know is there any way to do this" , " hi how can i help you today"]
input_ids = tokenizer(
    texts,
    padding=True,
    truncation=True,
    max_length=512,
    return_tensors="pt",
)
input_ids = { k: v.cuda() for k, v in input_ids.items() }

# packing the tokens
input_ids_with_packing = tokenizer(
    texts,
    padding=False,
    truncation=False,
    max_length=512,
)
input_with_packing = {"input_ids" : [], "position_ids" : []}
index =1
for i in range(len(input_ids_with_packing["input_ids"])):
    input__ = input_ids_with_packing["input_ids"][i]
    input_with_packing['input_ids'] += input__

    index = index+1
    input_with_packing['position_ids'] += list(range(len(input_ids_with_packing["input_ids"][i])))
    

input_with_packing = {k: torch.tensor([v]).cuda() for k, v in input_with_packing.items()}

```

so the final inputs will be like

```python
# without packing
{'input_ids': tensor([[ 6023, 1246, 525, 498, 419, 374, 1128, 773, 429,
             374, 429, 1128, 374, 1414, 374, 1052, 894, 1616,
             311, 653, 419, 151643],
         [151643, 151643, 151643, 151643, 151643, 151643, 151643, 151643, 151643,
          151643, 151643, 151643, 151643, 151643, 15588, 1246, 646, 600,
            1492, 498, 3351, 151643]], device='cuda:0'),
 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]],
        device='cuda:0')}

# with packing
{'input_ids': tensor([[ 6023, 1246, 525, 498, 419, 374, 1128, 773, 429,
             374, 429, 1128, 374, 1414, 374, 1052, 894, 1616,
             311, 653, 419, 151643, 15588, 1246, 646, 600, 1492,
             498, 3351, 151643]], device='cuda:0'),
 'position_ids': tensor([[ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17,
          18, 19, 20, 21, 0, 1, 2, 3, 4, 5, 6, 7]], device='cuda:0')}

```

when i run with these inputs

```python
with torch.no_grad():
    direct_outputs = model(**input_ids)
    packing_outputs = model(**input_with_packing)

```

the last\_hidden\_state outputs are

```auto
# without packing
 tensor([[[ 2.4062e+00, -1.4609e+00, -1.8848e-01, ..., -8.4375e+00,
           -1.1625e+01, 1.3504e-03],
          [-3.9375e+00, -3.2656e+00, -1.2500e+00, ..., 1.5078e+00,
           -1.7090e-01, 1.0000e+00],
          [-2.4219e+00, -5.9688e+00, -1.4766e+00, ..., 8.4375e-01,
           -1.0156e+00, -2.3281e+00],
          ...,
          [-1.6406e+00, -8.6875e+00, -1.3281e+00, ..., -5.7812e-01,
           -1.4844e+00, -1.5625e-01],
          [-1.7344e+00, -6.6875e+00, -1.3594e+00, ..., -3.1982e-02,
           -5.2344e-01, -1.7480e-01],
          [-9.8047e-01, -1.7188e+00, -9.6484e-01, ..., -9.9121e-02,
            2.2344e+00, 4.6250e+00]],
 
         [[-4.3438e+00, 1.6797e-01, 2.5586e-01, ..., 3.7969e+00,
            6.6250e+00, -3.9688e+00],
          [-4.3438e+00, 1.6797e-01, 2.5586e-01, ..., 3.7969e+00,
            6.6250e+00, -3.9688e+00],
          [-4.3438e+00, 1.6797e-01, 2.5586e-01, ..., 3.7969e+00,
            6.6250e+00, -3.9688e+00],
          ...,
          [ 5.7031e-01, -5.2188e+00, -1.1094e+00, ..., 1.0000e+00,
           -4.0938e+00, -1.3984e+00],
          [-2.5000e+00, -3.6719e+00, -1.2812e+00, ..., 1.1328e+00,
           -4.5508e-01, -2.4375e+00],
          [-1.6016e+00, 5.1953e-01, -1.1016e+00, ..., 9.2188e-01,
            5.0000e-01, 1.7266e+00]]], device='cuda:0', dtype=torch.bfloat16))

# with packing
(tensor([[[2.5312, -1.4531, -0.1982, ..., -8.4375, -11.7500, -0.0121],
          [-3.9688, -3.1562, -1.2656, ..., 1.5703, -0.1865, 0.8359],
          [-2.3906, -5.9062, -1.4766, ..., 0.9258, -1.0234, -2.2188],
          ...,
          [0.5078, -5.1250, -1.1250, ..., 1.0938, -4.0625, -1.3594],
          [-2.5625, -3.5938, -1.2656, ..., 1.1484, -0.4316, -2.4219],
          [-1.7188, 0.5547, -1.0938, ..., 0.9609, 0.5977, 1.7031]]],
        device='cuda:0', dtype=torch.bfloat16),

```

is there any reason for this ? or am i doing something wrong here

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [March 6, 2026, 4:26am UTC](https://discuss.huggingface.co/t/different-output-when-we-inference-through-packing-with-flash-attention-in-bf16/174030/2 "2026-03-06T04:26:26Z")

</div>

> or am i doing something wrong here

Seems nothing wrong. Well, it’s common for differences in float precision or attention backends to slightly alter the output. [Occasionally, there are also cases stemming from actual bugs](https://github.com/EleutherAI/gpt-neox/issues/1003).

* * *

Yes. There is a reason, and in your case it is **most likely expected behavior** , not a serious mistake.

## Main answer

Your two forwards are **not the same low-level computation** , even if they represent the same logical text inputs:

- **without packing** : padded batch, shape roughly `[2, 22]`, with an `attention_mask`
- **with packing** : flattened batch, shape `[1, 30]`, with reset `position_ids` and **no** `attention_mask`

For recent Hugging Face `transformers`, that packed form is the intended padding-free style for FlashAttention-based packing: Qwen3 forwards `attention_mask` and `position_ids` into `create_causal_mask(...)`, and the packed-sequence path is triggered from `position_ids` when `attention_mask is None`. TRL’s padding-free docs and code describe the same pattern: flatten sequences and return `position_ids` instead of `attention_mask`. ([GitHub](https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3/modeling_qwen3.py "transformers/src/transformers/models/qwen3/modeling\_qwen3.py at main · huggingface/transformers · GitHub"))

So the packed input structure you built is **conceptually valid** for a modern stack.

* * *

## Why BF16 can differ while FP16 looked the same

There are two separate effects here.

### 1) Different tensor layout means different accumulation order

PyTorch explicitly says that mathematically equivalent batched computations are **not guaranteed to be bitwise identical** to slice-by-slice or differently-shaped computations. It also says that applying an operation to a slice is not guaranteed to match slicing the result of the full operation. Your padded run and packed run have different shapes and different masking/layout behavior, so this warning applies directly. ([PyTorch Documentation](https://docs.pytorch.org/docs/stable/notes/numerical_accuracy.html "Numerical accuracy — PyTorch 2.10 documentation"))

### 2) Fused attention backends can change the result slightly

PyTorch’s SDPA docs state that because floating-point operations are fused, the output **may differ depending on which backend kernel is chosen**. The docs also note that the math backend keeps intermediates in `torch.float` for `torch.half` and `torch.bfloat16` inputs, which is one reason fused FA-style paths and higher-precision reference paths do not match exactly. ([PyTorch Documentation](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html "torch.nn.functional.scaled\_dot\_product\_attention — PyTorch 2.10 documentation"))

BF16 also has a PyTorch-documented reduced-precision GEMM reduction mode that is enabled by default and can be disabled if you see unwanted numerical effects. ([PyTorch Documentation](https://docs.pytorch.org/docs/stable/backends.html "torch.backends — PyTorch 2.10 documentation"))

So “FP16 looked identical to me, BF16 shows a small drift” is plausible. It does **not** automatically mean the BF16 packed path is wrong. ([PyTorch Documentation](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html "torch.nn.functional.scaled\_dot\_product\_attention — PyTorch 2.10 documentation"))

* * *

## There is a very similar upstream report

A close public analogue exists in the FlashAttention repo: a user compared `flash_attn_func` and `flash_attn_varlen_func` in a BF16 setup and found **small numerical differences** when strict equality was expected. That is very close to your situation, because “normal path” versus “varlen / packed path” is essentially the same class of comparison. ([GitHub](https://github.com/Dao-AILab/flash-attention/issues/1772 "Slight numerical discrepancy between flash\_attn\_func and flash\_attn\_varlen\_func · Issue #1772 · Dao-AILab/flash-attention · GitHub"))

There is also a Transformers issue reporting that `flash_attention_2` is not deterministic in the same way as `sdpa` or `eager`, which further supports the idea that strict equality is the wrong expectation once FA2 is involved. ([GitHub](https://github.com/huggingface/transformers/issues/31787 "Transformer models are not deterministic when using Flash Attention 2 · Issue #31787 · huggingface/transformers · GitHub"))

* * *

## What you are doing right

### Left padding

Qwen’s own model card recommends using `padding_side="left"` together with `flash_attention_2` for this model family. ([Hugging Face](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B "Qwen/Qwen3-Embedding-0.6B · Hugging Face"))

### Packing via flattened `input_ids` + reset `position_ids`

That matches the HF packing design for models that use `position_ids`. The packing blog and TRL docs describe padding-free packing exactly in those terms. ([Hugging Face](https://huggingface.co/blog/packing-with-FA2 "Improving Hugging Face Training Efficiency Through Packing with Flash Attention 2"))

### Not passing `attention_mask` in the packed case

That is aligned with TRL’s padding-free implementation, which returns `position_ids` instead of `attention_mask`. ([GitHub](https://github.com/huggingface/trl/blob/main/trl/trainer/sft_trainer.py "trl/trl/trainer/sft\_trainer.py at main · huggingface/trl · GitHub"))

* * *

## What is weak in the current comparison

## 1) You are comparing raw `last_hidden_state`

For an embedding model, that is the **most sensitive** thing you can compare. Qwen’s model card shows the intended inference path as:

1. run the model,
2. take `last_hidden_state`,
3. do **last-token pooling** using the attention mask,
4. normalize the embedding. ([Hugging Face](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B "Qwen/Qwen3-Embedding-0.6B · Hugging Face"))

That means the more meaningful comparison is the **final pooled normalized embedding** , not raw hidden states at every token position.

## 2) You need to compare the correct slices

Your padded batch and packed batch are laid out differently.

For your example:

- sequence 1 in the padded batch is `direct_outputs.last_hidden_state[0, :len1]`
- sequence 2 in the padded batch is `direct_outputs.last_hidden_state[1, -len2:]` because of left padding
- sequence 1 in the packed batch is `packing_outputs.last_hidden_state[0, :len1]`
- sequence 2 in the packed batch is `packing_outputs.last_hidden_state[0, len1:len1+len2]`

If you compare whole tensors by eye, you are mixing “different positions in different layouts,” which makes the discrepancy look worse than it is. This is a comparison problem, not necessarily a model problem.

* * *

## My judgment for your exact snippet

I would rank the explanations like this:

### Most likely

**Expected BF16 + FlashAttention2 numeric drift** due to different shapes, different masking/layout, and fused-kernel accumulation order. ([PyTorch Documentation](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html "torch.nn.functional.scaled\_dot\_product\_attention — PyTorch 2.10 documentation"))

### Also likely

**Your comparison method overstates the difference** , because you are visually comparing raw `last_hidden_state` rather than aligned token slices and final pooled embeddings. ([Hugging Face](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B "Qwen/Qwen3-Embedding-0.6B · Hugging Face"))

### Less likely, but worth checking

**Version-sensitive packing behavior**. This code path has been under active change, and there are public issues around packed-mask handling and `position_ids`-based detection in `transformers`. ([GitHub](https://github.com/huggingface/transformers/issues/41856 "`allow\_is\_causal\_skip` incorrectly disabled when `use\_cache= ..."))

* * *

## What I would change first

### 1) Put the model in eval mode

Do this before comparing anything:

```python
model.eval()

```

That removes training-time randomness as a confounder.

### 2) Compare aligned token slices, not whole tensors

```python
with torch.no_grad():
    direct = model(**input_ids).last_hidden_state
    packed = model(**input_with_packing).last_hidden_state

len1 = len(input_ids_with_packing["input_ids"][0])
len2 = len(input_ids_with_packing["input_ids"][1])

direct_1 = direct[0, :len1]
packed_1 = packed[0, :len1]

direct_2 = direct[1, -len2:]
packed_2 = packed[0, len1:len1+len2]

print("seq1 max abs diff:", (direct_1.float() - packed_1.float()).abs().max().item())
print("seq2 max abs diff:", (direct_2.float() - packed_2.float()).abs().max().item())

print("seq1 allclose:", torch.allclose(direct_1.float(), packed_1.float(), atol=1e-2, rtol=1e-2))
print("seq2 allclose:", torch.allclose(direct_2.float(), packed_2.float(), atol=1e-2, rtol=1e-2))

```

### 3) Compare the actual embeddings

For Qwen3-Embedding, compare the pooled embeddings, not just hidden states. The model card uses last-token pooling for left-padded inputs and then L2 normalization. ([Hugging Face](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B "Qwen/Qwen3-Embedding-0.6B · Hugging Face"))

For your packed tensor, the pooled positions are:

- seq1 pooled token: `len1 - 1`
- seq2 pooled token: `len1 + len2 - 1`

Then normalize and compare cosine similarity.

### 4) Build a stricter reference baseline

If you want to know whether the problem is semantic or only numeric, compare against a higher-precision non-fused reference:

- `torch_dtype=torch.float32`
- `attn_implementation="eager"` or a math-style reference path

PyTorch’s SDPA docs explicitly say the math backend is the higher-precision fallback and keeps float intermediates for half/BF16 inputs. ([PyTorch Documentation](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html "torch.nn.functional.scaled\_dot\_product\_attention — PyTorch 2.10 documentation"))

If packed and non-packed match there but drift slightly only in BF16 + FA2, that strongly points to “expected numeric difference.”

### 5) Try disabling BF16 reduced-precision reductions

PyTorch documents this flag for BF16 GEMMs:

```python
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

```

If the discrepancy shrinks, that is further evidence that your issue is numeric, not semantic. ([PyTorch Documentation](https://docs.pytorch.org/docs/stable/backends.html "torch.backends — PyTorch 2.10 documentation"))

* * *

## Practical conclusion

For **your case** , I would interpret the result like this:

- **No** , your code is not obviously wrong just because BF16 packed output is slightly different.
- **Yes** , there is a solid reason: different layout + packed masking path + FA2 fused kernels + BF16 reduction behavior. ([GitHub](https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3/modeling_qwen3.py "transformers/src/transformers/models/qwen3/modeling\_qwen3.py at main · huggingface/transformers · GitHub"))
- The right standard is **close within tolerance** , especially on the final pooled normalized embeddings, not exact equality of raw hidden states. ([Hugging Face](https://huggingface.co/Qwen/Qwen3-Embedding-0.6B "Qwen/Qwen3-Embedding-0.6B · Hugging Face"))

The result would become concerning if:

- the second packed segment is much worse than the first,
- pooled embeddings differ materially,
- or the mismatch stays large even in an FP32 / eager reference run.

* * *

```python
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

MODEL_NAME = "Qwen/Qwen3-Embedding-0.6B"
TEXTS = [
    "hi how are you this is what so that is that what is know is there any way to do this",
    "hi how can i help you today",
]
MAX_LENGTH = 512
DEVICE = "cuda"

# ----------------------------
# helpers
# ----------------------------
def last_token_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    # Qwen3-Embedding uses left padding with FA2, so the last token is the rightmost valid token.
    lengths = attention_mask.sum(dim=1) - 1
    batch_idx = torch.arange(last_hidden_states.size(0), device=last_hidden_states.device)
    return last_hidden_states[batch_idx, lengths]

def build_packed_inputs(tokenizer, texts, max_length, device):
    # Important: use the SAME truncation settings as the padded path.
    unpadded = tokenizer(
        texts,
        padding=False,
        truncation=True,
        max_length=max_length,
        return_attention_mask=False,
    )

    flat_ids = []
    flat_pos = []
    lengths = []

    for ids in unpadded["input_ids"]:
        flat_ids.extend(ids)
        flat_pos.extend(range(len(ids)))
        lengths.append(len(ids))

    packed = {
        "input_ids": torch.tensor([flat_ids], device=device),
        "position_ids": torch.tensor([flat_pos], device=device),
    }
    return packed, lengths

def compare_case(model_name, texts, dtype, attn_impl, max_length=512):
    print(f"\n=== dtype={dtype}, attn_impl={attn_impl} ===")

    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
    model = AutoModel.from_pretrained(
        model_name,
        attn_implementation=attn_impl,
        torch_dtype=dtype,
    ).to(DEVICE)
    model.eval()

    # Padded inputs
    padded = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=max_length,
        return_tensors="pt",
    )
    padded = {k: v.to(DEVICE) for k, v in padded.items()}

    # Packed inputs
    packed, packed_lengths = build_packed_inputs(tokenizer, texts, max_length, DEVICE)

    with torch.inference_mode():
        out_padded = model(**padded).last_hidden_state
        out_packed = model(**packed).last_hidden_state

    # Compare per-sequence hidden states
    lengths = padded["attention_mask"].sum(dim=1).tolist()
    start = 0
    pooled_packed = []

    print("\nPer-sequence hidden-state comparison")
    for i, L in enumerate(lengths):
        # left padding => valid tokens are the last L positions
        hs_padded = out_padded[i, -L:]
        hs_packed = out_packed[0, start:start + L]
        start += L

        diff = (hs_padded.float() - hs_packed.float()).abs()
        max_abs = diff.max().item()
        mean_abs = diff.mean().item()
        ok = torch.allclose(hs_padded.float(), hs_packed.float(), atol=1e-2, rtol=1e-2)

        print(f"seq{i}: len={L:3d} | max_abs={max_abs:.6f} | mean_abs={mean_abs:.6f} | allclose={ok}")

        pooled_packed.append(hs_packed[-1])

    # Compare final pooled normalized embeddings
    pooled_padded = last_token_pool(out_padded, padded["attention_mask"])
    pooled_packed = torch.stack(pooled_packed, dim=0)

    emb_padded = F.normalize(pooled_padded.float(), p=2, dim=1)
    emb_packed = F.normalize(pooled_packed.float(), p=2, dim=1)

    cos = F.cosine_similarity(emb_padded, emb_packed, dim=1)
    emb_diff = (emb_padded - emb_packed).abs()

    print("\nFinal embedding comparison")
    for i in range(len(texts)):
        print(
            f"seq{i}: cosine={cos[i].item():.8f} | "
            f"max_abs={emb_diff[i].max().item():.8f} | "
            f"mean_abs={emb_diff[i].mean().item():.8f}"
        )

if __name__ == " __main__":
    if not torch.cuda.is_available():
        raise RuntimeError("This script requires CUDA.")

    # Main case: the one you care about
    compare_case(
        model_name=MODEL_NAME,
        texts=TEXTS,
        dtype=torch.bfloat16,
        attn_impl="flash_attention_2",
        max_length=MAX_LENGTH,
    )

    # Reference case: useful to see whether differences are only from BF16 + FA2
    compare_case(
        model_name=MODEL_NAME,
        texts=TEXTS,
        dtype=torch.float32,
        attn_impl="eager",
        max_length=MAX_LENGTH,
    )

```

## What this script checks

It tests the two comparisons that matter:

1. **Per-sequence hidden states**

2. **Final pooled embeddings**

## How to read the output

For the **BF16 + FlashAttention2** block, the usual healthy pattern is:

- `max_abs` on hidden states is **not zero**
- `allclose=True` often holds with `atol=1e-2, rtol=1e-2`
- pooled embedding cosine is **very close to 1.0**

For the **FP32 + eager** block, differences should usually be smaller. If FP32/eager matches closely but BF16/FA2 shows only slight drift, that points to expected numeric behavior rather than a packing bug.

## Two small but important fixes compared with your original snippet

- It uses **the same truncation settings** in both paths.
- It compares the **correct slices** for each sequence instead of eyeballing full tensors.