FSDP Worker (single-GPU) optimizations

Hello,

I see here that FSDP worker always loads on CPU and uses float32:

https://github.com/NVIDIA/NeMo-RL/blob/1db1d75e17665c52d02aff4d623414f10a410334/nemo_rl/models/policy/fsdp1_policy_worker.py#L88

Is there still a reason for this? At least in the 1 GPU case, it seems I can switch this to cuda:0 without any ill effect, and it speeds up my checkpoint loading 3x (need to list device #, just cuda won't work because it still CPU copies) - does it make sense to make it configurable?

Additionally, I see it always loads in float32 because of a bug. That bug would appear fixed - and can verify that at least switching to bfloat16 works fine (float16 still diverges after a single policy update though).

Speedups for Qwen3 1.7B on 1x A100 80GB:

# CPU loading:
Loading checkpoint shards: 50%|█████ | 1/2 [00:04<00:04, 4.46s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:05<00:00, 2.62s/it]

# GPU loading:
Loading checkpoint shards: 50%|█████ | 1/2 [00:01<00:01, 1.53s/it]
Loading checkpoint shards: 100%|██████████| 2/2 [00:01<00:00, 1.17it/s]

# float32 (3 trials):

• Total step time: 42.58s
• Total step time: 42.42s
• Total step time: 42.18s

# bf16 (3 trials):

• Total step time: 35.50s
• Total step time: 34.03s
• Total step time: 34.50s

Thank you!