refactor: remove orchestrator abstraction from API (#289)

* refactor: remove orchestrator abstraction from API

* Remove orchestrator in GPT-J config

* Add `reward_fn` arg to NeMo constructor to match base trainer API

* Initial support for `make_experience` in NeMo ILQL

* Run pre-commit

* Remove unused sampling util

Author: jon-tow

configs/ilql_config.yml

@@ -8,7 +8,6 @@ train:
   eval_interval: 100
 
   pipeline: "PromptPipeline"
-  orchestrator: "OfflineOrchestrator"
   trainer: "AccelerateILQLTrainer"
   seed: 1000
 

configs/nemo_ilql_config.yml

@@ -7,7 +7,6 @@ train:
   eval_interval: 20
 
   pipeline: "PromptPipeline"
-  orchestrator: "OfflineOrchestrator"
   trainer: "NeMoILQLTrainer"
   trainer_kwargs:
     pretrained_model: "/mnt/nvme/home/uwu/nemo-megatron-gpt-20B/"

configs/ppo_config.yml

@@ -8,7 +8,6 @@ train:
   eval_interval: 100
 
   pipeline: "PromptPipeline"
-  orchestrator: "PPOOrchestrator"
   trainer: "AcceleratePPOTrainer"
 
 model:

configs/ppo_gptj.yml

@@ -8,7 +8,6 @@ train:
   eval_interval: 16
 
   pipeline: "PromptPipeline"
-  orchestrator: "PPOOrchestrator"
   trainer: "AcceleratePPOTrainer"
 
 model:

configs/sft_config.yml

@@ -8,7 +8,6 @@ train:
   eval_interval: 100
 
   pipeline: "PromptPipeline"
-  orchestrator: "PPOOrchestrator"
   trainer: "AccelerateSFTTrainer"
 
 model:

configs/test_config.yml

@@ -8,7 +8,6 @@ train:
   eval_interval: 128 # eval interval
 
   pipeline: "PromptPipeline" # prompt pipeline to load
-  orchestrator: "PPOOrchestrator" # orchestrator to load
   trainer: "AcceleratePPOTrainer" # Name of model trainer to load
 
 model:
@@ -36,7 +35,7 @@ scheduler:
 method:
   name: "ppoconfig" # Name of RL method config
   num_rollouts: 128 # Number of rollouts to collect per epoch
-  chunk_size: 128 # Number of rollouts to collect in one loop of orchestrator
+  chunk_size: 128 # Number of rollouts to collect in one loop
   ppo_epochs: 4 # Number of ppo epochs
   init_kl_coef: 0.2 # init kl coefficient
   target: 6 # target kl coefficient, set None for fixed kl coef

docs/source/index.rst

@@ -14,7 +14,6 @@ currently supports training using PPO or ILQL for models up to 20B using Acceler
 
    data
    models
-   orchestrator
    configs
    pipeline
    examples

docs/source/orchestrator.rst

@@ -4,7 +4,7 @@ Pipelines
 ************************
 
 Pipelines are how you read from a dataset with trlX. Rollout stores are how models store experiences created
-for them by the orchestrator. It is these experiences in their rollout store that they are trained on.
+for them. It is these experiences in their rollout store that they are trained on.
 
 **General**
 

docs/source/pipeline.rst

@@ -8,7 +8,6 @@ train:
   eval_interval: 16
 
   pipeline: "PromptPipeline"
-  orchestrator: "PPOOrchestrator"
   trainer: "AcceleratePPOTrainer"
 
 model: