[Tcc0403]

Summary

Fix #761, #742.

This PR found and addressed couple of issues

• The tolerance for DyT's alpha.grad was too strict.
  • Setting its tolerance alone and lowering the values.
• MultitokenAttention couldn't pass on xpu platform.
  • Disabling it for now to unblock the ci.
• Llama4 in convergence test shouldn't be imported in bwd-test ci (transformers==4.44.2).
  • Adding a import guard for Llama4.
• Fix casting issue for geglu/swiglu mlp layer.
  • fix(geglu, swiglu): match the casting in MLP layer for bfloat16 #783
• Some decprecated lce_forward functions didn't accept skip_logits, causing failures for bwd-test ci.
  • Adding skip_logits to lce_forward_deprecated()s.
• Some models are no longer supported in transformers<4.49.0
  • Skipping those test cases if transformers version don't meet the requirement.
• Llama4_multimodal implicitly casts image input to bfloat16, couldn't work with float32 model.
  • Commenting llama4 out from fp32 multimodal test.

Minor change:

• Enhancing testing utility function assert_verbose_allclose() to show addional information at the start of the error message.

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[shimizust]

@Tcc0403 thanks for fixing, although I see it says do not merge for now. Are there more changes needed?

[shimizust]

Trying to square with #778, although I'd prefer not commenting out the tests if possible

[Tcc0403]

@shimizust I was trying to make convergence tests passed as well but didn't have time. I'll try making it ready today.

[Tcc0403]

Trying to square with #778,

Oops you mean #775 right? My bad, I didn't see that before creating this PR.

although I'd prefer not commenting out the tests if possible

+1. Since we didn't change tolerance for logits when replacing them with logrpobs, I think some of them become too strict for logrpobs.

[Tcc0403]

@shimizust
Some test cases are too flaky, but commenting them out or setting high tolerance just for them don't make sense to me.

Instead, I propose a multi-stage tolerance testing:
Given a list of tolerance values, we perform all_close testing from low to high tolerance. With this method, we can

1. set a universal standard for all models, and
2. models can be further categorized into different levels based on which stage they can pass, like a patching confident level.

I've made a simplest implementation of the proposal and applied it to fp32/test_mini_models.py. I'll apply this method to all convergence tests if you think it's applicable.

For the second usage I mentioned, we can redirect output in the future.

[shimizust]

Hey @Tcc0403 thanks for the proposal on making the test tolerances more systematic. The problem I see is that the tolerance values I don't think would be expected to have a universal standard and should be expected to be a function of model architecture (e.g. MoE, multimodal, activation function), random initialized weights, the mini model config used for testing, etc.. We can set some initial buckets, but it's still arbitrary, and the addition of a new model type may necessitate creating another bucket or redefining the low/med/high tolerances so tests pass.

Another potential issue is that it may obscure real regressions, like if a model was initially in the high tolerance category, but a bug caused it to move to low tolerance, it wouldn't be immediately obvious. Thus, I think each model requires tracking its own tolerance, with some reasons added for why a model needs higher tolerance.

Another approach is to empirically find an acceptable tolerance for each model by running multiple experiments with different random seeds, and measuring the natural range of variation (atol/rtol) across N runs of the original model.

For now, I think it would be reasonable to proceed by tuning tolerances incrementally to unblock the CI, while figuring out a more reasoned approach in the meantime. What do you think?

@lancerts did you have any thoughts?

[Tcc0403]

Thanks for the feedback.

The problem I see is that the tolerance values I don't think would be expected to have a universal standard and should be expected to be a function of model architecture (e.g. MoE, multimodal, activation function), random initialized weights, the mini model config used for testing, etc.. We can set some initial buckets, but it's still arbitrary, and the addition of a new model type may necessitate creating another bucket or redefining the low/med/high tolerances so tests pass.

I see, model architecture, initial weights, the config may have a huge impact on numerical errors as well, so universal tolerances are only viable if we can find the settings that can generate similar results, i.e. in the same floating range, for each models. I can see that the model architecture would affect the most. If we are going to set a unviersal standard, we need to find the tolerance for different model architecture at the very least.

Another potential issue is that it may obscure real regressions, like if a model was initially in the high tolerance category, but a bug caused it to move to low tolerance, it wouldn't be immediately obvious. Thus, I think each model requires tracking its own tolerance, with some reasons added for why a model needs higher tolerance.

