Rejection Sampling on GSM8k

[AlexPiche]

Rejection Sampling on GSM8k

• Implement a simple reasoning architecture. COTMathAgent
• Fix vLLM get log probs OOM OOM when getting logprobs from vLLM reference model #77
• Do not include empty user message in nodes.py
• Launch training jobs with accelerate to handle multi-gpu. better multi-gpu inference and training in the RL example #76 (There is an NCCL timeout for large gradient accumulation passes NCCL timeout on multi-gpu #104)
• Get dataset stats (max, min, mean length) in finetune.py
• Fix finetune to start from the latest model instead of the base one.
• Move finetune/rl helper functions into finetune/rl/utils.py
• Implicit KL like Eq3 https://arxiv.org/abs/2402.14740
• Clean up rl orchestrator: remove discounting and max steps

Tape Browser PR from @rizar

• add a script to launch the tape browser
• save tapes in a single json file, which is a format that the default tape browser can load
• add a script to gather legacy tapes in single file so that we can view legacy tapes in the browser
• load LLM Calls 100 at a time to speed up tape browser loading time (the speedup was huge when the browser did not know where to look for LLM Calls; now that I added the search for llm_calls.sqlite to browser.py, loading is taking some time again)

Examples:
run tape browser like this: python -m examples.rl_gsm8k.browse outputs/yolo2_4gpu_lr1e-6/tapes/train/0/all

Reasoning Architecture

Learning curve

python examples/rl_gsm8k/orchestrate_rl.py finetune.rl.algo=reinforce finetune.train_batch_size=4 finetune.gradient_accumulation_passes=1 finetune.rl.implicit_kl_coef=0.0 finetune.rl.kl_coef=0.0 finetune.rl.use_advantages=false +finetune.rl.relu_weights=true use_rejection_sampling=true test_every_n_iterations=5 finetune.learning_rate=0.000001 finetune.gradient_clipping_threshold=1.0 finetune.save_checkpoint_steps=8 finetune.weight_decay=0.1 max_agent_forks=5000 attempts=8

Reproducing GSM8k

[rizar]

will address #77 and #76

[rizar]

Looks mostly good! We should not traverse the whole training set every time the script is launched though.

Note for the re-review:

• check ReLU weights citation
• check running aggregation of input lengths

[rizar]

lgtm!