added iw

efficient_reasoning/extras/preemptive_vllm_client.py

@@ -21,22 +21,37 @@ def __init__(
         sampler: Sampler,            # <-- take a Sampler
         batch_size: int,
         num_generations: int,
+        num_iterations: int,
+        per_device_train_batch_size: int,
+        preemptive_steps: int = 0,
+        iw: bool = False,
+        iw_first_step: bool = True,
     ):
+        self.iw = iw
+        self.iw_first_step = iw_first_step
         # Underlying multi‐client
         self.multi_client = MultiVLLMClient(server_configs)
 
         # How many macro‐batches to pre‐fetch
         self.grad_accum_steps = gradient_accumulation_steps
+
+        self.preemptive_steps = preemptive_steps
 
         # materialize the prompts column into a list of str
         self.train_prompts = train_dataset["prompt"]
 
         # materialize the sampler into a list of indices
-        self.sampler_cycle = itertools.cycle(list(sampler)[::num_generations])
-        # for _ in range(7744):
+        # print(f"Sampler: {list(sampler)[:300]}")
+        list_temp = list(sampler)[::num_generations] # remove the extra num_generations of samples (can't use training dataset because of potentially shuffling)
+        # print(f"Sampler after removing extra num_generations: {list_temp[:300]}")
+        list_final = [x for idx, x in enumerate(list_temp) if (idx // (per_device_train_batch_size * gradient_accumulation_steps)) % num_iterations == 0] # remove the extra num_iterations of samples
+        # print(f"Sampler after removing extra num_iterations: {list_final[:300]}")
+        self.sampler_cycle = itertools.cycle(list_final)
+        # print(f"Sampler after removing extra num_iterations: {list_final[:1200]}")
+        # for _ in range(8000):
         #     next(self.sampler_cycle) 
-        print(list(sampler))
-        print(list(sampler)[::num_generations])
+        # print(list(sampler))
+        # print(list(sampler)[::num_generations])
         self.batch_size = batch_size
 
         # Buffer to hold all pre‐fetched token‐ID lists
@@ -49,7 +64,7 @@ def _next_batch_prompts(self) -> List[str]:
         idxs = [next(self.sampler_cycle) for _ in range(self.batch_size)]
         return [self.train_prompts[i] for i in idxs]
 
-    def generate(self, prompts: List[str], n: int, eval_only: bool = False, **sampling_kwargs) -> List[List[int]]:
+    def generate(self, prompts: List[str], n: int, read_from_file: bool = False,  eval_only: bool = False, **sampling_kwargs) -> List[List[int]]:
         """
         On the very first call of each gradient‐accumulation cycle (or if
         sampling hyperparameters change), grab
@@ -60,57 +75,101 @@ def generate(self, prompts: List[str], n: int, eval_only: bool = False, **sampli
         On each call, slice out the next `batch_size * n` token‐ID lists
         and return them.
         """
-        if eval_only:
-            print("Eval only mode: skipping pre‐sampling.")
-            return self.multi_client.generate(prompts, n=n, **sampling_kwargs)
-        else:
-            # 1) Refill if at start of cycle, or if sampling settings have changed
-            if self.current_cycle_calls == 0:
-                self.current_sampling_kwargs = sampling_kwargs.copy()
-
-                # Gather all prompts for the cycle
-                all_prompts: List[str] = []
-                for _ in range(self.grad_accum_steps):
-                    batch_prompts = self._next_batch_prompts()
-                    all_prompts.extend(batch_prompts)
-
-                # Sanity check: the very first macro‐batch should match `prompts`
-                if all_prompts[: len(prompts)] != prompts:
-                    print(all_prompts[: len(prompts)])
-                    print(prompts)
-                    raise ValueError(
-                        "Sampler out of sync: fetched prompts != provided prompts"
+        if self.iw:
+            if self.iw_first_step:
+                self.iw_first_step = False
+
+                if not read_from_file:
+                    all_samples = self.multi_client.generate(
+                        self.train_prompts,
+                        n=n,
+                        **sampling_kwargs
                     )
+
+                    with open("buffer.txt", "w") as f:
+                        for sample in all_samples:
+                            f.write(f"{sample}\n")
 
-                # One big call for everything
-                all_samples = self.multi_client.generate(
-                    all_prompts,
-                    n=n,
-                    **sampling_kwargs
-                )
-                expected = len(all_prompts) * n
-                if len(all_samples) != expected:
+                else: 
+                    # Read from file
+                    all_samples = []
+                    with open("buffer.txt", "r") as f:
+                        for line in f:
+                            line = line.strip()
+                            if line:  # Skip empty lines
+                                sample = eval(line)
+                                all_samples.append(sample)
+
+                if len(all_samples) != 12000 * n:
                     raise ValueError(
-                        f"Expected {expected} token lists, got {len(all_samples)}"
+                        f"Expected {12000 * n} generation lists, got {len(all_samples)}"
                     )
 
-                # Load into buffer
                 self.buffer.clear()
                 self.buffer.extend(all_samples)
-                self.current_cycle_calls = 0
 
-            # 2) Now slice out this step’s completions
             slice_size = len(prompts) * n
             if len(self.buffer) < slice_size:
-                raise RuntimeError("Buffer underflow: not enough samples")
+                raise RuntimeError(f"Buffer underflow: not enough samples: remaining buffer size {len(self.buffer)} < requested slice size {slice_size}")
             next_batch = [self.buffer.popleft() for _ in range(slice_size)]
-            self.current_cycle_calls += 1
+            return next_batch
 
-            # 3) If we’ve completed the full micro‐batch cycle, reset
+        # 1) Refill if at start of cycle, or if sampling settings have changed
+        if self.current_cycle_calls == 0:
+            self.current_sampling_kwargs = sampling_kwargs.copy()
+
+            # Gather all prompts for the cycle
+            all_prompts: List[str] = []
+            if self.preemptive_steps > 0:
+                for _ in range(self.preemptive_steps):
+                    batch_prompts = self._next_batch_prompts()
+                    all_prompts.extend(batch_prompts)
+            else:
+                for _ in range(self.grad_accum_steps):
+                    batch_prompts = self._next_batch_prompts()
+                    all_prompts.extend(batch_prompts)
+
+            # Sanity check: the very first macro‐batch should match `prompts`
+            if all_prompts[: len(prompts)] != prompts:
+                print("Fetched prompts:", all_prompts[: len(prompts)])
+                print("Provided prompts:", prompts)
+                raise ValueError(
+                    "Sampler out of sync: fetched prompts != provided prompts"
+                )
+
+            # One big call for everything
+            all_samples = self.multi_client.generate(
+                all_prompts,
+                n=n,
+                **sampling_kwargs
+            )
+            expected = len(all_prompts) * n
+            if len(all_samples) != expected:
+                raise ValueError(
+                    f"Expected {expected} token lists, got {len(all_samples)}"
+                )
+
+            # Load into buffer
+            self.buffer.clear()
+            self.buffer.extend(all_samples)
+            self.current_cycle_calls = 0
+
+        # 2) Now slice out this step’s completions
+        slice_size = len(prompts) * n
+        if len(self.buffer) < slice_size:
+            raise RuntimeError("Buffer underflow: not enough samples")
+        next_batch = [self.buffer.popleft() for _ in range(slice_size)]
+        self.current_cycle_calls += 1
+
+        # 3) If we’ve completed the full micro‐batch cycle, reset
+        if self.preemptive_steps > 0:
+            if self.current_cycle_calls >= self.preemptive_steps:
+                self.current_cycle_calls = 0
+        else:
             if self.current_cycle_calls >= self.grad_accum_steps:
                 self.current_cycle_calls = 0
 
-            return next_batch
+        return next_batch
 
     def reset_buffer(self):
         """Clear out everything and restart the sampler iterator."""

efficient_reasoning/grpo_config.py

@@ -457,6 +457,19 @@ class GRPOConfig(TrainingArguments):
     # use old model
     use_old_model: bool = field(default=True, metadata={"help": "Whether to use old model."})
 
+    # preemptive steps
+    preemptive_steps: int = field(default=0, metadata={"help": "Number of preemptive steps."})
+
+    # importance weighting
+    iw: bool = field(
+        default=False,
+        metadata={
+            "help": "Controls whether to use importance weighting (IW) or not. "
+        },
+    )
+
+    read_from_file: bool = field(default=False, metadata={"help": "When doing iw, whether to directly read generations from pre-generated rollouts in a file (buffer.txt)"})
+
     def __post_init__(self):
         super().__post_init__()