Update KVCache maximum sequence length configuration in PPO recipe (#2412)

Author: SalmanMohammadi

recipes/configs/mistral/7B_full_ppo_low_memory.yaml

@@ -30,7 +30,7 @@ output_dir: /tmp/torchtune/mistral_7B/full_ppo_low_memory # /tmp may be deleted
 tokenizer:
   _component_: torchtune.models.mistral.mistral_tokenizer
   path: /tmp/Mistral-7B-Instruct-v0.2/tokenizer.model
-  max_seq_len: null
+  max_seq_len: 512
 
 # Dataset
 dataset:

recipes/ppo_full_finetune_single_device.py

@@ -234,14 +234,13 @@ def setup(self, cfg: DictConfig) -> None:
 
         # setup a context manager for enabling KV-cacheing during
         # trajectory generation if enabled in the config
-        self.cache_ctx_manager = lambda enable_kv_cache: (
+        self.cache_ctx_manager = lambda enable_kv_cache, decoder_max_seq_len: (
             local_kv_cache(
                 self._policy_model,
                 batch_size=self._forward_batch_size,
                 dtype=self._dtype,
-                decoder_max_seq_len=self._tokenizer.max_seq_len
-                + self._max_generated_tokens,
                 device=self._device,
+                decoder_max_seq_len=decoder_max_seq_len,
             )
             if enable_kv_cache
             else contextlib.nullcontext()
@@ -770,9 +769,12 @@ def generate_trajectory(self, input_ids: torch.Tensor) -> Trajectory:
             Trajectory: An instance of :class:`~torchtune.rlhf.Trajectory` comprising
                 the current trajectory.
         """
-
+        _, context_length = input_ids.shape
         # step 1: generate responses, and logits corresponding to the responses using the current policy
-        with self.cache_ctx_manager(self.enable_kv_cache):
+        with self.cache_ctx_manager(
+            self.enable_kv_cache,
+            decoder_max_seq_len=context_length + self._max_generated_tokens,
+        ):
             query_responses, logits = generation.generate(
                 model=self._policy_model,
                 prompt=input_ids,
@@ -782,7 +784,6 @@ def generate_trajectory(self, input_ids: torch.Tensor) -> Trajectory:
                 pad_id=self._tokenizer.pad_id,
                 rng=self._rng,
             )
-        _, context_length = input_ids.shape
         responses = query_responses[:, context_length:].clone()
         query_response_padding_masks = query_responses != self._tokenizer.pad_id