Fix LAMB norm reduction with ZeRO CPU offload

[sdjasj]

What changed

• move LAMB's ZeRO data-parallel norm scalar to the working parameter device before the collective
• keep the norm scalar and working-parameter norm on one device for the subsequent stacked reduction
• add a distributed regression case using FP16 parameters, ZeRO-2, initial_scale=2**20, and CPU offload

Why

With CPU offload, the LAMB master shard and its update live on CPU, so g_sum was also created on CPU. The data-parallel group uses NCCL, which cannot all-reduce that CPU tensor and raised No backend type associated with device type cpu.

Impact

LAMB can now take optimizer steps with LowLevelZeroPlugin(cpu_offload=True) while preserving CPU-offloaded master parameters.

Fixes #6400.

Validation

• python -m pytest tests/test_optimizer/test_dist_lamb.py -q (4 GPUs; 1 passed)