fix(training): all-reduce grad-norm scalars on the mesh device (#3461)

* fix(training): all-reduce grad-norm scalars on the mesh device

Under FSDP2 CPUOffloadPolicy the local shard of ``p.grad`` lives in host RAM,
so ``_clip_grad_norm_impl`` derives ``target_device = cpu`` and allocates its
norm accumulators there. Those scalars were then all-reduced over the mesh's
NCCL process group, which has no CPU backend, aborting training on every rank
at the first optimizer step with:

    RuntimeError: No backend type associated with device type cpu

Split the compute device from the communication device: the norm math stays on
the gradients' own device, and only the 0-dim accumulator hops to
``mesh.device_type`` for the collective before coming straight back. Simply
allocating the scalar on CUDA is not sufficient -- the scaled-sum pass divides
the CPU gradient by that scalar, and PyTorch's 0-dim promotion is one-way
(a CPU scalar mixes with a CUDA tensor, a CUDA scalar does not mix with a CPU
tensor).

All five collectives in the clip path are routed through the new helper; fixing
only the first one moves the failure to the second.

When the scalar already sits on the mesh device -- every configuration without
CPU offload, plus genuinely-CPU gloo meshes -- the helper performs the same
in-place all_reduce on the same tensor object, so those paths are unchanged.

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

* test(training): give the mocked PP mesh a real device_type

test_clip_grad_norm_with_pp_and_tp stands in for the pipeline-parallel mesh with
a bare Mock, so `pp_mesh.device_type` auto-generated a Mock rather than a string.
The norm reduction now compares that against the accumulator's device to choose
the collective's device, and a Mock never compares equal to "cpu", so the test
drove the code down the device-hop path and failed in `Tensor.to(device=<Mock>)`.

A real DeviceMesh always exposes `device_type` as a string, so the double was
simply incomplete. Give it one rather than teaching production code to tolerate
a non-string device type.

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>

* perf(training): avoid host sync in gradient clipping

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* test(training): reuse initialized process group

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

* test(training): scope profiler check to host scalars

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

---------

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
Co-authored-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>