Repository: pytorch/PiPPy

Issue:
Pipeline Driver slowing down pipeline execution
### Problem Statement:
As Figure 1 shows, the backward stages were not executed as soon as the previous chunk was completed. Instead, they waited for the driver to issue the `invoke` command.

<img width="979" alt="Screen Shot 2022-08-08 at 2 09 52 PM" src="https://user-images.githubusercontent.com/6676466/183515427-8fcc2b9d-c40e-4afc-b798-d9df29525833.png">

### Analysis:
Zooming in, there are two factors that cause the delay of invoke calls.
1. Current Pipeline Interpreter uses sync RPC calls. Thus, it blocks the driver till the remote executor finishes the invoke call and returns. See "IB_7,0" for example.
2. Each `getitem` invocation is a separate RPC call. If a stage has many output indices, then the driver would have to issue many RPC calls. See "G7,0" for example.
<img width="1767" alt="Screen Shot 2022-07-25 at 2 01 05 PM" src="https://user-images.githubusercontent.com/6676466/183516757-596135e7-fb53-4e95-af6e-f79c530fe43a.png">

Figures are from T5 3B model, with 8 chunks (8 PP ranks), and small chunk size.


Discussion:

Resolving patch:
diff --git a/examples/hf/t5/pippy_t5.py b/examples/hf/t5/pippy_t5.py
--- a/examples/hf/t5/pippy_t5.py
+++ b/examples/hf/t5/pippy_t5.py
@@ -210,7 +210,7 @@ def run_master(pp_ranks, args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--world_size', type=int, default=int(os.getenv("WORLD_SIZE", 16)))
+    parser.add_argument('--world_size', type=int, default=int(os.getenv("WORLD_SIZE", 8)))
     parser.add_argument('--rank', type=int, default=int(os.getenv("RANK", -1)))
     parser.add_argument('--master_addr', type=str, default=os.getenv('MASTER_ADDR', 'localhost'))
     parser.add_argument('--master_port', type=str, default=os.getenv('MASTER_PORT', '29500'))
diff --git a/pippy/PipelineDriver.py b/pippy/PipelineDriver.py
--- a/pippy/PipelineDriver.py
+++ b/pippy/PipelineDriver.py
@@ -419,7 +419,7 @@ def update_run_list(self, runlist_key, arg_idx, value):
                     work_item.state = SchedState.READY
                     self.ready_runlist[runlist_key] = self.waiting_runlist.pop(runlist_key)
                     self.ready_runlist_cv.notify()
-                logging.info(f'[{self.rank}][{work_item.microbatch_id}] all operands ready')
+                logging.info(f'[{self.rank}][{work_item.microbatch_id}] all operands ready: {runlist_key}')
 
 
 class PipeStageExecutor(EventRecorder):
@@ -441,6 +441,7 @@ def __init__(self, stage_id, mod, rank_worker, _record_mem_dumps=False):
         self.fwd_cache : Dict[int, Tuple[Any, List[torch.Tensor]]] = {}
 
         self.value_store_lock = threading.Lock()
+        self.value_store_cv = threading.Condition(self.value_store_lock)
         self.value_store : Dict[str, RefcountedFuture] = {}
 
         self.peer_executors : Dict[int, torch._C._distributed_rpc.PyRRef] = None
@@ -516,15 +517,13 @@ def init_data_parallel(self, n_stages, dp_group_size, dp_pg_cb=None):
 
     def invoke(self, output_unique_key : str, phase : Phase, args, kwargs, cur_microbatch : int, debug_str : str,
                output_refcount : int, batch_id : int, num_microbatches : int):
-        ts = time.time()
-        forward_name = event_name(Phase.FORWARD, self.stage_id, cur_microbatch)
-        forward_id = event_id(Phase.FORWARD, self.stage_id, cur_microbatch, batch_id)
-        name = f"R{forward_name}"
-        id = f"R{forward_id}"
-        self.record_event(rank=self.rank_worker.rank, start_ts=ts, finish_ts=ts, id=id, name=name, type='received', mbid=cur_microbatch)
-        self.record_event_dependency(from_id=name, to_id=forward_name, type='waiting')
+        start_ts = time.time()
+        target_name = event_name(phase, self.stage_id, cur_microbatch)
+        target_id = event_id(phase, self.stage_id, cur_microbatch, batch_id)
+        name = f"R{target_name}"
+        id = f"R{target_id}"
         if self._record_mem_dumps:
-            self._record_dumps_on_all_peer_executors(f'M{id}_invoke', ts)
+            self._record_dumps_on_all_peer_executors(f'M{id}_invoke', start_ts)
         # TODO: do we need to serialize calls to invoke() to preserve the order in which WorkItems appear for
         # static schedules?
 
@@ -552,8 +551,8 @@ def extract_value_ref_args(arg):
         # HACK: we assume the module has at least one parameter
         param = next(self.mod.parameters(), None)
         if param is None:
-            warnings.warn(f"Module of stage {self.stage_id} has 0 parameters, "
-                          f"cannot figure out device. Setting it to cpu")
+            logging.warning(f"Module of stage {self.stage_id} has 0 parameters, "
+                            f"cannot figure out device. Setting it to cpu")
         else:
             device = param.device
 
@@ -594,32 +593,62 @@ def extract_value_ref_args(arg):
             for fut in _futures:
                 fut.wait()
 
-        with self.value_store_lock:
-            assert output_unique_key not in self.value_store, f"output_unique_key = {output_unique_key}"
+        with self.value_store_cv:
+            assert output_unique_key not in self.value_store, (f"[{self.stage_id}] Output key {output_unique_key} "
+                                                               f"already exists or is not consumed from previous batch")
             self.value_store[output_unique_key] = RefcountedFuture(future, output_refcount)
+            self.value_store_cv.notify_all()
+
+        finish_ts = time.time()
+        self.record_event(rank=self.rank_worker.rank, start_ts=start_ts, finish_ts=finish_ts, id=id, name=name, type='received', mbid=cur_microbatch)
+        self.record_event_dependency(from_id=name, to_id=target_name, type='waiting')
 
         return ValueReference(self.stage_id, output_unique_key)
 
-    def index_value(self, output_unique_key : str, output_refcount : int, value_ref, idx):
-        # For the purposes of refcounting, decrement this use
-        with self.value_store_lock:
-            refcounted_future = self.value_store[value_ref.unique_key]
-            if refcounted_future.release():
-                self.value_store.pop(value_ref.unique_key)
+    def coalesced_index_value(self, indices : List[Tuple[str, int, ValueReference, int]]):
+        for index_tuple in indices:
+            (output_unique_key, output_refcount, value_ref, idx) = index_tuple
+            logging.info(f'[{self.stage_id}] Received getitem call: {(output_unique_key, output_refcount, value_ref, idx)}')
+            with self.value_store_cv:
+                if output_unique_key in self.value_store:
+                    logging.warning(f'[{self.stage_id}] Indexed value already in store: {(output_unique_key, output_refcount, value_ref, idx)}')
+                    # raise RuntimeError(f'Repeated index value call detected, potentially due to getitem calls not consumed in previous batch')
 
-            indexed = refcounted_future.future.then(lambda f: f.value()[idx])
+                # Wait for the future representing the stage output to be created
+                while value_ref.unique_key not in self.value_store:
+                    self.value_store_cv.wait()
+                # Now the stage output future is created
+                refcounted_future = self.value_store[value_ref.unique_key]
 
-            self.value_store[output_unique_key] = RefcountedFuture(indexed, output_refcount)
+                # For the purposes of refcounting, decrement this use
+                if refcounted_future.release():
+                    self.value_store.pop(value_ref.unique_key)
+
+                # Create an indexed future that represents a specific output arg
+                # Here we use an attach functon so that the index passed to the lambda changes in every loop
+                def attach_index(fut, index):
+                    indexed_fut = fut.then(lambda f: f.value()[index])
+                    return indexed_fut
+
+                indexed = attach_index(refcounted_future.future, idx)
+
+                # Enqueue the indexed future
+                # And notify places that may be waiting for it to be created, such as get_value
+                self.value_store[output_unique_key] = RefcountedFuture(indexed, output_refcount)
+                self.value_store_cv.notify_all()
 
-        return ValueReference(self.stage_id, output_unique_key)
 
     def get_value(self, caller_stage, runlist_key, microbatch, value_ref_arg):
         callee_stage = value_ref_arg.stage_id
-        logging.info(f'[{callee_stage}][{microbatch}] Executing async transfer of value '
+        logging.info(f'[{callee_stage}][{microbatch}] Executing transfer of value '
                      f'{value_ref_arg} initiated by stage {caller_stage} for {runlist_key}')
         assert callee_stage == self.stage_id, "Mismatch between ValueRef and stage executor"
 
-        with self.value_store_lock:
+        with self.value_store_cv:
+            # Waiting for the indexed future for this arg to be created
+            while value_ref_arg.unique_key not in self.value_store:
+                self.value_store_cv.wait()
+            # Now the indexed future is created
             refcounted_future = self.value_store[value_ref_arg.unique_key]
 
         value = refcounted_future.future.wait()
@@ -631,15 +660,16 @@ def get_value(self, caller_stage, runlist_key, microbatch, value_ref_arg):
         return value
 
     def async_transfer(self, microbatch, value_ref_arg, arg_idx, runlist_key):
-        logging.info(f'[{self.stage_id}][{microbatch}] Starting transfer')
+        logging.info(f'[{self.stage_id}][{microbatch}] Requesting transfer of value {value_ref_arg} '
+                     f'for runlist item {runlist_key} arg_idx {arg_idx}')
         value_ref_executor_rref = self.peer_executors[value_ref_arg.stage_id]
         fut = value_ref_executor_rref.rpc_async().get_value(
             self.stage_id, runlist_key, microbatch, value_ref_arg)
 
         def bottom_half(fut):
-            logging.info(f'[{self.stage_id}][{microbatch}] Completing transfer of value {value_ref_arg} '
-                         f'for runlist item {runlist_key}')
             value = fut.value()
+            logging.info(f'[{self.stage_id}][{microbatch}] Completing transfer of value {value_ref_arg} '
+                         f'for runlist item {runlist_key} arg_idx {arg_idx}')
             self.rank_worker.update_run_list(runlist_key, arg_idx, value)
 
         return fut.then(bottom_half)
@@ -679,6 +709,12 @@ def instantiate_lr_scheduler(self, lr_sched_class, *args, **kwargs):
         logging.info(f"[{self.stage_id}] Creating learning rate scheduler")
         return lr_sched_class(self.optimizer, *args, **kwargs)
 
+    def _check_cleanup(self) -> bool:
+        if len(self.value_store):
+            logging.warning(f"[{self.stage_id}] Unclean value store: {self.value_store}")
+            return False
+        return True
+
     def _record_dump(self, dump_id, ts):
         first_param = next(self.mod.parameters(), None)
         device: torch.device = first_param.device if first_param is not None else torch.device('cpu')
@@ -1068,6 +1104,12 @@ def retrieve_events(self) -> EventsContext:
         events_context.events.sort(key=lambda e: e.start_ts)
         return events_context
 
+    def _check_stages_cleanup(self) -> bool:
+        clean = True
+        for executor in self.stage_to_executor.values():
+            clean &= executor.rpc_sync()._check_cleanup()
+        return clean
+
 
 class RemoteInterpreter(pippy.fx.Interpreter, EventRecorder):
     def __init__(self, remote_stage_executor_rrefs, stage_to_executor, module, cur_microbatch : int,
@@ -1078,6 +1120,7 @@ def __init__(self, remote_stage_executor_rrefs, stage_to_executor, module, cur_m
         self.cur_microbatch = cur_microbatch
         self.pc = 0
         self.node_list = list(self.module.graph.nodes)
+        logging.info(f'[root] RemoteInterpreter created with {len(self.node_list)} nodes')
 
         # Process args/kwargs
 
@@ -1095,6 +1138,8 @@ def __init__(self, remote_stage_executor_rrefs, stage_to_executor, module, cur_m
         self.args_iter = iter(self.args)
         self.batch_id = batch_id
         self.num_microbatches = num_microbatches
+        # Dict from stage id to a list holding the coalesced getitem indices
+        self.stage_output_indices : Dict[int, List[Tuple[str, int, ValueReference, int]]] = {}
 
     def call_module(self, target, args, kwargs):
         assert isinstance(target, str)
@@ -1110,17 +1155,19 @@ def call_module(self, target, args, kwargs):
             logging.info(f'[root][{self.cur_microbatch}] Issuing {Phase.FORWARD} '
                          f'invocation for target {target} on stage {stage_id}')
             invocation_key = f'{self.cur_microbatch}_{node.name}'
-            ts = time.time()
+            start_ts = time.time()
             forward_name = event_name(Phase.FORWARD, stage_id, self.cur_microbatch)
             forward_id = event_id(Phase.FORWARD, stage_id, self.cur_microbatch, self.batch_id)
             name = f"I{forward_name}"
             id = f"I{forward_id}"
-            self.record_event(rank=0, start_ts=ts, finish_ts=ts, id=id, name=name, type='invoke',
-                              mbid=self.cur_microbatch)
-            self.record_event_dependency(from_id=name, to_id=f"R{forward_name}", type='invoke')
-            return stage_executor.rpc_sync().invoke(
+            stage_executor.rpc_async().invoke(
                 invocation_key, Phase.FORWARD, args, kwargs, self.cur_microbatch, debug_str=node.format_node(),
                 output_refcount=len(users), batch_id=self.batch_id, num_microbatches=self.num_microbatches)
+            finish_ts = time.time()
+            self.record_event(rank=0, start_ts=start_ts, finish_ts=finish_ts, id=id, name=name, type='invoke',
+                              mbid=self.cur_microbatch)
+            self.record_event_dependency(from_id=name, to_id=f"R{forward_name}", type='invoke')
+            return ValueReference(stage_id, invocation_key)
         else:
             logging.info(f'[root][{self.cur_microbatch}] Running local operation {target} from driver')
             return super().call_module(target, args, kwargs)
@@ -1136,30 +1183,37 @@ def call_function(self, target, args, kwargs):
             node.users.keys())) if not torch.is_grad_enabled() else node.users.keys()
         if target is operator.getitem and isinstance(args[0], ValueReference):
             stage_id = args[0].stage_id
-            if torch.is_grad_enabled() or args[0].unique_key != "noop":
-                stage_executor = self.stage_to_executor[stage_id]
-                return stage_executor.rpc_sync().index_value(
-                    output_unique_key=invocation_key, value_ref=args[0], output_refcount=len(users),
-                    idx=args[1])
-            else:
+            num_users = len(users)
+            if (not torch.is_grad_enabled() or
+                    args[0].unique_key == "noop" or
+                    num_users == 0):
+                # TODO: investigate why there are getitem calls with 0 users
                 return ValueReference(stage_id, "noop")
+            else:
+                indices = self.stage_output_indices.setdefault(stage_id, [])
+                index_tuple = (invocation_key, num_users, args[0], args[1])
+                logging.info(f'[root][{self.cur_microbatch}] Appending getitem tuple to stage {stage_id}: {index_tuple}')
+                indices.append(index_tuple)
+                return ValueReference(stage_id, invocation_key)
         elif target is stage_backward:
             assert 'fw_stage' in node.meta
             stage_id, stage_executor = self.remote_stage_executor_rrefs[node.meta['fw_stage']]
             if torch.is_grad_enabled():
                 logging.info(f'[root][{self.cur_microbatch}] Issuing BW invocation '
                              f'for target {node.meta["fw_stage"]} on stage {stage_id}')
-                ts = time.time()
+                start_ts = time.time()
                 backward_name = event_name(Phase.BACKWARD, stage_id, self.cur_microbatch)
                 backward_id = event_id(Phase.BACKWARD, stage_id, self.cur_microbatch, self.batch_id)
                 name = f"I{backward_name}"
                 id = f"I{backward_id}"
-                self.record_event(rank=0, start_ts=ts, finish_ts=ts, id=id, name=name, type='invoke',
-                                  mbid=self.cur_microbatch)
-                self.record_event_dependency(from_id=name, to_id=backward_name, type='invoke')
-                return stage_executor.rpc_sync().invoke(
+                stage_executor.rpc_async().invoke(
                     invocation_key, Phase.BACKWARD, args, kwargs, self.cur_microbatch, debug_str=node.format_node(),
                     output_refcount=len(users), batch_id=self.batch_id, num_microbatches=self.num_microbatches)
+                finish_ts = time.time()
+                self.record_event(rank=0, start_ts=start_ts, finish_ts=finish_ts, id=id, name=name, type='invoke',
+                                  mbid=self.cur_microbatch)
+                self.record_event_dependency(from_id=name, to_id=backward_name, type='invoke')
+                return ValueReference(stage_id, invocation_key)
             else:
                 return ValueReference(stage_id, "noop")
         elif target is sync_barrier:
@@ -1167,29 +1221,52 @@ def call_function(self, target, args, kwargs):
             stage_id, stage_executor = self.remote_stage_executor_rrefs[executor_keys[0]]
             logging.info(f'[root][{self.cur_microbatch}] Issuing sync invocation '
                          f'on stage {stage_id}')
-            return stage_executor.rpc_sync().invoke(
+            stage_executor.rpc_async().invoke(
                 invocation_key, Phase.SYNC_BARRIER, args, kwargs, self.cur_microbatch, debug_str=node.format_node(),
                 output_refcount=len(users), batch_id=self.batch_id, num_microbatches=self.num_microbatches)
+            return ValueReference(stage_id, invocation_key)
         elif target is _null_coalesce_accumulate:
             assert 'fw_stage' in node.meta
             stage_id, stage_executor = self.remote_stage_executor_rrefs[node.meta['fw_stage']]
             if torch.is_grad_enabled():
                 logging.info(f'[root][{self.cur_microbatch}] Issuing accumulate grad invocation '
                              f'for target {node.meta["fw_stage"]} on stage {stage_id}')
-                return stage_executor.rpc_sync().invoke(
+                stage_executor.rpc_async().invoke(
                     invocation_key, Phase.ACCUMULATE_GRAD, args, kwargs, self.cur_microbatch,
                     debug_str=node.format_node(),
                     output_refcount=len(users), batch_id=self.batch_id, num_microbatches=self.num_microbatches)
+                return ValueReference(stage_id, invocation_key)
             else:
                 return ValueReference(stage_id, "noop")
         else:
             raise AssertionError(f'Unknown operator {torch.typename(target)}')
 
+
+    def issue_coalesced_getitem_calls(self):
+        if len(self.stage_output_indices) == 0:
+            return
+        logging.info(f'[root][{self.cur_microbatch}] Issuing getitem calls to stage: {self.stage_output_indices.keys()}')
+
+        for stage_id in self.stage_output_indices:
+            stage_executor = self.stage_to_executor[stage_id]
+            name = f"G{stage_id},{self.cur_microbatch}"
+            id = f"{name},{self.batch_id}"
+            start_ts = time.time()
+            stage_executor.rpc_async().coalesced_index_value(self.stage_output_indices[stage_id])
+            finish_ts = time.time()
+            self.record_event(rank=0, start_ts=start_ts, finish_ts=finish_ts, id=id, name=name, type='invoke',
+                              mbid=self.cur_microbatch)
+
+        self.stage_output_indices.clear()
+
+
     def run_until(self, predicate: Callable[[pippy.fx.Node], bool]) -> Optional[pippy.fx.Node]:
         while self.pc < len(self.node_list):
             node = self.node_list[self.pc]
 
             if predicate(node):
+                # Issue coalesced getitem calls as we pause issuing stage calls
+                self.issue_coalesced_getitem_calls()
                 return node
 
             self.run_one(node)
@@ -1207,6 +1284,8 @@ def run_one(self, node):
         #  to the callees) as an optimization
         # TODO: is it possible for there to be a blocking version of this API?
         def wait_for_confirmation(n):
+            # The following if will not be true as we are using our own ValueRef
+            # instead of RPC's RRef
             if isinstance(n, torch._C._distributed_rpc.PyRRef):
                 while not n.confirmed_by_owner():
                     pass
@@ -1221,6 +1300,28 @@ def wait_for_confirmation(n):
         return node
 
 
+class _run_until_criteria:
+    def __init__(self):
+        self.seen_stages = 0
+
+    # Run the node we start with including all nodes that are tuple
+    # indexing, then stop
+    def hitting_next_stage(self, node):
+        if node.op == 'output':
+            return True
+
+        if node.target != operator.getitem and node.target != _null_coalesce_accumulate:
+            self.seen_stages += 1
+
+        if self.seen_stages > 1:
+            return True
+        elif self.seen_stages == 1 and node.target == _null_coalesce_accumulate:
+            # We are hitting the accumulate call of the next (backward) stage, stop
+            return True
+        else:
+            return False
+
+
 class PipelineDriverFillDrain(PipelineDriverBase):
     def __init__(self, pipe: Pipe, chunks: int, args_chunk_spec, kwargs_chunk_spec, output_chunk_spec, world_size: int,
                  all_ranks: List[int] = None, single_loss: bool = False, _debug_mask_minibatches: bool = False,
@@ -1276,17 +1377,8 @@ def forward(self, *args, **kwargs):
         for ramp_up_idx in range(len(self.microbatch_interpreters)):
             for i in range(ramp_up_idx + 1):
                 interp = self.microbatch_interpreters[i]
-                start_node = interp.node_list[min(interp.pc, len(interp.node_list) - 1)]
-
-                def run_including_indexing(n):
-                    if n.op == 'output':
-                        return True
-
-                    # Run the node we start with including all nodes that are tuple
-                    # indexing, then stop
-                    return n != start_node and n.target != operator.getitem
-
-                interp.run_until(run_including_indexing)
+                criteria = _run_until_criteria()
+                interp.run_until(criteria.hitting_next_stage)
 
         # Steady-state. We have a full diagonal in the matrix; keep dispatching
         # across the diagonal
@@ -1296,19 +1388,8 @@ def run_including_indexing(n):
             any_valid = False
             for interp in self.microbatch_interpreters:
                 start_node = interp.node_list[min(interp.pc, len(interp.node_list) - 1)]
-
-                def run_including_indexing(n):
-                    # Run the node we start with including all nodes that are
-                    # tuple indexing, but also stop at the output node. Because
-                    # we are on a diagonal, interpreters as lower microbatch IDs
-                    # will be invoked when their pc's point to the output node
-                    # multiple times.
-                    if n.op == 'output':
-                        return True
-
-                    return n != start_node and n.target != operator.getitem
-
-                interp.run_until(run_including_indexing)
+                criteria = _run_until_criteria()
+                interp.run_until(criteria.hitting_next_stage)
 
                 any_valid |= interp.node_list[interp.pc] != start_node
 
@@ -1324,6 +1405,9 @@ def run_including_indexing(n):
             # (by way of the synchronization dependency earlier)
             self._sync_replicated_params()
 
+        if DEBUG:
+            self._check_stages_cleanup()
+
         return merge_chunks(local_results, self.output_chunk_spec, self._debug_mask_minibatches)
 
 
diff --git a/pippy/utils.py b/pippy/utils.py
--- a/pippy/utils.py
+++ b/pippy/utils.py
@@ -60,7 +60,7 @@ def run_worker(rank, run_master, args, *extra_args):
 
     # TODO: Move to training args, blocked by: cannot pickle 'TensorPipeRpcBackendOptions' object
     # Exclude IB for metadata transport due to lack of EFA support on AWS
-    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256,
+    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=512,
                                               rpc_timeout=1800,
                                               _transports=tp_transports())
     if args.cuda:


Test patch:
diff --git a/.github/workflows/pippy_tests.yaml b/.github/workflows/pippy_tests.yaml
--- a/.github/workflows/pippy_tests.yaml
+++ b/.github/workflows/pippy_tests.yaml
@@ -158,18 +158,18 @@ jobs:
         run: python examples/hf/bert/pippy_bert.py --replicate ${{ matrix.replicate }} -s ${{ matrix.schedule }}
       - name: Run T5 example
         run: python examples/hf/t5/pippy_t5.py --replicate ${{ matrix.replicate }} -s ${{ matrix.schedule }} --checkpoint ${{ matrix.checkpoint }}
-      - name: "HF Causal Language Modeling: fine-tune GPT-2 on WikiText-2"
-        run: python examples/hf/language-modeling/run_clm.py --dp_group_size=2 --pp_group_size=8 --model_name_or_path gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --do_train --do_eval --output_dir /tmp/test-clm --max_steps=3 --overwrite_output_dir
-      - name: "HF Masked Language Modeling: fine-tune RoBERTa on WikiText-2"
-        run: python examples/hf/language-modeling/run_mlm.py --dp_group_size=2 --pp_group_size=8 --model_name_or_path roberta-base --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --do_train --do_eval --output_dir /tmp/test-mlm --max_steps=3 --overwrite_output_dir
-      - name: "HF Text classification: fine-tune BERT on the GLUE benchmark"
-        run: python examples/hf/text-classification/run_glue.py --dp_group_size=2 --pp_group_size=8 --model_name_or_path bert-base-cased --task_name mrpc --do_train --do_eval --max_seq_length 128 --per_device_train_batch_size 32 --learning_rate 2e-5 --num_train_epochs 3 --output_dir /tmp/mrpc/ --max_steps=3 --overwrite_output_dir
-      - name: "HF Translation: fine-tune T5 model translation English to Romanian"
-        run: >
-          python examples/hf/translation/run_translation.py --model_name_or_path t5-small --do_train --source_lang en --target_lang ro --source_prefix "translate English to Romanian: " --dataset_name wmt16 --dataset_config_name ro-en --output_dir /tmp/tst-translation --per_device_train_batch_size=8 --per_device_eval_batch_size=8 --overwrite_output_dir --predict_with_generate --max_steps=10  --dp_group_size=2 --pp_group_size=8
-      - name: "HF Translation: fine-tune BART model translation English to Romanian"
-        run: >
-          python examples/hf/translation/run_translation.py --model_name_or_path facebook/bart-base --do_train --source_lang en --target_lang ro --source_prefix "translate English to Romanian: " --dataset_name wmt16 --dataset_config_name ro-en --output_dir /tmp/tst-translation --per_device_train_batch_size=8 --per_device_eval_batch_size=8 --overwrite_output_dir --predict_with_generate --max_steps=10  --dp_group_size=2 --pp_group_size=8
+      #- name: "HF Causal Language Modeling: fine-tune GPT-2 on WikiText-2"
+      #run: python examples/hf/language-modeling/run_clm.py --dp_group_size=2 --pp_group_size=8 --model_name_or_path gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --do_train --do_eval --output_dir /tmp/test-clm --max_steps=3 --overwrite_output_dir
+      #- name: "HF Masked Language Modeling: fine-tune RoBERTa on WikiText-2"
+      #run: python examples/hf/language-modeling/run_mlm.py --dp_group_size=2 --pp_group_size=8 --model_name_or_path roberta-base --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --do_train --do_eval --output_dir /tmp/test-mlm --max_steps=3 --overwrite_output_dir
+      #- name: "HF Text classification: fine-tune BERT on the GLUE benchmark"
+      #run: python examples/hf/text-classification/run_glue.py --dp_group_size=2 --pp_group_size=8 --model_name_or_path bert-base-cased --task_name mrpc --do_train --do_eval --max_seq_length 128 --per_device_train_batch_size 32 --learning_rate 2e-5 --num_train_epochs 3 --output_dir /tmp/mrpc/ --max_steps=3 --overwrite_output_dir
+      #- name: "HF Translation: fine-tune T5 model translation English to Romanian"
+      #run: >
+      #  python examples/hf/translation/run_translation.py --model_name_or_path t5-small --do_train --source_lang en --target_lang ro --source_prefix "translate English to Romanian: " --dataset_name wmt16 --dataset_config_name ro-en --output_dir /tmp/tst-translation --per_device_train_batch_size=8 --per_device_eval_batch_size=8 --overwrite_output_dir --predict_with_generate --max_steps=10  --dp_group_size=2 --pp_group_size=8
+      #- name: "HF Translation: fine-tune BART model translation English to Romanian"
+      #run: >
+      #  python examples/hf/translation/run_translation.py --model_name_or_path facebook/bart-base --do_train --source_lang en --target_lang ro --source_prefix "translate English to Romanian: " --dataset_name wmt16 --dataset_config_name ro-en --output_dir /tmp/tst-translation --per_device_train_batch_size=8 --per_device_eval_batch_size=8 --overwrite_output_dir --predict_with_generate --max_steps=10  --dp_group_size=2 --pp_group_size=8
       - name: Run visualizer test
         run: python test/local_test_visualizer.py --replicate ${{ matrix.replicate }} -s ${{ matrix.schedule }}
 
diff --git a/test/pippy_wrapper.sh b/test/pippy_wrapper.sh
--- a/test/pippy_wrapper.sh
+++ b/test/pippy_wrapper.sh
@@ -9,4 +9,4 @@ export LOCAL_RANK=${SLURM_LOCALID}
 export WORLD_SIZE=${SLURM_NTASKS}
 export RANK=${SLURM_PROCID}
 
-exec "$@"
+python -u "$@" 2>&1