feat: optimize refit by reducing set of IPC handles sent to each device

[ZhiyuLi-Nvidia]

What does this PR do ?

Optimization for refitting:

• avoid duplicate ipc passing across ray: it removes one linearity factor in TP scaling in refitting. Instead of broadcast all ipc handles just sent the necessary ones. The overhead in argument serialization is constant to TP scaling now.
  • see ~25% gain in refitting speed in small scale and expect to be much more impactful with with large TP size.
  • in dsv3 w/ 64 tp. Found 3.5x speed (from 420s to 120s) with the change.
• avoid waiting for update_weights_from_ipc_handles for better overlap:
  • avoid unnecessary waiting time for better overlap ~8% gain in refitting speed in small scale
  • reverted given unsafe memory control

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

• ...

[ZhiyuLi-Nvidia]

This change is compatible to async_llm.
TODO: fix conflicts in async LLM

[guyueh1]

LGTM

[yuki-97]

@ZhiyuLi-Nvidia report_device_id is already supported in asyncLLM, see report_device_id_async.

[ZhiyuLi-Nvidia]

@ZhiyuLi-Nvidia report_device_id is already supported in asyncLLM, see report_device_id_async.

Yeap. I have seen this implementation. I had some issues to correctly call it within VllmGenerationWorker:

• report_device_id_async is now used in VllmGeneration along with worker_group functions like run_all_workers_single_data
• I want to call and wait for the result within the VllmGenerationWorker scope. I have tried but either returning future/coroutine objects or not correctly call the async report_device_id_async function. Still need sometime to investigating.

Let me know if you have any suggestions