Repository: pytorch/PiPPy

Issue:
[Bug] RemoteInterpreter.call_function doesn't handle _null_coalesce_accumulate

```
# Copyright (c) Meta Platforms, Inc. and affiliates
import argparse
import os
import socket

import torch
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp

from pippy.IR import Pipe, pipe_split, TrivialLossWrapper, _null_coalesce_accumulate
from pippy.PipelineDriver import PipelineDriverFillDrain, PipelineDriver1F1B, PipelineDriverBase
from pippy.microbatch import TensorChunkSpec, CustomReducer

PROFILING_ENABLED = True
CHECK_NUMERIC_EQUIVALENCE = True

schedules = {
    'FillDrain': PipelineDriverFillDrain,
    '1F1B': PipelineDriver1F1B,
}

torch.fx.Tracer.proxy_buffer_attributes = True


def run_master(args):
    all_ranks = list(range(1, args.world_size))  # exclude master rank = 0
    chunks = len(all_ranks)
    batches = 1
    bs = 4 * chunks
    hid_dim = 50

    class Code(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = torch.nn.Linear(hid_dim, hid_dim)

        def forward(self, x):
            x = self.linear(x)
            pipe_split()
            y = torch.relu(x)
            pipe_split()
            z = torch.sigmoid(x)
            pipe_split()
            return y + z

    c = Code()
    c.train()
    mse_loss = torch.nn.MSELoss()
    wrapper = TrivialLossWrapper(c, mse_loss)
    accum_pipe = Pipe.from_tracing(wrapper)
    assert 4 == len(list(accum_pipe.split_gm.children()))
    assert any(n.target == _null_coalesce_accumulate for n in accum_pipe.split_gm.graph.nodes)
    input = torch.randn(bs, hid_dim)
    target = torch.randn(bs, hid_dim)
    accum_pipe(input, target)

    args_chunk_spec = (TensorChunkSpec(0), TensorChunkSpec(0))
    kwargs_chunk_spec = {}
    output_chunk_spec = CustomReducer(torch.tensor(0.0), lambda a, b: a + b)
    pipe_driver: PipelineDriverBase = schedules[args.schedule](accum_pipe, args_chunk_spec, kwargs_chunk_spec,
                                                               output_chunk_spec, args.world_size - 1,
                                                               all_ranks=all_ranks, _debug_mask_minibatches=True)

    for i in range(batches):
        pipe_driver.run(chunks, input, target)


def run_worker(rank, world_size, args):
    print(f"rank = {rank} host/pid = {socket.gethostname()}/{os.getpid()}")
    os.environ['MASTER_ADDR'] = args.master_addr
    os.environ['MASTER_PORT'] = args.master_port
    if args.rank == -1:  # run via mp.spawn
        # each worker will see its GPU as `cuda:0`
        try:
            import subprocess
            device_count = int(subprocess.getoutput('nvidia-smi --list-gpus | wc -l'))
            os.environ['CUDA_VISIBLE_DEVICES'] = str(rank % device_count)
        except ValueError:
            pass
    assert not torch.cuda.is_available() or torch.cuda.device_count() == 1, \
        "Do not use torch.cuda.* before setting CUDA_VISIBLE_DEVICES"
    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256, _transports=["uv"])  # uv for AWS EFA instances
    if args.use_cuda and torch.cuda.is_available():
        for i in range(world_size):
            options.set_device_map(f"worker{i}", {0: 0})
    rpc.init_rpc(
        f"worker{rank}",
        rank=rank,
        world_size=world_size,
        rpc_backend_options=options
    )
    if rank == 0:
        run_master(args)
    rpc.shutdown()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--world_size', type=int, default=int(os.getenv("WORLD_SIZE", 5)))
    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'))
    parser.add_argument('-s', '--schedule', type=str, default=list(schedules.keys())[0], choices=schedules.keys())
    parser.add_argument('--replicate', type=int, default=int(os.getenv("REPLICATE", '0')))
    parser.add_argument('--use_cuda', type=int, default=1)
    args = parser.parse_args()

    if args.rank == -1:
        mp.spawn(run_worker, args=(args.world_size, args,), nprocs=args.world_size, join=True)
    elif args.rank < args.world_size:
        run_worker(args.rank, args.world_size, args)
    else:
        print("I'm unused, exiting")

```
```
Traceback (most recent call last):
  File "/Users/pbelevich/miniconda3/envs/PiPPy/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/Users/pbelevich/PycharmProjects/PiPPy/examples/slurm/hf/t5/test_null_coalesce_accumulate.py", line 93, in run_worker
    run_master(args)
  File "/Users/pbelevich/PycharmProjects/PiPPy/examples/slurm/hf/t5/test_null_coalesce_accumulate.py", line 65, in run_master
    pipe_driver.run(chunks, input, target)
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/PipelineDriver.py", line 871, in run
    last_nodes.append(interp.run_until(lambda n: n.op == 'output'))
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/PipelineDriver.py", line 811, in run_until
    self.env[node] = super().run_node(node)
  File "/Users/pbelevich/miniconda3/envs/PiPPy/lib/python3.9/site-packages/torch/fx/interpreter.py", line 152, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/Users/pbelevich/PycharmProjects/PiPPy/pippy/PipelineDriver.py", line 800, in call_function
    raise AssertionError(f'Unknown operator {torch.typename(target)}')
AssertionError: Unknown operator pippy.IR._null_coalesce_accumulate
```