My idea is to show the category for each models on our future regression tracking site one day, so moving between categories can be captured (maybe make a dedicated benchmarking for this purpose). Since tolerances are also represented as the confidence level of our patching, if a small bug causes the model dropping into the lower bracket, it is still acceptable. We can always keep track of it and decide whether we should start investigating or put it off. Plus, if we set the tolerance too high, such small bugs are more difficult to be detected.

Another approach is to empirically find an acceptable tolerance for each model by running multiple experiments with different random seeds, and measuring the natural range of variation (atol/rtol) across N runs of the original model.

It might work but it also increases the burden of supporting new models for contributors.

For now, I think it would be reasonable to proceed by tuning tolerances incrementally to unblock the CI, while figuring out a more reasoned approach in the meantime. What do you think?

Gotcha. I'll go with the original method and unblock the ci first.

[Tcc0403]

Some lce_forward_deprecated() don't support argument skip_logits. Adding them in this pr as well.

[Tcc0403]

@shimizust I found that the casting in GEGLUMLP computation can result in huge discrepancy in bf16 scenario. See #783.

(Update: SwiGLUMLP as well)

We can merge that PR first. The change of gelu kernel is included in this PR just to help fixing ci.

Before the fix

We can see there's a huge gap between two models.

E           AssertionError: [Top k logprobs]Number of mismatched elements: 1
E           Mismatch at index (0, 1015, 0): tensor1[(0, 1015, 0)] = -1.2275793552398682, tensor2[(0, 1015, 0)] = -1.5559215545654297

Some statistics:

max values at the mismatch row differ, resulting a large discrepancy after logsoftmax computation

# original model
(expected_output['logits'][0, 1015, :]).max()=tensor(10.2500, device='cuda:0', grad_fn=<MaxBackward1>) 
(expected_output['topk_logprobs'][0, 1015, :])=tensor([-1.2276, -6.9463, -6.9463, -6.9776, -7.0713, -7.0713, -7.1338, -7.1338,
        -7.1651, -7.1963, -7.2588, -7.3213, -7.3838, -7.4151, -7.4151, -7.4151,
        -7.4463, -7.4463, -7.4463, -7.4463], device='cuda:0')
# patched model
(actual_output['logits'][0, 1015, :]).max()=tensor(9.8125, device='cuda:0', grad_fn=<MaxBackward1>)
actual_output['topk_logprobs'][0, 1015, :]=tensor([-1.5070, -6.7570, -6.8820, -6.8820, -6.9445, -6.9445, -7.0070, -7.0070,
        -7.1008, -7.1320, -7.1320, -7.1633, -7.1945, -7.2258, -7.2570, -7.2883,
        -7.3508, -7.3508, -7.3508, -7.3508], device='cuda:0')

After the fix

Sucessfully passed the test

test/convergence/bf16/test_mini_models_multimodal.py::test_mini_model_multimodal[mini_gemma3-32-1e-05-dtype7-0.05-0.05-0.1-0.1-0.01-0.01] PASSED [100%]

Some statistics

max values are now the same, resulting in closer outputs after logsoftmax computation

# original model
(expected_output['logits'][0, 1015, :]).max()=tensor(10.2500, device='cuda:0', grad_fn=<MaxBackward1>) 
(expected_output['topk_logprobs'][0, 1015, :])=tensor([-1.2276, -6.9463, -6.9463, -6.9776, -7.0713, -7.0713, -7.1338, -7.1338,
        -7.1651, -7.1963, -7.2588, -7.3213, -7.3838, -7.4151, -7.4151, -7.4151,
        -7.4463, -7.4463, -7.4463, -7.4463], device='cuda:0')


# patched model
(actual_output['logits'][0, 1015, :]).max()=tensor(10.2500, device='cuda:0', grad_fn=<MaxBackward1>)
actual_output['topk_logprobs'][0, 1015, :]=tensor([-1.2278, -6.8840, -6.9778, -7.0090, -7.1028, -7.1028, -7.1028, -7.1340,
        -7.1653, -7.1965, -7.2278, -7.3215, -7.3840, -7.3840, -7.4153, -7.4465,
        -7.4465, -7.4465, -7.4465, -7.4778], device='cuda:0')

[Tcc0403]

@shimizust Ready for review!

[shimizust]

Thanks so much for getting the ci issues resolved!