refactor: make ilql respect the config (#22)

Author: maxreciprocate

examples/ilql_randomwalks.py

@@ -95,12 +95,12 @@ def reward_fn(samples):
 
         return rewards
 
-    gpt_config_or_path = GPT2Config(
+    config.model.model_path = GPT2Config(
         n_layer=4, n_embd=144, vocab_size=logit_mask.shape[0]
     )
 
     model = ILQLModel(
-        config=config, gpt_config_or_path=gpt_config_or_path, logit_mask=logit_mask
+        config=config, logit_mask=logit_mask
     )
 
     orch = OfflineOrchestrator(

examples/ilql_sentiments.py

@@ -1,52 +1,38 @@
 import math
-from typing import Callable, Iterable, List
+from typing import List
 
 import numpy as np
 import torch
 from datasets import load_dataset
-from tqdm import tqdm
 from transformers import AutoTokenizer, pipeline
 
 from trlx.data.configs import TRLConfig
 from trlx.model.accelerate_ilql_model import ILQLModel
 from trlx.orchestrator.offline_orchestrator import OfflineOrchestrator
 
-
-def batch_map(fn: Callable, xs: Iterable, bsize: int, desc=None):
-    out = []
-    for ind in tqdm(range(math.ceil(len(xs) / bsize)), desc=desc, disable=not desc):
-        batch = xs[ind * bsize : min(len(xs), (ind + 1) * bsize)]
-        out.extend(fn(batch))
-
-    return out
-
-
 if __name__ == "__main__":
     config = TRLConfig.load_yaml("configs/ilql_config.yml")
     sentiment_pipe = pipeline(
         "sentiment-analysis", "lvwerra/distilbert-imdb", device=torch.device(0)
     )
 
-    gpt_config_or_path = "gpt2"
-    tokenizer = AutoTokenizer.from_pretrained(gpt_config_or_path)
+    tokenizer = AutoTokenizer.from_pretrained(config.model.tokenizer_path)
     tokenizer.pad_token = tokenizer.eos_token
 
     def reward_fn(samples: List[str]) -> List[float]:
         if isinstance(samples[0], torch.Tensor):
             samples = tokenizer.batch_decode(samples, skip_special_tokens=True)
 
-        desc = "sentiment pipeline" if len(samples) > 1024 else None
-        sentiments = batch_map(
-            lambda batch: sentiment_pipe(batch), samples, bsize=1024, desc=desc
-        )
-        return [
-            1 - s["score"] if s["label"] == "NEGATIVE" else s["score"]
-            for s in sentiments
-        ]
-
-    model = ILQLModel(
-        config=config, gpt_config_or_path=gpt_config_or_path, tokenizer=tokenizer
-    )
+        sent_kwargs = {
+            "return_all_scores": True,
+            "function_to_apply": None,
+            "batch_size": 1024,
+        }
+        pipe_outputs = sentiment_pipe(samples, **sent_kwargs)
+        scores = torch.tensor([output[1]["score"] for output in pipe_outputs])
+        return scores
+
+    model = ILQLModel(config=config, tokenizer=tokenizer)
 
     n_prompts = 128
     eval_prompts = torch.tensor([model.tokenizer.bos_token_id] * n_prompts).view(

trlx/model/accelerate_ilql_model.py

@@ -4,13 +4,13 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
+import wandb
 from accelerate import Accelerator
 from torch.utils.data import DataLoader
 from transformers import AutoConfig, AutoTokenizer
 
-import wandb
 from trlx.model import BaseRLModel, register_model
-from trlx.model.nn.ilql_models import QVModel
+from trlx.model.nn.ilql_models import CausalLMWithValueHeads
 from trlx.pipeline.offline_pipeline import (OfflinePipeline,
                                             OfflineRolloutStorage)
 from trlx.utils import Clock, rampup_decay, safe_mkdir, topk_mask
@@ -24,20 +24,23 @@ class ILQLModel(BaseRLModel):
     def __init__(
         self,
         config,
-        gpt_config_or_path,
         tokenizer=None,
         logit_mask=None,
         train_mode=True,
     ):
         super().__init__(config, train_mode)
 
-        self.model = QVModel(gpt_config_or_path, config.method)
+        self.model = CausalLMWithValueHeads(
+            config.model.model_path,
+            params=config.method,
+            num_layers_unfrozen=config.model.num_layers_unfrozen,
+        )
         self.max_length = config.train.gen_size
-        self.tokenizer = tokenizer
+
         self.logit_mask = logit_mask
+        self.tokenizer = tokenizer
 
         self.accelerator = Accelerator(log_with="wandb")
-        self.accelerator.print(os.environ)
 
         if WORLD_SIZE > 1:
             torch.distributed.barrier(device_ids=[LOCAL_RANK])
@@ -121,7 +124,7 @@ def learn(self):
                     samples = self.accelerator.gather(torch.vstack(all_samples))
 
                     if self.accelerator.is_main_process:
-                        rewards = torch.tensor(self.reward_fn(samples), dtype=float)
+                        rewards = torch.as_tensor(self.reward_fn(samples), dtype=float)
                         reward = rewards.mean()
 
                         if self.stats_fn:
@@ -134,7 +137,7 @@ def learn(self):
                             )
                             pairs = list(zip(texts, rewards))
                             logs["samples"] = wandb.Table(
-                                columns=["samples", "reward"], rows=pairs[:16]
+                                columns=["samples", "reward"], rows=pairs[:128]
                             )
                             if os.environ.get("DEBUG"):
                                 print(

trlx/model/nn/ilql_models.py

@@ -7,25 +7,17 @@
 import accelerate
 import deepspeed
 import numpy as np
-import torch as th
+import torch
 import torch.nn.functional as F
 import transformers
 from accelerate.utils import compute_module_sizes
 from torch import nn, tensor
 from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig
 
 
-def topk_mask(xs: th.FloatTensor, k: int):
-    mintop = th.topk(xs, k)[0][:, -1].unsqueeze(-1)
-    return th.where(xs < mintop, -np.inf * th.ones_like(xs, dtype=xs.dtype), xs)
-
-
-class QVOutput(Tuple):
-    logits: th.FloatTensor
-    qs: th.FloatTensor
-    target_qs: th.FloatTensor
-    vs: th.FloatTensor
-    past_key_values: Tuple[th.FloatTensor]
+def topk_mask(xs: torch.FloatTensor, k: int):
+    mintop = torch.topk(xs, k)[0][:, -1].unsqueeze(-1)
+    return torch.where(xs < mintop, -np.inf * torch.ones_like(xs, dtype=xs.dtype), xs)
 
 
 def make_head(n_embd: int, out: int):
@@ -34,8 +26,12 @@ def make_head(n_embd: int, out: int):
     )
 
 
-class QVModel(nn.Module):
-    def __init__(self, config: Union[PretrainedConfig, str], params):
+class CausalLMWithValueHeads(nn.Module):
+    """This is a wrapper around huggingface AutoModelForCausalLM with two additional scalar heads"""
+
+    def __init__(
+        self, config: Union[PretrainedConfig, str], params, num_layers_unfrozen=-1
+    ):
         super().__init__()
 
         # enable zero3 init within from_pretrained
@@ -49,15 +45,26 @@ def __init__(self, config: Union[PretrainedConfig, str], params):
         else:
             self.gpt = AutoModelForCausalLM.from_pretrained(config)
 
-        for block in self.gpt.transformer.h:
-            block.requires_grad_(False)
-
-        if hasattr(self.gpt.config, "hidden_size"):
+        if hasattr(self.gpt, "gpt_neox"):
+            self.gpt.transformer = self.gpt.gpt_neox
+            self.gpt.lm_head = self.gpt_embed_out
             self.n_embd = self.gpt.config.hidden_size
+            gpt_blocks = self.gpt.gpt_neox.layers
         else:
             self.n_embd = self.gpt.config.n_embd
-        self.vocab_size = self.gpt.config.vocab_size
+            gpt_blocks = self.gpt.transformer.h
+
+        if num_layers_unfrozen == 0:
+            gpt_blocks_to_freeze = list(gpt_blocks)
+        elif num_layers_unfrozen > 0:
+            gpt_blocks_to_freeze = list(gpt_blocks)[:-num_layers_unfrozen]
+        else:
+            gpt_blocks_to_freeze = []
+
+        for m in gpt_blocks_to_freeze:
+            m.requires_grad_(False)
 
+        self.vocab_size = self.gpt.config.vocab_size
         self.v_head = make_head(self.n_embd, 1)
         self.q1_head = make_head(self.n_embd, self.vocab_size)
         self.target_q1_head = deepcopy(self.q1_head)
@@ -77,11 +84,7 @@ def __init__(self, config: Union[PretrainedConfig, str], params):
             self.target_q2_head.requires_grad_(False)
 
     def forward(self, **x):
-        if hasattr(self.gpt, "gpt_neox"):
-            out = self.gpt.gpt_neox(**x)
-        else:
-            out = self.gpt.transformer(**x)
-
+        out = self.gpt.transformer(**x)
         hs = out.last_hidden_state
 
         if self.two_qs:
@@ -91,12 +94,10 @@ def forward(self, **x):
             qs = self.q1_head(hs)
             target_qs = self.target_q1_head(hs)
 
-        if hasattr(self.gpt, "gpt_neox"):
-            logits = self.gpt.embed_out(hs)
-        else:
-            logits = self.gpt.lm_head(hs)
+        logits = self.gpt.lm_head(hs)
+        vs = self.v_head(hs)
 
-        return QVOutput((logits, qs, target_qs, self.v_head(hs), out.past_key_values))
+        return logits, qs, target_qs, vs, out.past_key_values
 
     def loss(self, batch):
         tokens = batch.input_ids.to(self.device)
@@ -115,7 +116,7 @@ def loss(self, batch):
 
             targetQ1 = target_qs[0][:, :-1].gather(-1, actions).squeeze(-1).detach()
             targetQ2 = target_qs[1][:, :-1].gather(-1, actions).squeeze(-1).detach()
-            targetQ = th.minimum(targetQ1, targetQ2)
+            targetQ = torch.minimum(targetQ1, targetQ2)
         else:
             Q = qs[:, :-1].gather(-1, actions).squeeze(-1)
             targetQ = target_qs[:, :-1].gather(-1, actions).squeeze(-1).detach()
@@ -212,7 +213,7 @@ def sync_target_q_heads(self):
         else:
             self._sync_target_q_heads(self.alpha)
 
-    @th.inference_mode()
+    @torch.inference_mode()
     def sample(
         self,
         query,
@@ -228,32 +229,32 @@ def sample(
         past_key_values = None
         tensors = defaultdict(list)
 
-        finished = th.zeros(input.shape[0], 1, dtype=th.long, device=query.device)
+        finished = torch.zeros(input.shape[0], 1, dtype=torch.long, device=query.device)
 
         for _ in range(max_length - 1):
             logits, _, target_qs, vs, past_key_values = self.forward(
                 input_ids=input, past_key_values=past_key_values
             )
 
             if self.two_qs:
-                qs = th.minimum(target_qs[0][:, -1], target_qs[1][:, -1])
+                qs = torch.minimum(target_qs[0][:, -1], target_qs[1][:, -1])
             else:
                 qs = target_qs[:, -1]
 
             logits = logits[:, -1]
 
             if logit_mask is not None:
-                logits[th.where(logit_mask[input[:, -1]])] = -np.inf
+                logits[torch.where(logit_mask[input[:, -1]])] = -np.inf
 
             adv = qs - vs[:, -1, :]
             pi = F.log_softmax(logits, -1)
             modpi = topk_mask(pi + beta * adv, top_k)
             ps = F.softmax(modpi / temperature, -1)
 
-            tokens = th.multinomial(ps, 1)
+            tokens = torch.multinomial(ps, 1)
             tokens = (1 - finished) * tokens + finished * eos_token_id
 
-            query = th.hstack((query, tokens))
+            query = torch.hstack((query, tokens))
 
             input = tokens
             finished = (tokens == eos_token_id).long()
@@ -265,21 +266,21 @@ def sample(
 
         stats = {}
         for name, xs in tensors.items():
-            xs = th.vstack(xs)
+            xs = torch.vstack(xs)
             stats.update(
                 {
                     f"{name}-min": xs.min(),
                     f"{name}-max": xs.max(),
                     f"{name}-std": xs.std(),
-                    f"{name}-avg": xs.mean(),
+                    f"{name}-mean": xs.mean(),
                 }
             )
 
         return query, stats
 
     @property
     def dummy_inputs(self):
-        return {"input_ids": th.ones(1, 1, device=self.gpt.device, dtype=th.long)}
+        return {"input_ids": torch.ones(1, 1, device=self.gpt.device, dtype=torch.long)}
 
     @property
     def device(self):

trlx/pipeline/offline_pipeline.py

@@ -1,4 +1,3 @@
-from functools import partial, reduce
 from typing import Callable, Iterable, Tuple
 
 import torch