Discussion:

Resolving patch:
diff --git a/pippy/PipelineDriver.py b/pippy/PipelineDriver.py
--- a/pippy/PipelineDriver.py
+++ b/pippy/PipelineDriver.py
@@ -12,7 +12,7 @@
 import torch.distributed.rpc as rpc
 import torch.fx
 
-from pippy.IR import Pipe, stage_backward, sync_barrier
+from pippy.IR import Pipe, stage_backward, sync_barrier, _null_coalesce_accumulate
 from pippy.events import EventRecorder, EventsContext, Event
 from pippy.microbatch import split_args_kwargs_into_chunks, merge_chunks
 
@@ -50,7 +50,8 @@
 class Phase(Enum):
     FORWARD = 0
     BACKWARD = 1
-    SYNC_BARRIER = 2
+    ACCUMULATE_GRAD = 2
+    SYNC_BARRIER = 3
 
 # TODO: do we need this?
 class SchedState(Enum):
@@ -64,6 +65,7 @@ def event_name(ph, stage_id, mbid):
     phase_to_short_str = {
         Phase.FORWARD: 'F',
         Phase.BACKWARD: 'B',
+        Phase.ACCUMULATE_GRAD: 'A',
         Phase.SYNC_BARRIER: 'S',
     }
     return f"{phase_to_short_str[ph]}_{stage_id},{mbid}"
@@ -130,6 +132,9 @@ def __init__(self, stage_id, unique_key):
     stage_id : int
     unique_key : str
 
+    def __repr__(self):
+        return f'ValueReference({self.stage_id}, {self.unique_key})'
+
 
 class RefcountedFuture:
     future : torch.futures.Future
@@ -357,6 +362,9 @@ def extract_tensor_args(a):
 
                 # Schedule forward stage of a new micro-batch
                 self.outstanding -= 1
+            elif work_item.phase == Phase.ACCUMULATE_GRAD:
+                logging.info(f'[{self.local_rank}][{work_item.microbatch_id}] Running accumulate grad')
+                out_val = _null_coalesce_accumulate(*args, **kwargs)
             elif work_item.phase == Phase.SYNC_BARRIER:
                 logging.info(f'[{self.local_rank}][{work_item.microbatch_id}] Running sync_barrier')
                 out_val = sync_barrier(*args, **kwargs)
@@ -499,7 +507,6 @@ def get_value(self, caller_stage, runlist_key, microbatch, value_ref_arg):
 
         return value
 
-
     def get_grad(self, qualname):
         mod = self.mod
         if isinstance(mod, torch.nn.parallel.DistributedDataParallel):
@@ -564,6 +571,8 @@ class ExecutorDescriptor:
                 executor_descriptors[bw_idx].has_backward = True
                 node.meta['fw_stage'] = executor_descriptors[bw_idx].name
                 bw_idx -= 1
+            elif (node.op, node.target) == ('call_function', _null_coalesce_accumulate):
+                node.meta['fw_stage'] = executor_descriptors[bw_idx].name
 
         assert all(d.has_backward for d in executor_descriptors) or all(not d.has_backward for d in executor_descriptors)
 
@@ -744,6 +753,15 @@ def call_function(self, target, args, kwargs):
             return stage_executor.rpc_sync().invoke(
                 invocation_key, Phase.SYNC_BARRIER, args, kwargs, self.cur_microbatch, debug_str=node.format_node(),
                 output_refcount=len(node.users), batch_id=self.batch_id, num_microbatches=self.num_microbatches)
+        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']]
+            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(
+                invocation_key, Phase.ACCUMULATE_GRAD, args, kwargs, self.cur_microbatch,
+                debug_str=node.format_node(),
+                output_refcount=len(node.users), batch_id=self.batch_id, num_microbatches=self.num_microbatches)
         else:
             raise AssertionError(f'Unknown operator {torch.typename(target)}')
 


Test patch:
diff --git a/.github/workflows/ci_tests.yaml b/.github/workflows/ci_tests.yaml
--- a/.github/workflows/ci_tests.yaml
+++ b/.github/workflows/ci_tests.yaml
@@ -125,6 +125,8 @@ jobs:
         run: VERBOSE=${{ matrix.verbose }} python test/local_test_forward.py --replicate ${{ matrix.replicate }} -s ${{ matrix.schedule }}
       - name: Run forward-loss-backward integration test
         run: VERBOSE=${{ matrix.verbose }} python test/local_test_forward_backward.py --replicate ${{ matrix.replicate }} -s ${{ matrix.schedule }}
+      - name: Run null_coalesce_accumulate integration test
+        run: VERBOSE=${{ matrix.verbose }} python test/local_test_null_coalesce_accumulate.py --replicate ${{ matrix.replicate }} -s ${{ matrix.schedule }}
       - name: Run HF BERT forward-only integration test
         run: VERBOSE=${{ matrix.verbose }} python test/local_test_forward_hf_bert.py --replicate ${{ matrix.replicate }} -s ${{ matrix.schedule }}
       - name: Run HF GPT2 forward-only integration test
diff --git a/test/local_test_null_coalesce_accumulate.py b/test/local_test_null_coalesce_accumulate.py
new file mode 100644
--- /dev/null
+++ b/test/local_test_null_coalesce_accumulate.py
@@ -0,0 +1,104 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+import argparse
+import logging
+import os
+import socket
+
+import torch
+import torch.distributed.rpc as rpc
+import torch.multiprocessing as mp
+
+from pippy.IR import Pipe, TrivialLossWrapper, pipe_split, _null_coalesce_accumulate
+from pippy.PipelineDriver import PipelineDriverBase, PipelineDriverFillDrain, PipelineDriver1F1B
+from pippy.microbatch import TensorChunkSpec, CustomReducer
+
+PROFILING_ENABLED = True
+CHECK_NUMERIC_EQUIVALENCE = True
+
+schedules = {
+    'FillDrain': PipelineDriverFillDrain,
+    '1F1B': PipelineDriver1F1B,
+}
+
+VERBOSE = bool(int(os.environ.get('VERBOSE', False)))
+
+if VERBOSE:
+    logging.getLogger().setLevel(logging.DEBUG)
+
+torch.fx.Tracer.proxy_buffer_attributes = True
+
+
+def run_master(args):
+    all_ranks = list(range(1, args.world_size))  # exclude master rank = 0
+    chunks = len(all_ranks)
+    bs = 4 * chunks
+    hid_dim = 50
+
+    class Code(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(hid_dim, hid_dim)
+
+        def forward(self, x):
+            x = self.linear(x)
+            pipe_split()
+            y = torch.relu(x)
+            pipe_split()
+            z = torch.sigmoid(x)
+            pipe_split()
+            return y + z
+
+    c = Code()
+    c.train()
+    mse_loss = torch.nn.MSELoss()
+    wrapper = TrivialLossWrapper(c, mse_loss)
+    accum_pipe = Pipe.from_tracing(wrapper)
+    assert 4 == len(list(accum_pipe.split_gm.children()))
+    assert any(n.target == _null_coalesce_accumulate for n in accum_pipe.split_gm.graph.nodes)
+
+    input = torch.randn(bs, hid_dim)
+    target = torch.randn(bs, hid_dim)
+    accum_pipe(input, target)
+
+    args_chunk_spec = (TensorChunkSpec(0), TensorChunkSpec(0))
+    kwargs_chunk_spec = {}
+    output_chunk_spec = CustomReducer(torch.tensor(0.0), lambda a, b: a + b)
+    pipe_driver: PipelineDriverBase = schedules[args.schedule](accum_pipe, args_chunk_spec, kwargs_chunk_spec,
+                                                               output_chunk_spec, args.world_size - 1,
+                                                               all_ranks=all_ranks, _debug_mask_minibatches=True)
+
+    pipe_driver.run(chunks, input, target)
+
+
+def run_worker(rank, world_size, args):
+    print(f"rank = {rank} host/pid = {socket.gethostname()}/{os.getpid()}")
+    os.environ['MASTER_ADDR'] = args.master_addr
+    os.environ['MASTER_PORT'] = args.master_port
+    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256)
+    rpc.init_rpc(
+        f"worker{rank}",
+        rank=rank,
+        world_size=world_size,
+        rpc_backend_options=options
+    )
+    if rank == 0:
+        run_master(args)
+    rpc.shutdown()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--world_size', type=int, default=int(os.getenv("WORLD_SIZE", 5)))
+    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'))
+    parser.add_argument('-s', '--schedule', type=str, default=list(schedules.keys())[0], choices=schedules.keys())
+    parser.add_argument('--replicate', type=int, default=int(os.getenv("REPLICATE", '0')))
+    args = parser.parse_args()
+
+    if args.rank == -1:
+        mp.spawn(run_worker, args=(args.world_size, args,), nprocs=args.world_size, join=True)
+    elif args.rank < args.world_size:
+        run_worker(args.rank, args.world_size, args)
+    else:
+        print("I'm unused, exiting")