HF Qwen 2 with Thunder returns a slightly different loss function output

[IvanYashchuk]

🐛 Bug

We need to determine whether Thunder has real accuracy problems computing HF's Qwen 2 model.

The test added in #1406 might fail because the loss computed by the Thunder-generated function is slightly different from HF's implementation. Here's the snippet to reproduce the problem:

import torch
from thunder.dynamo import ThunderCompiler
from transformers import Qwen2Config, Qwen2ForCausalLM
torch.manual_seed(0)

# https://huggingface.co/Qwen/Qwen2.5-7B-Instruct/blob/main/config.json
configuration = Qwen2Config(
    # Qwen2.5-7B-Instruct uses Grouped-Query Attention, while the default
    # config uses Multi-Head Attention
    num_attention_heads=28,
    num_key_value_heads=4,
    # Scaled down for testing
    hidden_size=56,
    vocab_size=2,
    max_position_embeddings=32,
)
configuration.num_hidden_layers = 1
with torch.device("cuda"):
    model = Qwen2ForCausalLM(configuration).to(torch.bfloat16)

# thunder.jit doesn't work with Qwen2, so we use torch.compile
# https://github.com/Lightning-AI/lightning-thunder/issues/1405
backend = ThunderCompiler()
compiled_model = torch.compile(model, backend=backend, fullgraph=True)

input_ids = torch.randint(0, configuration.vocab_size, (1, configuration.max_position_embeddings), device="cuda")
# input_ids = torch.ones_like(input_ids) * 0
ref_output = model(input_ids=input_ids, labels=input_ids)
ref_loss = ref_output.loss

compiled_output = compiled_model(input_ids=input_ids, labels=input_ids)
compiled_loss = compiled_output.loss
torch.testing.assert_close(compiled_loss, ref_loss)

AssertionError: Scalars are not close!

Expected 0.7005462646484375 but got 0.7004587650299072.
Absolute difference: 8.749961853027344e-05 (up to 1e-05 allowed)
Relative difference: 0.00012490198427392128 (up to 1.3e-06 allowed)

Thunder may return a different result because upcasting and downcasting to bf16 are different. However, we need to know that Thunder's is indeed more accurate by comparing the distance to the fp64 result and the tolerances in the test may need to be tweaked.

cc @apaz-cli

[IvanYashchuk]

The accuracy problem could be related to #1338.

[riccardofelluga]

Quick update: dynamo doesn't seems to be able to capture a full graph anymore and the flag full_graph needs to be set to False. With that in mind the same error appears when using inductor as a backend with numerical differences on a similar scale as ThunderCompile

[riccardofelluga]

The graph break was introduced by the new loss function selection logic in HF transformers and it's going to (hopefully) be fixed in huggingface/transformers#34616

[riccardofelluga]

Update on this issue: the problem doesn't look like it's coming from the loss function but from the evaluation of the model instead. The snipped above can be modified to compare the logits and the same behavior appears.

Modified repro

import torch
from thunder.dynamo import ThunderCompiler
from transformers import Qwen2Config, Qwen2ForCausalLM
torch.manual_seed(0)

# https://huggingface.co/Qwen/Qwen2.5-7B-Instruct/blob/main/config.json
configuration = Qwen2Config(
    # Qwen2.5-7B-Instruct uses Grouped-Query Attention, while the default
    # config uses Multi-Head Attention
    num_attention_heads=28,
    num_key_value_heads=4,
    # Scaled down for testing
    hidden_size=56,
    vocab_size=2,
    max_position_embeddings=32,
)
configuration.num_hidden_layers = 1
with torch.device("cuda"):
    model = Qwen2ForCausalLM(configuration).to(torch.bfloat16)

# thunder.jit doesn't work with Qwen2, so we use torch.compile
# https://github.com/Lightning-AI/lightning-thunder/issues/1405
backend = ThunderCompiler()
compiled_model = torch.compile(model, backend=backend, fullgraph=True)

input_ids = torch.randint(0, configuration.vocab_size, (1, configuration.max_position_embeddings), device="cuda")
# input_ids = torch.ones_like(input_ids) * 0
ref_output = model(input_ids=input_ids, labels=None)

#labels to None so it does not trace into the loss fn.
compiled_output = compiled_model(input_ids=input_ids, labels=None)
torch.testing.assert_close(compiled_output.logits, ref_output.logits)

Interestingly enough, while trying multiple datatypes for the model's weights, the numerical mismatch happens only for bfloat16 and not for float16, float32 nor float64. Moreover, swapping backend to inductor still shows the mismatch. My idea here is that there is some other mixed precision issue deeper(like a couple of extra cast to from fp32)

