feat: code execution & tool use

[KiddoZhu]

Sorry I messed up commit history when clearing secrets in some commits (#204). Here is a new clean branch.

@SahilJain314 How shall we decide gen_length if generated logprobs can be 0?

---

Supported features

• On-the-fly multiple tool calls. It will continue generation until hitting an EOS token or user-specified stop strings.
• Batch code execution based on ray.remote.
• Stateful code executor. In the same generation, functions and variables will be passed to the next code snippet.
• A simple sandbox for code execution. Not fully secure for malicious code, but better than nothing.

Examples (test_vllm_tools.py)

<code>x = 3; y = 4</code>
This is some regular text.
<code>x + y</code>
<result>7</result>

<code>retrieve('Jen-Hsun Huang')</code>

<result>
['Nvidia was established in 1993 by Jen-Hsun Huang, Curtis Priem, and Chris '
 'Malachowsky. In 2000 Nvidia took intellectual possession of 3dfx, one of the '
 'biggest GPU producers in 1990s.']
</result>

A simple sandbox

• Filesystem: The code will be executed in a temporary directory. It can read or write any file in this directory. Access beyond this temporary directory is denied.
• Modules: Block modules like os, sys, multiprocessing, subprocess, etc. that may modify the filesystem or reside programs in the memory.

Tokenizer issue

There may be side cases where tokens don't split exactly at the end of </code>. For example, the tokenizer of GPT-4o will generate the following

... </##code##>x

If we tweak >x to be >, it will change the log prob in RL. My current solution is to not touch any generated token and directly append results afterwards

... </code>x<result> ...

[KiddoZhu]

Just added environments for code execution and tool use. Do we need to keep the previous generate_with_code_and_tools implementation? My feeling is that a policy.generate-like interface is more convenient when users just want to do inference.

[KiddoZhu]

@SahilJain314 any comment or suggestion?

[KiddoZhu]

@SahilJain314 I've rolled back both unnecessary implementations you mentioned. Now it's ready to merge except for a mypy check. Do you know how to solve it?

[KiddoZhu]

Just rolled back the logic of generation length in HFPolicy. I opened #499 for future discussions on this issue.

[KiddoZhu]

Hi @SahilJain314 I noticed HF path has been removed in test_rollout.py. How shall I adapt the HF path in tool use? Do we want to remove it as well?

[SahilJain314]

I think it's safe to remove HF generation as a supported rollout path if that's what you're asking

[KiddoZhu]

Thanks! I removed them and all relative tests are passed. Should be ready to merge.

[SahilJain314]

I'm happy to merge this as it stands, but I'd love to see an example (and convergence plots) with this soon (in follow up) to show users how to use it.

@KiddoZhu please also address the DCO.

[KiddoZhu]

DCO resolved.

I'm happy to merge this as it stands, but I'd love to see an example (and convergence plots) with this soon (in follow up) to show users how to use it.

Yes, I will add an example and a doc file once I use code execution in a more realistic setup. It'll be in a separate thread.

[terrykong]

@KiddoZhu a couple of other small things:

• can you mark the tests w/ pytest.mark.hf_gated for ones that need the gated models (like llama)
  • (see this PR for example): https://github.com/NVIDIA-NeMo/RL/pull/755/files
• for things decorated with ray.remote, can you add

@ray.remote  # pragma: no cover

[KiddoZhu]

Thanks Terry! I will do that for future commits.

[shamanez]

Any coding/tool use example, specially with a sandbox?