Accounting error in FLOPS calculation

[Harindhar10]

Reproduction

from transformers import AutoConfig
from trl.trainer.utils import compute_flops_per_token

DENSE_MODEL_ID = "trl-internal-testing/tiny-Qwen3ForCausalLM"

# Tied and untied should have the same FLOPS
cfg = AutoConfig.from_pretrained(DENSE_MODEL_ID)
cfg.tie_word_embeddings = True
f_tied = compute_flops_per_token(cfg, 16384)
cfg.tie_word_embeddings = False
f_untied = compute_flops_per_token(cfg, 16384)
expected_delta = 0
assert f_untied - f_tied == expected_delta

outputs:

AssertionError                            Traceback (most recent call last)
[/tmp/ipykernel_1600/1286663483.py](https://localhost:8080/#) in <cell line: 0>()
     11 f_untied = compute_flops_per_token(cfg, 16384)
     12 expected_delta = 0
---> 13 assert f_untied - f_tied == expected_delta

AssertionError:

System Info

• Platform: Linux-6.6.122+-x86_64-with-glibc2.35
• Python version: 3.12.13
• TRL version: 1.9.2
• PyTorch version: 2.11.0+cu128
• accelerator(s): Tesla T4
• Transformers version: 5.13.1
• Accelerate version: 1.14.0
• Accelerate config: not found
• Datasets version: 5.0.1
• HF Hub version: 1.23.0
• bitsandbytes version: not installed
• DeepSpeed version: not installed
• Liger-Kernel version: not installed
• PEFT version: 0.19.1
• vLLM version: not installed

Checklist

• I have checked that my issue isn't already filed (see open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible (more on MREs)
• Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
• Any traceback provided is complete