Fix context overflow (#131)

* feat(ppo_models): context truncating generation

* feat(base_model): segragate prompts and responses in logging

* fix(ppo_model): truncate left padded tokens

* revert(ppo_models): remove context overflowing generate()

* feat(configs): add max_new_tokens

* fix(pipeline): truncate prompts

* chore(base_model): remove whole samples from logging

* chore(configs): update the rest of configs

* fix(configs): update program synthesis config

* fix(base_model): prompts sizes

* revert(config): emulate old ppo_sentiments behavior

* fix(ppo): unequal generation lengths

* chore(ppo): put indexing on cpu

* revert(configl): old ilql_sentiments behavior

* fix(ppo): unzero clipfrac

* merge(configs): delete old options

* revert(config): old ppo_sentiment behavior

* refactor(ppo_orchestrator): remove unused ref_logprobs

* fix(base_model): pin rewards to single precision

* refactor(ppo): rename padding percentage

* feat(wandb): add git branch name to tags

* refactor(wandb): logging name hierarchy

* revert(wandb): merge tags into a single string

Author: maxreciprocate

configs/ilql_config.yml

@@ -40,5 +40,9 @@ method:
   awac_scale: 1
   alpha: 0.001
   steps_for_target_q_sync: 5
-  betas: [4]
   two_qs: true
+  gen_kwargs:
+    max_new_tokens: 56
+    top_k: 20
+    beta: 4
+    temperature: 1.0

configs/ppo_config.yml

@@ -1,5 +1,5 @@
 train:
-  seq_length: 48
+  seq_length: 1024
   epochs: 100
   total_steps: 10000
   batch_size: 128
@@ -48,8 +48,7 @@ method:
   ref_std: null
   cliprange_reward: 10
   gen_kwargs:
-    max_length: 48
-    min_length: 48
-    top_k: 0.0
+    max_new_tokens: 40
+    top_k: 0
     top_p: 1.0
     do_sample: True

configs/ppo_gptj.yml

@@ -48,8 +48,7 @@ method:
   ref_std: null
   cliprange_reward: 10
   gen_kwargs:
-    max_length: 48
-    min_length: 48
+    max_new_tokens: 48
     top_k: 0.0
     top_p: 0.7
     do_sample: True

examples/experiments/grounded_program_synthesis/configs/trlx_ppo_config.yml

@@ -48,9 +48,8 @@ method:
   ref_mean: null
   ref_std: null
   gen_kwargs:
-    max_length: 256
-    min_length: 48
-    top_k: 0.0
+    max_new_tokens: 256
+    top_k: 0
     top_p: 0.7
     do_sample: True
     temperature: 0.5

examples/randomwalks/configs/ilql_randomwalks.yml

@@ -40,5 +40,9 @@ method:
   awac_scale: 1
   alpha: 0.1
   steps_for_target_q_sync: 5
-  betas: [100]
   two_qs: true
+  gen_kwargs:
+    max_new_tokens: 9
+    top_k: 1
+    beta: 100
+    temperature: 1.0

examples/randomwalks/configs/ppo_randomwalks.yml

@@ -48,8 +48,7 @@ method:
   ref_std: null
   cliprange_reward: 1
   gen_kwargs:
-    max_length: 10
-    min_length: 2
+    max_new_tokens: 9
     top_k: 0.0
     top_p: 1.0
     do_sample: True

trlx/model/accelerate_base_model.py

@@ -147,6 +147,7 @@ def evaluate(self):
         """Samples model on `eval_prompts`, logs stats with `reward_fn` or `metric_fn` if provided"""
         stats = {}
         all_samples = []
+        prompts_sizes = []
         generate_time = time()
         for prompts in self.eval_dataloader:
             if isinstance(prompts, torch.Tensor):
@@ -165,34 +166,54 @@ def evaluate(self):
                     value=pad_token,
                 )
             )
-        stats["generate_time"] = time() - generate_time
+            sizes = torch.tensor(prompts.input_ids.shape[1]).repeat(
+                len(prompts.input_ids)
+            )
+            prompts_sizes.append(sizes.to(samples.device))
+
+        stats["time/generate"] = time() - generate_time
 
         samples = self.accelerator.gather(torch.vstack(all_samples))
+        prompts_sizes = self.accelerator.gather(torch.hstack(prompts_sizes))
 
         if self.accelerator.is_main_process:
             if self.tokenizer:
-                samples = self.tokenizer.batch_decode(samples, skip_special_tokens=True)
+                str_samples = self.tokenizer.batch_decode(
+                    samples, skip_special_tokens=True
+                )
+
+                prompts, responses = [], []
+                for sample, prompt_size in zip(samples, prompts_sizes):
+                    prompts.append(sample[:prompt_size])
+                    responses.append(sample[prompt_size:])
+
+                str_prompts = self.tokenizer.batch_decode(
+                    prompts, skip_special_tokens=True
+                )
+                str_responses = self.tokenizer.batch_decode(
+                    responses, skip_special_tokens=True
+                )
 
-            if isinstance(samples[0], str):
-                columns_data = [samples]
+            if isinstance(str_samples[0], str):
+                columns_data = [str_prompts, str_responses]
             else:
                 columns_data = [samples.tolist()]
-            columns = ["samples"]
+            columns = ["prompt", "response"]
 
             # in online setting, compute the reward for validation
             if self.reward_fn:
-                rewards = torch.as_tensor(self.reward_fn(samples), dtype=torch.float)
+                rewards = torch.tensor(self.reward_fn(str_samples), dtype=torch.float)
                 mean_reward = rewards.mean()
                 columns.append("reward")
                 columns_data.append(rewards)