For example, setting the following flag halves the amount of numerical mismatch measured, still does not solve the issue.

torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False

At this point the differences on the logits are:

# Inductor backend
Mismatched elements: 12 / 64 (18.8%)
Greatest absolute difference: 0.00146484375 at index (0, 31, 0) (up to 1e-05 allowed)
Greatest relative difference: 0.0439453125 at index (0, 20, 1) (up to 0.016 allowed)

# ThunderCompiler backend
Mismatched elements: 12 / 64 (18.8%)
Greatest absolute difference: 0.00146484375 at index (0, 31, 0) (up to 1e-05 allowed)
Greatest relative difference: 0.036865234375 at index (0, 21, 1) (up to 0.016 allowed)

[IvanYashchuk]

Interesting! We don't have control over Inductor, but can we find operations or layers when the drift starts in Thunder compared to Eager?

the numerical mismatch happens only for bfloat16 and not for float16, float32 nor float64

float16 is a more accurate representation of floating points than bfloat16 (10 bits vs 7 bits for storing significand). More errors are accumulated with bfloat16 operations.

[riccardofelluga]

Hmm yes, that is my next step! Actually I think that it is more likely that we find the drift in the layers rather than the single ops, tho it can't be excluded 100%.

[riccardofelluga]

It seems that the numerical differences do not come from a single module, for example: in the case of Qwen2DecoderLayer it's components, when tested individually, both show numerical mismatch with eager. Since both Qwen2Attention and Qwen2MLP can cause this issue I've restricted the search to the MLP module which is simpler to look at.

New repro with only Qwen2MLP-like module

import torch
import torch.nn as nn
from thunder.dynamo import ThunderCompiler
torch.manual_seed(0)

class QwenLIKE2MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden_size = 2
        self.intermediate_size = 3
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = nn.SiLU()

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


with torch.device("cuda"):
    model = QwenLIKE2MLP().to(torch.bfloat16)

backend = ThunderCompiler(executors=("torchcompile",))
# backend = ThunderCompiler(executors=("nvfuser",), nv_enable_linear=False)
compiled_model = torch.compile(model, backend=backend, fullgraph=True)

mlp_in = torch.randn(32, 2, dtype=torch.bfloat16, device="cuda")
ref_output = model(mlp_in)

compiled_output = compiled_model(mlp_in)

torch.testing.assert_close(compiled_output, ref_output)

When executing the above snippet with nvfuser instead of inductor as fusion executor, the numerical mismatch disappears. The clear difference being that nvFuser does not include linear in the fusion by default. Luckily we do have nv_enable_linear that when set to true allows nvfuser to capture the linears and therefore create a similar fusion as the one made by inductor. Here is where things get interesting tho, when adding the linear to the fusion the numerical mismatch reappears even tho nothing really changed since, under the hood, the same kernels are called(nvfuser relies on aten kernels for linear layers, verified by profiling).

nvFuser trace with nv_enable_linear=False

