partial merge

Author: bengioe

src/gflownet/config.py

@@ -112,6 +112,7 @@ class Config:
     num_workers: int = 0
     hostname: Optional[str] = None
     pickle_mp_messages: bool = False
+    mp_buffer_size: Optional[int] = 32 * 1024 ** 2  # 32Mb
     git_hash: Optional[str] = None
     overwrite_existing_exp: bool = True
     algo: AlgoConfig = AlgoConfig()

src/gflownet/data/sampling_iterator.py

@@ -12,6 +12,7 @@
 
 from gflownet.data.replay_buffer import ReplayBuffer
 from gflownet.envs.graph_building_env import GraphActionCategorical
+from gflownet.utils.multiprocessing_proxy import BufferPickler, SharedPinnedBuffer
 
 
 class SamplingIterator(IterableDataset):
@@ -44,7 +45,8 @@ def __init__(
         random_traj_prob: float = 0.0,
         hindsight_ratio: float = 0.0,
         init_train_iter: int = 0,
-        is_validation: bool = False
+        is_validation: bool = False,
+        mp_cfg = None,
     ):
         """Parameters
         ----------
@@ -110,6 +112,7 @@ def __init__(
         self.train_it = init_train_iter
         self.is_validation = is_validation
         self.do_validate_batch = False  # Turn this on for debugging
+        self.num_workers, _, self.mp_buffer_size = mp_cfg
 
         # Slightly weird semantics, but if we're sampling x given some fixed cond info (data)
         # then "offline" now refers to cond info and online to x, so no duplication and we don't end
@@ -125,6 +128,8 @@ def __init__(
         self.log = SQLiteLog()
         self.log_hooks: List[Callable] = []
 
+        self.setup_mp_buffers()
+
     def add_log_hook(self, hook: Callable):
         self.log_hooks.append(hook)
 
@@ -282,17 +287,21 @@ def __iter__(self):
                 # and sample replay_batch_size of them to add to the batch
 
                 # cond_info is a dict, so we need to convert it to a list of dicts
-                cond_info = [{k: v[i] for k, v in cond_info.items()} for i in range(num_offline + num_online)]
+                cond_info_ = [{k: v[i] for k, v in cond_info.items()} for i in range(num_offline + num_online)]
 
                 # push the online trajectories in the replay buffer and sample a new 'online' batch
                 for i in range(num_offline, len(trajs)):
+                    if not is_valid[i].item():
+                        continue
                     self.replay_buffer.push(
                         deepcopy(trajs[i]),
                         deepcopy(log_rewards[i]),
                         deepcopy(flat_rewards[i]),
-                        deepcopy(cond_info[i]),
+                        deepcopy(cond_info_[i]),
                         deepcopy(is_valid[i]),
                     )
+            if self.replay_buffer is not None and len(self.replay_buffer) > self.replay_buffer.warmup:
+                cond_info = [{k: v[i] for k, v in cond_info.items()} for i in range(num_offline + num_online)]
                 replay_trajs, replay_logr, replay_fr, replay_condinfo, replay_valid = self.replay_buffer.sample(
                     self.replay_batch_size
                 )
@@ -340,11 +349,11 @@ def __iter__(self):
 
             # TODO: need to change this for non-molecule environments
             try:
-                smiles = [Chem.MolToSmiles(self.ctx.graph_to_mol(traj["result"])) for traj in trajs]
+                smiles = [self.ctx.object_to_log_repr(traj["result"]) for traj in trajs]
             except:
                 smiles = [traj["result"].__repr__() for traj in trajs]
             # alternative: [traj["smi"] for traj in trajs]
-            yield batch, (smiles, flat_rewards)
+            yield self._maybe_put_in_mp_buffer((batch, (smiles, flat_rewards)))
 
     def validate_batch(self, batch, trajs):
         for actions, atypes in [(batch.actions, self.ctx.action_type_order)] + (
@@ -400,6 +409,18 @@ def log_generated(self, trajs, rewards, flat_rewards, cond_info):
 
         self.log.insert_many(data, data_labels)
 
+    def setup_mp_buffers(self):
+        if self.num_workers > 0 and self.mp_buffer_size:
+            self.result_buffer = [SharedPinnedBuffer(self.mp_buffer_size) for _ in range(self.num_workers)]
+        else:
+            self.mp_buffer_size = None
+
+    def _maybe_put_in_mp_buffer(self, batch):
+        if self.mp_buffer_size:
+            return (BufferPickler(self.result_buffer[self._wid]).dumps(batch), self._wid)
+        else:
+            return batch
+
 
 class SQLiteLog:
     def __init__(self, timeout=300):

src/gflownet/tasks/seh_frag.py

@@ -80,6 +80,7 @@ class SEHFragTrainer(StandardOnlineTrainer):
     def set_default_hps(self, cfg: Config):
         cfg.hostname = socket.gethostname()
         cfg.pickle_mp_messages = False
+        cfg.mp_buffer_size = 32 * 1024 ** 2  # 32Mb should be enough for this setup
         cfg.num_workers = 5
 
         cfg.opt.learning_rate = 1e-4

src/gflownet/trainer.py

@@ -19,7 +19,7 @@
 from gflownet.envs.graph_building_env import GraphActionCategorical, GraphBuildingEnv, GraphBuildingEnvContext
 from gflownet.envs.seq_building_env import SeqBatch
 from gflownet.utils.misc import create_logger
-from gflownet.utils.multiprocessing_proxy import mp_object_wrapper
+from gflownet.utils.multiprocessing_proxy import mp_object_wrapper, BufferUnpickler
 from gflownet.utils.misc import prepend_keys, average_values_across_dicts
 from gflownet.utils.metrics_final_eval import compute_metrics
 import wandb
@@ -219,7 +219,8 @@ def _wrap_for_mp(self, obj, send_to_device=False):
                 self.cfg.num_workers,
                 cast_types=(gd.Batch, GraphActionCategorical, SeqBatch),
                 pickle_messages=self.cfg.pickle_mp_messages,
-            )
+                sb_size=self.cfg.mp_buffer_size,
+            ).placeholder
             return placeholder, torch.device("cpu")
         else:
             return obj, self.device
@@ -248,6 +249,7 @@ def build_training_data_loader(self) -> DataLoader:
             random_action_prob=self.cfg.algo.train_random_action_prob,
             random_traj_prob=self.cfg.algo.train_random_traj_prob,
             hindsight_ratio=self.cfg.replay.hindsight_ratio,
+            mp_cfg=(self.cfg.num_workers, self.cfg.pickle_mp_messages, self.cfg.mp_buffer_size),
         )
         for hook in self.sampling_hooks:
             iterator.add_log_hook(hook)
@@ -258,7 +260,7 @@ def build_training_data_loader(self) -> DataLoader:
             persistent_workers=self.cfg.num_workers > 0,
             # The 2 here is an odd quirk of torch 1.10, it is fixed and
             # replaced by None in torch 2.
-            prefetch_factor=1 if self.cfg.num_workers else 2,
+            prefetch_factor=1 if self.cfg.num_workers else (None if torch.__version__.startswith('2') else 2),
             generator=g,
             worker_init_fn=seed_worker
         )
@@ -284,6 +286,7 @@ def build_validation_data_loader(self) -> DataLoader:
             stream=False,
             random_action_prob=self.cfg.algo.valid_random_action_prob,
             is_validation=True,
+            mp_cfg=(self.cfg.num_workers, self.cfg.pickle_mp_messages, self.cfg.mp_buffer_size),
         )
         for hook in self.valid_sampling_hooks:
             iterator.add_log_hook(hook)
@@ -292,7 +295,7 @@ def build_validation_data_loader(self) -> DataLoader:
             batch_size=None,
             num_workers=self.cfg.num_workers,
             persistent_workers=self.cfg.num_workers > 0,
-            prefetch_factor=1 if self.cfg.num_workers else 2,
+            prefetch_factor=1 if self.cfg.num_workers else (None if torch.__version__.startswith('2') else 2),
             generator=g,
             worker_init_fn=seed_worker
         )
@@ -322,6 +325,7 @@ def build_final_data_loader(self) -> DataLoader:
             hindsight_ratio=0.0,
             is_validation=True,
             # init_train_iter=self.cfg.num_training_steps,
+            mp_cfg=(self.cfg.num_workers, self.cfg.pickle_mp_messages, self.cfg.mp_buffer_size),
         )
         for hook in self.sampling_hooks:
             iterator.add_log_hook(hook)
@@ -330,11 +334,19 @@ def build_final_data_loader(self) -> DataLoader:
             batch_size=None,
             num_workers=self.cfg.num_workers,
             persistent_workers=self.cfg.num_workers > 0,
-            prefetch_factor=1 if self.cfg.num_workers else 2,
+            prefetch_factor=1 if self.cfg.num_workers else (None if torch.__version__.startswith('2') else 2),
             generator=g,
             worker_init_fn=seed_worker
         )
 
+    def _maybe_resolve_shared_buffer(self, batch, dl: DataLoader):
+        if dl.dataset.mp_buffer_size and isinstance(batch, (tuple, list)):
+            batch, wid = batch
+            batch = BufferUnpickler(dl.dataset.result_buffer[wid], batch, self.device).load()
+        elif isinstance(batch, (gd.Batch, SeqBatch)):
+            batch = batch.to(self.device)
+        return batch
+
     def train_batch(self, batch: gd.Batch, epoch_idx: int, batch_idx: int, train_it: int) -> Dict[str, Any]:
         try:
             loss, info = self.algo.compute_batch_losses(self.model, batch)
@@ -383,7 +395,8 @@ def run(self, logger=None):
         start = self.cfg.start_at_step + 1
         num_training_steps = self.cfg.num_training_steps
         logger.info("Starting training")
-        for it, (batch, _) in zip(range(start, 1 + num_training_steps), cycle(train_dl)):
+        for it, batch in zip(range(start, 1 + num_training_steps), cycle(train_dl)):
+            batch, _ = self._maybe_resolve_shared_buffer(batch, train_dl)
             epoch_idx = it // epoch_length
             batch_idx = it % epoch_length
             if self.replay_buffer is not None and len(self.replay_buffer) < self.replay_buffer.warmup:
@@ -404,7 +417,8 @@ def run(self, logger=None):
                 candidates_eval_infos = []
                 # for batch in valid_dl:
                 # validate on at least 10 batches
-                for valid_it, (batch, candidates_eval_info) in zip(range(10), cycle(valid_dl)):
+                for valid_it, batch in zip(range(10), cycle(valid_dl)):
+                    batch, candidates_eval_info = self._maybe_resolve_shared_buffer(batch, valid_dl)
                     # print("valid_it", valid_it)
                     candidates_eval_infos.append(candidates_eval_info)
                     metrics = self.evaluate_batch(batch.to(self.device), epoch_idx, batch_idx)
@@ -440,10 +454,11 @@ def run(self, logger=None):
         if num_final_gen_steps:
             gen_candidates_list = []
             logger.info(f"Generating final {num_final_gen_steps} batches ...")
-            for it, (_, gen_candidates_eval_info) in zip(
+            for it, batch in zip(
                 range(num_training_steps, num_training_steps + num_final_gen_steps + 1),
                 cycle(final_dl),
             ):
+                _, gen_candidates_eval_info = self._maybe_resolve_shared_buffer(batch, final_dl)
                 gen_candidates_list.append(gen_candidates_eval_info)
 
             info_final_gen = compute_metrics(gen_candidates_list, cand_type=self.task.cand_type, k=self.cfg.evaluation.k, reward_thresh=self.cfg.evaluation.reward_thresh, distance_thresh=self.cfg.evaluation.distance_thresh)