-                stats["mean_reward"] = mean_reward
+                stats["reward/mean"] = mean_reward
                 print(f"{mean_reward=}")
 
             # additionally log any other metrics
             if self.metric_fn:
                 metric_time = time()
-                metrics = self.metric_fn(samples)
-                stats["metric_time"] = time() - metric_time
+                metrics = self.metric_fn(str_samples)
+                stats["time/metric"] = time() - metric_time
 
                 mean_metrics = {
                     f"metrics/{k}": torch.as_tensor(xs).mean(-1)
@@ -258,8 +279,8 @@ def learn(self):
                     if self.iter_count % self.config.train.checkpoint_interval == 0:
                         self.save()
 
-                    stats["forward_time"] = forward_time
-                    stats["backward_time"] = backward_time
+                    stats["time/forward"] = forward_time
+                    stats["time/backward"] = backward_time
 
                     if self.iter_count % self.config.train.eval_interval == 0:
                         results = self.evaluate()

trlx/model/accelerate_ilql_model.py

@@ -32,6 +32,14 @@ def __init__(
 
         self.ilql: ILQLConfig = cast(ILQLConfig, config.method)
 
+        self.generate_kwargs = dict(
+            config.method.gen_kwargs,
+            max_length=self.max_length,
+            logit_mask=self.logit_mask,
+            eos_token_id=self.tokenizer.eos_token_id if self.tokenizer else 0,
+            pad_token_id=self.tokenizer.pad_token_id if self.tokenizer else 0,
+        )
+
     def get_arch(self, config):
         return CausalLMWithValueHeads(
             config.model.model_path,
@@ -87,11 +95,3 @@ def prepare_learning(self):
         self.n_updates_per_batch = 1
         self.total_steps = self.config.train.epochs * len(train_dataloader)
         self.total_steps = min(self.total_steps, self.config.train.total_steps)
-
-        self.generate_kwargs = {
-            "beta": self.config.method.betas[0],
-            "max_length": self.max_length,
-            "logit_mask": self.logit_mask,
-            "eos_token_id": self.tokenizer.eos_token_id if self.tokenizer else 0,
-            "pad_token_id": self.tokenizer.pad_token_id if self.tokenizer else 0,
-        }

trlx/model/accelerate_ppo_model.py

@@ -62,7 +62,9 @@ def get_model_inputs(
         query_tensors: TensorType["batch_size", "query_size"],
         response_tensors: TensorType["batch_size", "response_size"],
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        tokens = torch.cat((query_tensors, response_tensors), dim=1)
+        tokens = torch.cat((query_tensors, response_tensors), dim=1)[
+            :, -self.max_length :
+        ]
         attention_mask = (
             tokens.not_equal(self.tokenizer.pad_token_id).long().to(tokens.device)
         )
@@ -79,23 +81,30 @@ def loss(self, batch: PPORLBatch):
         old_values = batch.values.to(self.accelerator.device)
         old_rewards = batch.rewards.to(self.accelerator.device)
 
-        response_length = response_tensors.shape[-1]
+        response_length = old_rewards.shape[1]
+
         advantages, returns = self.config.method.get_advantages_and_returns(
             old_values, old_rewards, response_length
         )
 
         tokens, attention_mask, position_ids = self.get_model_inputs(
             query_tensors, response_tensors
         )
+
         logits, *_, values_pred = self.model(
             tokens, attention_mask=attention_mask, position_ids=position_ids
         )
+        values_pred = values_pred[:, :-1]
         logprobs = logprobs_from_logits(logits[:, :-1, :], tokens[:, 1:])
+        attention_mask = attention_mask[:, :-1]
+
         # Only the response part of the values/logprobs is needed
+        start = query_tensors.shape[1] - 1
+        end = start + response_length
         logprobs, values_pred, mask = (
-            logprobs[:, -response_length:],
-            values_pred[:, -response_length:],
-            attention_mask[:, -response_length:],
+            logprobs[:, start:end],
+            values_pred[:, start:end],
+            attention_mask[:, start:end],
         )
 
         loss, stats = self.config.method.loss(

trlx/model/nn/ilql_models.py

@@ -42,8 +42,8 @@ class ILQLConfig(MethodConfig):
     awac_scale: float
     alpha: float
     steps_for_target_q_sync: float
-    betas: Sequence[float]
     two_qs: bool
+    gen_kwargs: dict
 
     def heads(self, hidden_size: int, vocab_size: int):
         return ILQLHeads(self, hidden_size, vocab_size)
@@ -140,7 +140,6 @@ def forward(
         states_ixs: torch.Tensor = None,
         actions_ixs: torch.Tensor = None,
     ):
-
         if states_ixs is not None:
             states_hs = hs.gather(
                 dim=1, index=states_ixs.unsqueeze(-1).repeat(1, 1, hs.shape[-1])
@@ -260,7 +259,8 @@ def generate(
         position_ids=None,
         past_key_values=None,
         beta=1,
-        max_length=32,
+        max_new_tokens=32,
+        max_length=1024,
         temperature=1,
         top_k=20,
         logit_mask=None,
@@ -278,13 +278,12 @@ def generate(
             position_ids.masked_fill_(attention_mask.eq(0), 0)
 
         samples = input_ids.clone()
-        tensors = defaultdict(list)
-        n_new_tokens = max_length - input_ids.shape[1]
+        max_new_tokens = min(max_new_tokens, max_length - input_ids.shape[1])
 
         finished = torch.zeros(
             input_ids.shape[0], 1, dtype=torch.long, device=input_ids.device
         )
-        for _ in range(n_new_tokens):
+        for _ in range(max_new_tokens):
             out = self.forward(
                 input_ids=input_ids,
                 attention_mask=attention_mask,