v0 param server (using collectives not object store)

[mikaylagawarecki]

Stack from ghstack (oldest at bottom):

• -> v0 param server (using collectives not object store) #2865

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2865

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 7 New Failures, 1 Cancelled Job, 1 Unrelated Failure

As of commit 32c10d7 with merge base 04d70c1 ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 4dbc541ccd54ef94d84a3cc0540d7a13ee841f3b02373a45eb358c29eb1fcaee /exec failed with exit code 1
• Lint / python-source-and-configs / linux-job (gh)
  torchrl/collectors/vllm_weight_update.py:191:9: F401 'transformers.AutoModel' imported but unused
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
  test/test_env.py::TestNonTensorEnv::test_parallel[False-True]
• Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh)
  RuntimeError: Command docker exec -t 7a0121bec6daf52cfb9989f88871ca4d4764cf10fe4824304174f3bea5425632 /exec failed with exit code 1
• Unit-tests on Linux / tests-optdeps (3.11, 12.8) / linux-job (gh)
  test/test_cost.py::TestPPO4LLMs::test_hf[False]

CANCELLED JOB - The following job was cancelled. Please retry:

• LLM Tests on Linux / unittests (3.9, 12.8) / linux-job (gh)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh) (trunk failure)
  test/test_transforms.py::TestTimer::test_transform_env

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mikaylagawarecki]

@vmoens This is one part where I'm trying to call a method on the LLM object to init a process group with the vllm workers, the second part is below on L293

In this case
SyncDataCollector is remote and has an attribute .policy

.policy is the ModuleDict object returned by from_vllm, and the actual llm instance is in the generate key (the LLM instance is local to the SyncDataCollector)

How can I have a handle to the LLM instance within the SyncDataCollector to call remote methods on it without the hacky call_policy_method implementation below?

[mikaylagawarecki]

@vmoens I keep finding that I want to get info off vllm directly :/, what would you do here?

Should this vLLMRemoteWeightUpdaterBase be aware of all the vllm engines and tp size of each owned by its parent RayCollector in its __init__, I already needed to pass separate master_address and master_port I guess

[vmoens]

I think you're right, there's no way around it

We need for the main worker to know that stuff about the remote ones

[joecummings]

import find_spec might be preferred: https://docs.python.org/3/library/importlib.html#importlib.util.find_spec as it doesn't do the actual import until needed.

[vmoens]

I agree, that's the proper way to do it
All third party imports should be done locally even if not optional (otherwise that slows down multiproc / distributed start time and can cause bugs that are hard to debug)