TE: update test to be more stable

[kshitij12345]

Running pytest thunder/tests/test_transformer_engine_executor.py -k test_te_linear_forward_backward_multiple_iteration[defau --count 200 (or testing with mxfp8_e4m3 recipe).

Could lead to

>       assert_close(w1, te_linear1.weight)
E       AssertionError: Tensor-likes are not close!
E       
E       Mismatched elements: 16777216 / 16777216 (100.0%)
E       Greatest absolute difference: nan at index (0, 0) (up to 1e-05 allowed)
E       Greatest relative difference: nan at index (0, 0) (up to 0.016 allowed)

This was most likely due to not correctly initializing the linear and using dummy data. Have updated the test to initialize the Linear and also updated the target tensor for training.

Now running pytest thunder/tests/test_transformer_engine_executor.py -k test_te_linear_forward_backward_multiple_iteration[defau --count 200 passes for all cases. Tested on both H100 and B200 where default recipe is DelayedScaling and MXFP8BlockScaling respectively.

[t-vi]

Looks good to me, but let's have @IvanYashchuk take a look if he wants.

[IvanYashchuk]

How does the testing time change with the new target_value?

[kshitij12345]

On H100

Running - pytest thunder/tests/test_transformer_engine_executor.py -k test_te_linear_forward_backward_multiple_iteration[defau --count 200

Main Branch - ~12.3 seconds

Running 200 items in this shard

thunder/tests/test_transformer_engine_executor.py ........................................................................................................................................................................................................                                      [100%]

======================================================================================================================== 200 passed, 2000 deselected, 1200 warnings in 12.30s =========================================================================================================================

PR Branch - ~11.3 seconds

Running 200 items in this shard

thunder/tests/test_transformer_engine_executor.py ........................................................................................................................................................................................................                                      [100%]

======================================================================================================================== 200 passed, 2000 deselected, 1200 warnings in 11.32s =========================================================================================================================

Shouldn't be very different for B200.