# Constructed by Delete Last Used (took 0 milliseconds)
import torch
import torch.nn.functional
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(l_x_, l_self_modules_gate_proj_parameters_weight_, l_self_modules_up_proj_parameters_weight_, l_self_modules_down_proj_parameters_weight_):
  # l_x_: "cuda:0 bf16[32, 2]"
  # l_self_modules_gate_proj_parameters_weight_: "cuda:0 bf16[3, 2]"
  # l_self_modules_up_proj_parameters_weight_: "cuda:0 bf16[3, 2]"
  # l_self_modules_down_proj_parameters_weight_: "cuda:0 bf16[2, 3]"

  # <eval_with_key>.4:5: 	    linear = torch._C._nn.linear(l_x_, l_self_modules_gate_proj_parameters_weight_, None);  l_self_modules_gate_proj_parameters_weight_ = None
  t78 = torch.nn.functional.linear(l_x_, l_self_modules_gate_proj_parameters_weight_, None)  # t78: "cuda:0 bf16[32, 3]"
    # t78 = ltorch.linear(l_x_, l_self_modules_gate_proj_parameters_weight_, None)  # t78: "cuda:0 bf16[32, 3]"
      # t78 = prims.linear(l_x_, l_self_modules_gate_proj_parameters_weight_, None)  # t78: "cuda:0 bf16[32, 3]"
  [silu] = nvFusion0(t78)
    # t18 = prims.convert_element_type(t78, dtypes.float32)  # t18: "cuda:0 f32[32, 3]"
    # t19 = prims.neg(t18)  # t19: "cuda:0 f32[32, 3]"
    # t20 = prims.exp(t19)  # t20: "cuda:0 f32[32, 3]"
    # t21 = prims.add(1.0, t20)  # t21: "cuda:0 f32[32, 3]"
    # t22 = prims.reciprocal(t21)  # t22: "cuda:0 f32[32, 3]"
    # t26 = prims.mul(t18, t22)  # t26: "cuda:0 f32[32, 3]"
    # silu = prims.convert_element_type(t26, dtypes.bfloat16)  # silu: "cuda:0 bf16[32, 3]"

  # <eval_with_key>.4:7: 	    linear_1 = torch._C._nn.linear(l_x_, l_self_modules_up_proj_parameters_weight_, None);  l_x_ = l_self_modules_up_proj_parameters_weight_ = None
  t79 = torch.nn.functional.linear(l_x_, l_self_modules_up_proj_parameters_weight_, None)  # t79: "cuda:0 bf16[32, 3]"
    # t79 = ltorch.linear(l_x_, l_self_modules_up_proj_parameters_weight_, None)  # t79: "cuda:0 bf16[32, 3]"
      # t79 = prims.linear(l_x_, l_self_modules_up_proj_parameters_weight_, None)  # t79: "cuda:0 bf16[32, 3]"
  [mul] = nvFusion1(silu, t79)
    # t29 = prims.convert_element_type(silu, dtypes.float32)  # t29: "cuda:0 f32[32, 3]"
    # t30 = prims.convert_element_type(t79, dtypes.float32)  # t30: "cuda:0 f32[32, 3]"
    # t31 = prims.mul(t29, t30)  # t31: "cuda:0 f32[32, 3]"
    # mul = prims.convert_element_type(t31, dtypes.bfloat16)  # mul: "cuda:0 bf16[32, 3]"

  # <eval_with_key>.4:9: 	    down_proj = torch._C._nn.linear(mul, l_self_modules_down_proj_parameters_weight_, None);  mul = l_self_modules_down_proj_parameters_weight_ = None
  t80 = torch.nn.functional.linear(mul, l_self_modules_down_proj_parameters_weight_, None)  # t80: "cuda:0 bf16[32, 2]"
    # t80 = ltorch.linear(mul, l_self_modules_down_proj_parameters_weight_, None)  # t80: "cuda:0 bf16[32, 2]"
      # t80 = prims.linear(mul, l_self_modules_down_proj_parameters_weight_, None)  # t80: "cuda:0 bf16[32, 2]"
  return {'output': (t80,), 'flat_args': [l_x_, l_self_modules_gate_proj_parameters_weight_, l_self_modules_up_proj_parameters_weight_, l_self_modules_down_proj_parameters_weight_], 'flat_output': (t80,)}, ((l_self_modules_down_proj_parameters_weight_, l_x_, mul, silu, t78, t79), ())

nvFuser trace with nv_enable_linear=True

# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(l_x_, l_self_modules_gate_proj_parameters_weight_, l_self_modules_up_proj_parameters_weight_, l_self_modules_down_proj_parameters_weight_):
  # l_x_: "cuda:0 bf16[32, 2]"
  # l_self_modules_gate_proj_parameters_weight_: "cuda:0 bf16[3, 2]"
  # l_self_modules_up_proj_parameters_weight_: "cuda:0 bf16[3, 2]"
  # l_self_modules_down_proj_parameters_weight_: "cuda:0 bf16[2, 3]"
  [linear, t26, linear_1, mul, down_proj] = nvFusion0(l_x_, l_self_modules_gate_proj_parameters_weight_, l_self_modules_up_proj_parameters_weight_, l_self_modules_down_proj_parameters_weight_)
    # linear = prims.linear(l_x_, l_self_modules_gate_proj_parameters_weight_, None)  # linear: "cuda:0 bf16[32, 3]"
    # t18 = prims.convert_element_type(linear, dtypes.float32)  # t18: "cuda:0 f32[32, 3]"
    # t19 = prims.neg(t18)  # t19: "cuda:0 f32[32, 3]"
    # t20 = prims.exp(t19)  # t20: "cuda:0 f32[32, 3]"
    # t21 = prims.add(1.0, t20)  # t21: "cuda:0 f32[32, 3]"
    # t22 = prims.reciprocal(t21)  # t22: "cuda:0 f32[32, 3]"
    # t26 = prims.mul(t18, t22)  # t26: "cuda:0 f32[32, 3]"
    # linear_1 = prims.linear(l_x_, l_self_modules_up_proj_parameters_weight_, None)  # linear_1: "cuda:0 bf16[32, 3]"
    # t30 = prims.convert_element_type(linear_1, dtypes.float32)  # t30: "cuda:0 f32[32, 3]"
    # t31 = prims.mul(t26, t30)  # t31: "cuda:0 f32[32, 3]"
    # mul = prims.convert_element_type(t31, dtypes.bfloat16)  # mul: "cuda:0 bf16[32, 3]"
    # down_proj = prims.linear(mul, l_self_modules_down_proj_parameters_weight_, None)  # down_proj: "cuda:0 bf16[32, 2]"
  return {'output': (down_proj,), 'flat_args': [l_x_, l_self_modules_gate_proj_parameters_weight_, l_self_modules_up_proj_parameters_weight_, l_self_modules_down_proj_parameters_weight_], 'flat_output': (down_proj,)}, ((l_self_modules_down_proj_parameters_weight_, l_x_, linear, linear_1, mul, t26), ())

