feat: optimize refit by preparing refit info ahead of time

[yuki-97]

Separate the refit process changes from #613.

What does this PR do ?

1. Update refit process, prepare refit info ahead of time to avoid repeated sending metadata from train side to inference side each time.
2. Add NRL_REFIT_BUFFER_MEMORY_RATIO and update the default value from 10% to 20%.
3. Found the dtype of some tensors (e.g. e_score_correction_bias) will change during training, have some special handle with it, and refit_param_info_mcore is not cached for now because of this.

Test Result

convergence

dtensor	mcore	mcore w/ packing (dsv3 w/ 64 tp)

time cost
In mcore w/ packing (dsv3 w/ 64 tp)

1. Speedup from 420s to 105s with preparing metainfo.
2. Speedup from 105s to 55s with updating refit buffer memory ratio.

*The ~20s overhead is due to offload.

Refit Process Changes

Colocated

Previous

• Each time in refit
  1. prepare_weights_for_ipc in train side.
  2. get_weights_ipc_handles in train side and update_weights_from_ipc_handles in inference side.

Now

• Init
  1. prepare_refit_info in train side.
• Each time in refit
  1. prepare_weights_for_ipc in train side.
  2. get_weights_ipc_handles in train side and update_weights_from_ipc_handles in inference side.

Non-colocated

Previous

• Init
  1. init_collective in both train and inference side.
• Each time in refit
  1. prepare_info_for_collective in train side.
  2. broadcast_weights_for_collective in train side and update_weights_from_collective in inference side.

Now

• Init
  1. init_collective in both train and inference side.
  2. prepare_refit_info in both train and inference side.
• Each time in refit
  1. broadcast_weights_for_collective in train side and update_weights_from_collective in inference side.

[ZhiyuLi-Nvidia]

@YUki-666
just curious, does the refit_info correspond to the policy model (like DTensor or Mcore) or generation model (vllm one)?
At first glance, look like refit_info in the change is relevant to the policy model. We probably need the generation model (vllm one) to avoid duplicated meta data transfer in collective_rpc.
Do you know if we can extract the dtype and shape without actually converting the params (from policy model to the generation one).

[yuki-97]

Do you know if we can extract the dtype and shape without actually converting the params (from policy model to the generation one).

Yup, I added it in ebb874a.

[ZhiyuLi-Nvidia]

Thank you @YUki-666. LGTM!