fix(ppo): generalize and stage for api

Author: maxreciprocate

examples/ppo_sentiments.py

@@ -1,41 +1,25 @@
-from typing import List
+import trlx
 
-import torch
+from datasets import load_dataset
 from transformers import pipeline
 
-import wandb
-from trlx.data.configs import TRLConfig
-from trlx.model.accelerate_ppo_model import AcceleratePPOModel
-from trlx.orchestrator.ppo_orchestrator import PPOOrchestrator
-from trlx.pipeline.ppo_pipeline import PPOPipeline
-from trlx.utils.loading import get_model, get_orchestrator, get_pipeline
-
 if __name__ == "__main__":
-    cfg = TRLConfig.load_yaml("configs/ppo_config.yml")
-
-    sentiment_pipe = pipeline(
-        "sentiment-analysis", "lvwerra/distilbert-imdb", device=-1
+    sentiment_fn = pipeline("sentiment-analysis", "lvwerra/distilbert-imdb")
+
+    def reward_fn(samples):
+        outputs = sentiment_fn(samples, return_all_scores=True)
+        sentiments = [output[1]["score"] for output in outputs]
+        return sentiments
+
+    # Take few words off of movies reviews as prompts
+    imdb = load_dataset("imdb", split="train+test")
+    prompts = [" ".join(review.split()[:4]) for review in imdb["text"]]
+
+    model = trlx.train(
+        model_path="lvwerra/gpt2-imdb",
+        reward_fn=reward_fn,
+        prompts=prompts,
+        eval_prompts=["I don't know much about Hungarian underground"] * 64 + ["<|endoftext|>"] * 64
     )
 
-    def reward_fn(samples: List[str]):
-        sent_kwargs = {
-            "return_all_scores": True,
-            "function_to_apply": None,
-            "batch_size": cfg.method.chunk_size,
-        }
-        pipe_outputs = sentiment_pipe(samples, **sent_kwargs)
-        scores = torch.tensor([output[1]["score"] for output in pipe_outputs])
-        return scores
-
-    model: AcceleratePPOModel = get_model(cfg.model.model_type)(cfg)
-    if model.accelerator.is_main_process:
-        wandb.watch(model.model)
-
-    pipeline: PPOPipeline = get_pipeline(cfg.train.pipeline)(model.tokenizer, cfg)
-    orch: PPOOrchestrator = get_orchestrator(cfg.train.orchestrator)(
-        model, pipeline, reward_fn=reward_fn, chunk_size=cfg.method.chunk_size
-    )
-    orch.make_experience(cfg.method.num_rollouts)
-    model.learn()
 
-    print("DONE!")

trlx/model/accelerate_base_model.py

@@ -26,11 +26,16 @@ class AccelerateRLModel(BaseRLModel):
     def __init__(self, config, train_mode=True):
         super().__init__(config, train_mode)
 
-        self.store = rollout_storage  # Need to pass in rollout_storage to be loaded into accelerate object
+        self.accelerator = Accelerator(log_with="wandb")
+
+        if int(os.environ.get("WORLD_SIZE", 1)) > 1:
+            torch.distributed.barrier(device_ids=[int(os.environ.get("LOCAL_RANK", 0))])
+        else:
+            torch.random.manual_seed(1000)
+
+        # Retrieves model equipped for ppo, ilql, etc
+        self.model = self.get_arch(self.config)
 
-        self.model = self.get_arch(
-            self.config
-        )  # Retrieves model equipped for ppo, ilql, etc
         if self.config.model.num_layers_unfrozen > 0:
             for block in self.model.gpt.transformer.h[:-self.config.model.num_layers_unfrozen]:
                 for parameter in block.parameters():
@@ -43,19 +48,13 @@ def __init__(self, config, train_mode=True):
         else:
             self.tokenizer = None
 
-        self.max_length = config.train.gen_size
         config_dict = self.config.to_dict()
         if self.config.train.accelerate_config_path != "":
             with open(self.config.train.accelerate_config_path, mode="r") as file:
                 accelerate_config = yaml.safe_load(file)
             config_dict.update(accelerate_config)
 
-        self.accelerator = Accelerator(log_with="wandb")
-
-        if int(os.environ.get("WORLD_SIZE", 1)) > 1:
-            torch.distributed.barrier(device_ids=[int(os.environ.get("LOCAL_RANK", 0))])
-        else:
-            torch.random.manual_seed(1000)
+        self.max_length = config.train.gen_size
 
         if self.accelerator.is_main_process:
             self.accelerator.init_trackers(
@@ -97,33 +96,21 @@ def tokenize(self, text: Iterable[str]):
         )
 
     def act(
-        self, data: PromptBatch
+        self, prompts
     ) -> Tuple[
         TensorType["chunk_size", "input_length"],
         TensorType["chunk_size", "gen_size"],
         Iterable[str],
     ]:
-        query_tensors = data.tokens.to(
-            self.accelerator.device
-        )  # [B, N] #TODO(dahoas): This may need to be changed
         with torch.no_grad():
-            # TODO(dahoas): swap this out for custom generate to if this fixes issue
-            _ = self.model(
-                self.dummy_input.to(self.accelerator.device)
-            )  # Dummy pass to make things play nice with accelerate
-            # Removed synced gpus
-            response = self.model.generate(
-                query_tensors,
-                pad_token_id=self.tokenizer.eos_token_id,
-                **self.config.method.gen_kwargs,
+            samples = self.model.generate(
+                **prompts,
+                pad_token_id=self.tokenizer.pad_token_id,
+                **self.config.method.gen_kwargs
             )
-            response_tensors = response[
-                :,
-                query_tensors.size()[1] : query_tensors.size()[1]
-                + self.config.train.gen_size,
-            ]
-        response_text = self.tokenizer.batch_decode(response_tensors)
-        return query_tensors, response_tensors, response_text
+
+        texts = self.tokenizer.batch_decode(samples, skip_special_tokens=True)
+        return prompts.input_ids, samples[:, prompts.input_ids.shape[1]:], texts
 
     @torch.inference_mode()
     def sample(self, prompts: PromptBatch, gen_kwargs: dict) -> Iterable[str]:

trlx/model/accelerate_ilql_model.py

@@ -12,7 +12,7 @@
 import wandb
 from trlx.model import BaseRLModel, register_model
 from trlx.model.nn.ilql_models import CausalLMWithValueHeads
-from trlx.pipeline.offline_pipeline import (OfflinePipeline,
+from trlx.pipeline.offline_pipeline import (PromptPipeline,
                                             OfflineRolloutStorage)
 
 from .accelerate_base_model import AccelerateRLModel
@@ -111,23 +111,6 @@ def learn(self):
                                 columns=["samples", *metrics.keys()], rows=rows
                             )
 
-                            metric_time = time()
-                            metrics = self.metric_fn(samples)
-                            metric_time = time() - metric_time
-                            logs.update({"metric_time": metric_time})
-
-                            mean_metrics = {
-                                f"metrics/{k}/{beta}": torch.as_tensor(xs).mean(-1)
-                                for k, xs in metrics.items()
-                            }
-                            logs.update(tensor_stats)
-                            logs.update(mean_metrics)
-
-                            rows = list(zip(samples, *metrics.values()))
-                            logs[f"samples/{beta}"] = wandb.Table(
-                                columns=["samples", *metrics.keys()], rows=rows
-                            )
-
                             print(rows[0])
                             print(mean_metrics)
 

trlx/model/accelerate_ppo_model.py

@@ -33,8 +33,6 @@ def update(self, current, n_steps):
         mult = 1 + proportional_error * n_steps / self.horizon
         self.value *= mult
 
-# Cell
-
 class FixedKLController:
     """Fixed KL controller."""
     def __init__(self, kl_coef):
@@ -48,20 +46,21 @@ class AcceleratePPOModel(AccelerateRLModel):
     def __init__(self, config, train_mode=True):
         super().__init__(config, train_mode)
 
-        self.store = PPORolloutStorage()
+        self.store = PPORolloutStorage(self.tokenizer.pad_token_id)
 
         rollout_loader = self.store.create_loader(
             self.config.train.batch_size, shuffle=True
         )
+
         self.model, self.opt, self.scheduler, rollout_loader = self.accelerator.prepare(
             self.model, self.opt, self.scheduler, rollout_loader
         )
-        self.store.clear_history()
 
         self.dummy_input = self.tokenize("dummy input")[
             "input_ids"
         ]  # Hack to make acclerate distributed work with model generation
 
+        self.store.clear_history()
         if config.method.target is not None:
             self.kl_ctl = AdaptiveKLController(
                                                 config.method.init_kl_coef,
@@ -78,6 +77,7 @@ def get_arch(self, config: TRLConfig):
     def loss(
         self, query_tensors, response_tensors, all_logprobs, all_values, all_rewards
     ):
+
         lastgaelam = 0
         advantages_reversed = []
         gen_len = response_tensors.shape[1]
@@ -99,7 +99,11 @@ def loss(
         advantages = advantages.detach()
 
         all_tokens = torch.cat((query_tensors, response_tensors), dim=1)
-        logits, _, vpred = self.model(all_tokens)
+        attention_mask = all_tokens.not_equal(self.tokenizer.pad_token_id).long()
+        position_ids = attention_mask.cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask.eq(0), 0)
+
+        logits, _, vpred = self.model(all_tokens, attention_mask, position_ids=position_ids)
         logprob = logprobs_from_logits(logits[:, :-1, :], all_tokens[:, 1:])
 
         # only the generation part of the values/logprobs is needed
@@ -111,9 +115,12 @@ def loss(
             all_values + self.config.method.cliprange_value,
         )
 
+        vf_mask = attention_mask[:, -gen_len-1:-1]
+        pg_mask = attention_mask[:, -gen_len:]
+
         vf_losses1 = (vpred - returns) ** 2
         vf_losses2 = (vpredclipped - returns) ** 2
-        vf_loss = 0.5 * torch.mean(torch.max(vf_losses1, vf_losses2))
+        vf_loss = 0.5 * torch.sum(torch.max(vf_losses1, vf_losses2) * vf_mask) / vf_mask.sum()
 
         kl = logprob - all_logprobs
         # Record mean_kl for kl coef adjustment
@@ -127,57 +134,58 @@ def loss(
             1.0 + self.config.method.cliprange,
         )
 
-        pg_loss = torch.mean(torch.max(pg_losses, pg_losses2))
+        pg_loss = torch.sum(torch.max(pg_losses, pg_losses2) * pg_mask) / pg_mask.sum()
 
         model_loss = pg_loss + self.config.method.vf_coef * vf_loss
         return model_loss, pg_loss, vf_loss
 
     def post_epoch_callback(self):
-        # TODO(dahoas): are experiences being made for dataloaders on each process or same dataloader
         self.epoch += 1
         self.store.clear_history()
         self.orch.make_experience(
             self.config.method.num_rollouts, self.iter_count
         )  # Collect more rollouts for training
 
     def post_backward_callback(self):
-        batch = self.logs["batch"]
         # Update kl_coefficient
         self.kl_ctl.update(self.mean_kl ,self.config.train.batch_size)
-        # Run evaluation
+
+        all_samples = []
+        for prompts in self.eval_dataloader:
+            query, response, _ = self.act(prompts)
+            pad_token = self.tokenizer.eos_token_id if self.tokenizer else 0
+            samples = torch.hstack((query, response))
+            all_samples.append(F.pad(samples, (0, self.max_length-samples.shape[1]), value=pad_token))
+
+        samples = self.accelerator.gather(torch.vstack(all_samples))
+
         if self.accelerator.is_main_process:
-            if (
-                self.iter_count % self.config.train.eval_interval == 0
-                or self.iter_count <= self.config.method.ppo_epochs
-            ):
-                text = self.tokenizer.batch_decode(batch.query_tensors)
-                eval_batch: PromptBatch = PromptBatch(
-                    text=text, tokens=batch.query_tensors
-                )
-                query_tensors, response_tensors, response_text = self.act(eval_batch)
-                gen_texts = [q + r for q, r in zip(eval_batch.text, response_text)]
-                scores = self.orch.score(gen_texts)
-                mean_score = torch.mean(scores).item()
-                rows = list(zip(gen_texts, scores.tolist()))
-                stats = {
-                    "mean_score": mean_score,
-                    "responses": wandb.Table(columns=["response", "score"], rows=rows),
-                    "pg_loss": self.logs["pg_loss"],
-                    "vf_loss": self.logs["vf_loss"],
-                    "kl_coef": self.kl_ctl.value,
-                }
-                self.accelerator.log(stats, step=self.iter_count)
-                self.accelerator.print(
-                    "Step: {}, Mean score: {}, pg_loss: {}, vf_loss: {}, kl_coef: {}".format(
-                        self.iter_count, mean_score, stats["pg_loss"], stats["vf_loss"], self.kl_ctl.value,
-                    )
+            samples = self.tokenizer.batch_decode(samples, skip_special_tokens=True)
+            scores = self.orch.score(samples)
+            mean_score = torch.mean(torch.as_tensor(scores)).item()
+            rows = list(zip(samples, scores))
+            stats = {
+                "mean_score": mean_score,
+                "responses": wandb.Table(columns=["response", "score"], rows=rows),
+                "pg_loss": self.logs["pg_loss"],
+                "vf_loss": self.logs["vf_loss"],
+                "kl_coef": self.kl_ctl.value,
+            }
+
+            self.accelerator.log(stats, step=self.iter_count)
+            self.accelerator.print(
+                "Step: {}, Mean score: {}, pg_loss: {}, vf_loss: {}, kl_coef: {}".format(
+                    self.iter_count, mean_score, stats["pg_loss"], stats["vf_loss"], self.kl_ctl.value,
                 )
+            )
 
     def learn(self):
+        self.eval_dataloader = self.eval_pipeline.create_loader(self.config.train.batch_size)
+
         rollout_loader = self.store.create_loader(
             self.config.train.batch_size, shuffle=True
         )
-        rollout_loader = self.accelerator.prepare(rollout_loader)
+        rollout_loader, self.eval_dataloader = self.accelerator.prepare(rollout_loader, self.eval_dataloader)
 
         self.iter_count = 0
         self.epoch = 0
@@ -204,8 +212,7 @@ def learn(self):
                         "batch": batch,
                         "rewards": rewards,
                     }
-                    # self.post_backward_callback()
-                    # exit()
+
                     self.opt.zero_grad()
                     self.accelerator.backward(loss)
                     self.opt.step()

trlx/model/nn/ilql_models.py

@@ -137,14 +137,16 @@ def loss(self, batch):
 
         _V = vs[:, :-1].squeeze()
         V = vs[:, 1:].squeeze() * dones[:, 1:]
-        Q_ = rewards + self.gamma * V
+        Q_ = rewards + self.gamma * V.detach()
 
         if self.two_qs:
-            loss_q1 = ((Q1 - Q_.detach()) * terminal_mask).pow(2).sum() / n_nonterminal
-            loss_q2 = ((Q2 - Q_.detach()) * terminal_mask).pow(2).sum() / n_nonterminal
+            loss_q1 = ((Q1 - Q_) * terminal_mask).pow(2).sum() / n_nonterminal
+            loss_q2 = ((Q2 - Q_) * terminal_mask).pow(2).sum() / n_nonterminal
             loss_q = loss_q1 + loss_q2
         else:
-            loss_q = ((Q - Q_.detach()) * terminal_mask).pow(2).sum() / n_nonterminal
+            loss_q = ((Q - Q_) * terminal_mask).pow(2).sum() / n_nonterminal
+
+        targetQ = targetQ.detach()
 
         loss_v = (
             (