[riccardofelluga]

Update! After splitting the model apart, the drift can be detected in this two sections of code: Qwen2MLP(previous comment) and apply_rotary_pos_emb function in Qwen2Attention module(note that it is the only part of the module that results in numerical mismatch). In the end, the mismatch for the whole model boils down to the following snippets:

# Qwen2MLP equivalent
import torch
import torch.nn.functional as F
from thunder.dynamo import thunderfx
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
torch.manual_seed(0)

l0 = torch.randn(2, 2, dtype=torch.bfloat16, device="cuda")
l1 = torch.randn(2, 2, dtype=torch.bfloat16, device="cuda")

def bar(a, b):
    t0 = F.silu(a, inplace=False)
    t1 = b * t0
    return t1

cbar = thunderfx(bar, fullgraph=True)

compiled = cbar(l0, l1)
ref = bar(l0, l1)

torch.testing.assert_close(compiled, ref, rtol=0.0, atol=0.0)

# apply_rotary_pos_emb equivalent
import torch
from thunder.dynamo import thunderfx
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
torch.manual_seed(0)

def foo(x1, x2, cos, sin):
    t0 = -x2
    q = torch.cat((t0, x1), dim=-1)
    t1 = q * sin
    t2 = cos + t1
    return t2

cfoo = thunderfx(foo, fullgraph=True)

x1 = torch.rand(1, 28, 32, 1, dtype=torch.float32, device="cuda")
x2 = torch.rand(1, 28, 32, 1, dtype=torch.float32, device="cuda")
cos = torch.randn(1, 1, 32, 2, dtype=torch.float32, device="cuda")
sin = torch.randn(1, 1, 32, 2, dtype=torch.float32, device="cuda")

compiled = cfoo(x1, x2, cos, sin)
ref = foo(x1, x2, cos, sin)

torch.testing.assert_close(compiled, ref, rtol=0.0, atol=0.0)

For both snippets, all the compiled versions differ from eager at the same rate(torch.compile, thunderfx, and thunder.jit). Interestingly enough, in the second snippet the mismatch happens for both fp32 and bf16 inputs, tho it is worst in the case of half precision.

Next question is, is eager less accurate than thunder or the opposite?

[mruberry]

Executing the code eagerly in double precision is probably the best "golden value" estimator.

[riccardofelluga]

After comparing the snippets with the golden standard of float64, it really looks like Thunder is more accurate than eager. And this makes sense because when the inputs are not in fp32, up-casts and down-casts are needed to actually execute the computation in fp32. The difference that sets apart Thunder and eager is that Thunder is able to minimize the quantity of up/down-casts whereas eager cannot. (here for implementation)

Results for the MLP example(bf16 vs fp64 reference):

# Eager Mismatched elements: 4 / 4 (100.0%)
# Greatest absolute difference: 0.022391211651423326 at index (1, 0)
# Greatest relative difference: 0.005960235582927172 at index (1, 0)

# Thunder mismatched elements: 4 / 4 (100.0%)
# Greatest absolute difference: 0.006766211651423326 at index (1, 0)
# Greatest relative difference: 0.0018010733887134785 at index (1, 0)

Results for the rope transformation(fp32 vs fp64 reference):

# Eager Mismatched elements: 1792 / 1792 (100.0%)
# Greatest absolute difference: 2.0036165437886666e-07 at index (0, 19, 26, 1)
# Greatest relative difference: 1.8068834159046413e-05 at index (0, 2, 9, 1)

# Thunder Mismatched elements: 1792 / 1792 (100.0%)
# Greatest absolute difference: 1.1765928320528474e-07 at index (0, 18, 22, 0)
# Greatest relative difference: 5.862098371807772e-08 at index (0, 4, 6, 1)

As we can see Thunder has a lower numerical drift than eager. This makes the point for the importance of paying attention to half precision operations when writing modules like the apply_rotary_pos_emb. If the practitioners were to upcast the inputs, execute all the ops in fp32 and then downcast the output, performing ops in eager wouldn't be different than Thunder/Inductor. With this I'd mark this issue as closed 🎉