Move amax and scale update later in TEv2 grad transform

Nobody has claimed this yet.

## 🚀 Feature

This is an improvement idea for TEv2, if we were to join the forward call's amax and scale update with the backward, there might not be the need for the hidden checkpointing specific transform [here](https://github.com/Lightning-AI/lightning-thunder/blob/5c18b33917586e388e9526eb28cfd48cf8755915/thunder/__init__.py#L549-L552) as there would not be the need to move the checkpointed forward's amax and scale update in the backward anymore.

The new grad transform form [here](https://github.com/Lightning-AI/lightning-thunder/blob/5c18b33917586e388e9526eb28cfd48cf8755915/thunder/executors/transformer_engine_v2ex_impl.py#L202-L248) would look something like this:

```python
def _te_linear_grad_transform(a, w, bias):
[ ... ]

primal, quantized_a, quantized_w = _te_linear_fwd(a, w, bias, input_quantizer, weight_quantizer)

# No more amax scale and update here

grad_out = get_grad(primal)

backward_recipe_state = _get_te_fp8_state(recipe, "backward", 1)

(grad_output_quantizer,) = _get_te_fp8_quantizers(backward_recipe_state, 1)

grad_a, grad_w = _te_linear_bwd(
grad_out,
quantized_a,
quantized_w,
input_quantizer,
weight_quantizer,
grad_output_quantizer,
)

grad_a, grad_w = _te_fp8_amax_and_scale_update(
recipe,
# added forward state here
states=(forward_recipe_state, backward_recipe_state,),
tokens=(grad_a, grad_w),
)

[ ... ]

return primal
```

Another potential benefit would be a simpler [TransformerEngineTransformV2](https://github.com/Lightning-AI/lightning-thunder/blob/5c18b33917586e388e9526eb28cfd48cf8755915/thunder/executors/transformer_engine_v2ex_impl.py#L282C7-L282C35)

Read thunder/executors/transformer_engine_v2ex_impl.py, especially _te_linear_grad_transform and TransformerEngineTransformV2, then inspect the checkpoint-specific transform in thunder/__init__.py around the linked lines. Compare where forward and backward amax/scale updates occur; done means the update moves into the gradient transform without requiring that hidden checkpointing transform and the TEv2 behavior remains correct.

Written by the indexing model from the issue text.