First clip_grad call takes too long

Describe the bug

First call to clip_grad is too slow. In particular spends time in the following:

Thread 3994769 (active+gil): "MainThread"
    redistribute_cost (torch/distributed/tensor/_collective_utils.py:342)
    generate_redistribute_costs (torch/distributed/tensor/_ops/utils.py:268)
    stack_strategy (torch/distributed/tensor/_ops/_tensor_ops.py:753)
    propagate_op_sharding_non_cached (torch/distributed/tensor/_sharding_prop.py:331)
    __call__ (torch/distributed/tensor/_sharding_prop.py:46)
    propagate (torch/distributed/tensor/_sharding_prop.py:311)
    dispatch (torch/distributed/tensor/_dispatch.py:150)
    __torch_dispatch__ (torch/distributed/tensor/_api.py:358)
    _fn (torch/_dynamo/eval_frame.py:1005)
    inner (torch/_compile.py:53)
    _get_total_norm (torch/nn/utils/clip_grad.py:107)
    _no_grad_wrapper (torch/nn/utils/clip_grad.py:43)
    _clip_grad_norm_impl (nemo_automodel/components/training/utils.py:101)
    decorate_context (torch/utils/_contextlib.py:120)
    clip_grad_norm (nemo_automodel/components/training/utils.py:200)
    decorate_context (torch/utils/_contextlib.py:120)
    scale_grads_and_clip_grad_norm (nemo_automodel/components/training/utils.py:306)
    decorate_context (torch/utils/_contextlib.py:120)
    _run_train_optim_step (train_ft.py:1272)
    run_train_validation_loop (train_ft.py:1127)
    main (finetune.py:29)
    <module> (finetune.py:33)

Steps/Code to reproduce bug

Please list minimal steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

Expected behavior

A clear and concise description of what you expected to happen.

Additional context

Add any other context about the problem here.

Start with the reported call path in nemo_automodel/components/training/utils.py, especially _clip_grad_norm_impl and clip_grad_norm, then inspect torch/nn/utils/clip_grad.py and the distributed tensor files named in the stack. Establish a minimal reproduction, measure the first clip_grad call, and confirm that its startup overhead is reduced without changing gradient-clipping behavior.