fix: maintain fp32 mlp.router.expert_bias even with bf16 enabled

[ZhiyuLi-Nvidia]

What does this PR do ?

fix: maintain fp32 mlp.router.expert_bias even with bf16 enabled

model.layers.x.mlp.gate.e_score_correction_bias dtype is changing between step 1 to the following steps, which make it impossible to cache para dtype with inconsistency issue:

f"ERROR: {key}: {shape} != {shape_ref}", where shape is about 1st step, while shape_ref is the 2nd step

('ERROR: model.layers.3.mlp.gate.e_score_correction_bias: torch.bfloat16 != torch.float32',)

Details

• hf's model.layers.x.mlp.gate.e_score_correction_bias originates from decoder.layers.x.mlp.router.expert_bias in mcore
• decoder.layers.x.mlp.router.expert_bias should be starting with fp32 here and always fp32 here
• there seems something wrong in setup_megatron_model
  • https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/policy/megatron_policy_worker.py#L192-L199
  • There's hidden cast from fp32 to bf16 (param_dtype) despite fp32 param initialization: https://github.com/NVIDIA/NeMo/blob/33259f2540af6eef375d43fc48bdcbd7ec490c29/nemo/tron/model.py#L119-L121

avoid serializing rebuild_cuda_tensor function
serializing funciton is costly in transfer cudaipc between workers. Got some performance improvement with the change.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[yuki-97]

Target to yukih/prepare-refit-info first for comparing the modification easily.
Will target to main after #638 merged.

The base branch was changed.

[ZhiyuLi-Nvidia]

Close thie PR since it is part of #686.